avoid fsdp2 shard of lm_head to enable fused_kernels

Author: wqwqazwsxedc

apertus/launch/multinode_async_sandbox/_verl_training.sbatch

@@ -13,12 +13,14 @@
 set -xeuo pipefail
 
 clear_inherited_pyxis_options() {
+  set +x
   local name
   while IFS='=' read -r name _; do
     case "${name}" in
       SLURM_SPANK__SLURM_SPANK_OPTION_pyxis_*) unset "${name}" ;;
     esac
   done < <(env)
+  set -x
 }
 
 clear_inherited_pyxis_options

apertus/launch/multinode_async_sandbox/launch.sh

@@ -73,9 +73,9 @@ ASYNC_STEADY_WARMUP_STEPS="${ASYNC_STEADY_WARMUP_STEPS:-}"
 ###############################################################################
 
 # Set REASONING_GYM_DIR="" to install reasoning-gym from PyPI.
-REASONING_GYM_DIR="${REASONING_GYM_DIR:-${SCRATCH_HOME}/projects/r-gym}"
-TOOL_GYM_DIR="${TOOL_GYM_DIR:-${SCRATCH_HOME}/projects/tool-gym}"
-TOOL_GYM_FUNCTION_TOOL_PATH="${TOOL_GYM_FUNCTION_TOOL_PATH:-/capstor/store/cscs/swissai/infra01/reasoning/data/RL-prod/toolgym_test_v2/apertus_function_tools.py}"
+REASONING_GYM_DIR=""
+TOOL_GYM_DIR=""
+TOOL_GYM_FUNCTION_TOOL_PATH="${TOOL_GYM_FUNCTION_TOOL_PATH:-/capstor/store/cscs/swissai/infra01/reasoning/data/RL-prod/toolgym_test_v3/apertus_function_tools_v3.py}"
 SANDBOX_BACKEND="kubernetes"  # kubernetes, codegym, or none
 KUBERNETES_SANDBOX_URL="https://sandbox-dev.swissai.svc.cscs.ch"
 CODE_GYM_DIR="" # ${SCRATCH_HOME}/projects/code-gym}  # Not needed if using kubernetes

tests/checkpoint_engine/test_correctness_on_gpu.py

@@ -42,7 +42,7 @@ async def test_nccl_checkpoint_engine(
     num_gpus_per_node=_ngpus,
     bucket_size_mb=128,
     check_allclose=True,
-    model_path="~/models/Qwen/Qwen3-8B-Base",
+    model_path="swiss-ai/Apertus-8B-Instruct-2509",
 ):
     model_path = os.path.expanduser(model_path)
     ray.init(

tests/checkpoint_engine/test_utils.py

@@ -39,6 +39,7 @@ def __init__(self, config: TrainingWorkerConfig, checkpoint_engine_config: Check
         if torch.distributed.get_rank() == 0:
             engine_kwargs["is_master"] = True
         self.checkpoint_engine = CheckpointEngineRegistry.new(backend, bucket_size=bucket_size, **engine_kwargs)
+        
 
     @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
     async def update_weights(self, global_steps: int = None, mode: str = "auto"):
@@ -142,7 +143,7 @@ def check_weights(self):
 def create_trainer_worker_group(
     resource_pool: RayResourcePool, model_config: HFModelConfig, checkpoint_engine_config: CheckpointEngineConfig
 ) -> RayWorkerGroup:
-    engine_config = FSDPEngineConfig(forward_only=True, fsdp_size=resource_pool.world_size, strategy="fsdp")
+    engine_config = FSDPEngineConfig(forward_only=True, fsdp_size=resource_pool.world_size, strategy="fsdp2")
     trainer_config = TrainingWorkerConfig(
         model_type="language_model",
         model_config=model_config,

verl/experimental/fully_async_policy/config/async.yaml

@@ -66,19 +66,29 @@ actor_rollout_ref:
     use_remove_padding: true
     enable_gradient_checkpointing: true
     use_shm: false
+    use_fused_kernels: true
+    fused_kernel_options:
+      impl_backend: torch
     override_config:
       attn_implementation: flash_attention_3
 
   actor:
     use_rollout_log_probs: ${oc.select:async_training.use_rollout_log_probs,true}
     ppo_mini_batch_size: 128  # NOTE: train_batch_size = require_batches * ppo_mini_batch_size
     use_dynamic_bsz: true
-    ppo_max_token_len_per_gpu: 32768
+    ppo_max_token_len_per_gpu: 16384
     ulysses_sequence_parallel_size: 1
     entropy_from_logits_with_chunking: false
     entropy_checkpointing: false
-
-
+    profiler:
+      enable: False
+      all_ranks: False
+      ranks: [0]
+      tool: torch
+      tool_config:
+        torch:
+          contents: [cpu, memory, cuda, shapes, stack]
+          discrete: True
     optim:
       optimizer: _AdamW
       lr: 1e-6
@@ -92,6 +102,9 @@ actor_rollout_ref:
       strategy: fsdp2
       param_offload: false
       optimizer_offload: false
+      reshard_after_forward: false
+      forward_prefetch: true
+      fsdp_size: 4
       entropy_from_logits_with_chunking: ${oc.select:actor_rollout_ref.actor.entropy_from_logits_with_chunking,false}
       entropy_checkpointing: ${oc.select:actor_rollout_ref.actor.entropy_checkpointing,false}
       model_dtype: bf16
@@ -164,6 +177,14 @@ trainer:
   rollout_data_dir: ${trainer.default_local_dir}/rollout/
   validation_data_dir: ${trainer.default_local_dir}/validation/
 
+global_profiler:
+  steps: [4]
+  save_path: "/iopsstor/scratch/cscs/atazza/verl_profile"
+  tool: torch
+  tool_config:
+    torch:
+      contents: [cpu, memory, cuda, stack]
+
 ray_kwargs:
   ray_init:
     num_cpus: ${oc.decode:${oc.env:SLURM_CPUS_PER_TASK,null}}