Revert "chore: remove fsdp_cpu_offload"

docs/scaling.md

@@ -85,6 +85,7 @@ FSDP2 is the default model sharding strategy. By default the trainer fully shard
 |---|---|
 | `trainer.model.dp_replicate` | Number of dimensions to **replicate** instead of shard. Set to 2 to run 2-way DP replication × FSDP sharding within each replica — useful for very large clusters where pure FSDP communication dominates. |
 | `trainer.model.reshard_after_forward` | If `true` (default), parameters are resharded after the forward pass to free memory; the backward pass re-gathers. Set `false` to keep params resident — faster but more memory. |
+| `trainer.model.fsdp_cpu_offload` | Offload params + grads + optimizer state to CPU. Big memory win, large throughput hit. |
 | `trainer.model.optim_cpu_offload` | Offload only optimizer state. Mid-ground — small throughput cost, decent memory savings, especially at low GPU count. |
 
 ### Expert Parallelism
@@ -143,7 +144,7 @@ Offloading optimizer states to CPU is enabled by default (`optim_cpu_offload = t
 optim_cpu_offload = true   # already the default
 ```
 
-Incompatible with `trainer.max_concurrent_runs > 1` (multi-tenant training) — set `optim_cpu_offload = false` for multi-run.
+Mutually exclusive with `fsdp_cpu_offload`. Also incompatible with `trainer.max_concurrent_runs > 1` (multi-tenant training) — set `optim_cpu_offload = false` for multi-run. Muon doesn't support `fsdp_cpu_offload` but does support `optim_cpu_offload`.
 
 ### LM Head Chunking
 

packages/prime-rl-configs/src/prime_rl/configs/sft.py

@@ -355,6 +355,12 @@ def validate_lora_adapter_saving(self):
                 )
         return self
 
+    @model_validator(mode="after")
+    def validate_opt_and_fsdp_offload(self):
+        if self.optim.type == "muon" and self.model.fsdp_cpu_offload:
+            raise ValueError("Muon optimizer does not support FSDP CPU offload")
+        return self
+
     @model_validator(mode="after")
     def validate_and_disable_chunked_loss(self):
         self.model.fused_lm_head_token_chunk_size = "disabled"

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -135,6 +135,9 @@ class ModelConfig(BaseModelConfig):
     ac_offloading: ActivationOffloadingConfig | None = ActivationOffloadingConfig()
     """Activation offloading configuration. If None, activation offloading is disabled."""
 
+    fsdp_cpu_offload: bool = False
+    """Enable FSDP CPU offloading for parameters, gradients, and optimizer states. Uses pinned memory for efficient CPU↔GPU transfers."""
+
     optim_cpu_offload: bool = True
     """Offload only optimizer states (momentum, variance) to CPU, keeping weights on GPU. Avoids the H2D all-gather overhead of FSDP CPU offload while still saving GPU memory."""
 
@@ -239,6 +242,12 @@ def selective_ac_only_with_custom_impl(self):
             raise ValueError("Selective activation checkpointing requires model.impl='custom' or 'auto'")
         return self
 
+    @model_validator(mode="after")
+    def cpu_offload_mutual_exclusion(self):
+        if self.fsdp_cpu_offload and self.optim_cpu_offload:
+            raise ValueError("Cannot enable both fsdp_cpu_offload and optim_cpu_offload. Use one or the other.")
+        return self
+
     @model_validator(mode="after")
     def flash_attention_4_only_with_custom_impl(self):
         if self.attn == "fa4" and self.impl != "custom":
@@ -631,6 +640,12 @@ def validate_lora_adapter_saving(self):
                 )
         return self
 
+    @model_validator(mode="after")
+    def validate_opt_and_fsdp_offload(self):
+        if self.optim.type == "muon" and self.model.fsdp_cpu_offload:
+            raise ValueError("Muon optimizer does not support FSDP CPU offload")
+        return self
+
     @model_validator(mode="after")
     def validate_optim_cpu_offload_single_run(self):
         if self.model.optim_cpu_offload and self.max_concurrent_runs > 1:

src/prime_rl/trainer/model.py

@@ -19,7 +19,7 @@
 from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageReader
 from torch.distributed.checkpoint.state_dict_loader import load as dcp_load
 from torch.distributed.device_mesh import DeviceMesh
-from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, fully_shard
+from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, OffloadPolicy, fully_shard
 from torch.distributed.tensor.parallel import parallelize_module
 from torchtitan.distributed.expert_parallel import ExpertParallel
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, PretrainedConfig
@@ -666,9 +666,11 @@ def setup_tokenizer(config: TokenizerConfig) -> PreTrainedTokenizer:
 
 def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDims):
     mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=DTYPE_MAP[config.reduce_dtype])
+    offload_policy: OffloadPolicy = CPUOffloadPolicy(pin_memory=True) if config.fsdp_cpu_offload else OffloadPolicy()
 
     fsdp_config = {
         "mp_policy": mp_policy,
+        "offload_policy": offload_policy,
         "reshard_after_forward": config.reshard_after_forward,
     }
 
@@ -709,6 +711,7 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim
                 block_mlp.router,
                 mesh=hsdp_mesh,
                 mp_policy=MixedPrecisionPolicy(param_dtype=torch.float32, reduce_dtype=torch.float32),
+                offload_policy=offload_policy,
                 reshard_after_forward=config.reshard_after_forward,
             )
 
@@ -733,6 +736,7 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim
             [model.lm_head, norm_module],
             mesh=hsdp_mesh,
             mp_policy=mp_policy,
+            offload_policy=offload_policy,
             reshard_after_forward=False,
         )
     else:
@@ -742,6 +746,7 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim
         model,
         mesh=hsdp_mesh,
         mp_policy=mp_policy,
+        offload_policy=offload_policy,
         reshard_after_forward=config.reshard_after_forward,
     )
 
@@ -800,7 +805,8 @@ def setup_fsdp(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDim
 
 
 def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDims):
-    model.to_empty(device="cuda")
+    device = "cpu" if config.fsdp_cpu_offload else "cuda"
+    model.to_empty(device=device)
     torch.distributed.barrier()
 
     def _init_buffers_post_meta():
@@ -813,6 +819,7 @@ def _init_buffers_post_meta():
     if config.debug.random_init:
         logger.warning("Randomly initializing model. Skipping loading weights from HF.")
         _init_buffers_post_meta()
+        _move_buffers_to_cuda(model, config)
         return
 
     if not Path(config.name).exists():
@@ -878,6 +885,8 @@ def _init_buffers_post_meta():
         model.tie_weights()
     _init_buffers_post_meta()
 
+    _move_buffers_to_cuda(model, config)
+
     lora_modules = [m for m in model.modules() if hasattr(m, "_init_lora_parameters")]
     if lora_modules:
         generator: torch.Generator | None = None
@@ -1044,6 +1053,15 @@ def apply_ep(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDims)
             )
 
 
+def _move_buffers_to_cuda(model: nn.Module, config: ModelConfig) -> None:
+    """FSDP CPU offloading only manages parameters, not buffers. Move buffers to CUDA."""
+    if not config.fsdp_cpu_offload:
+        return
+    for _, buffer in model.named_buffers():
+        if buffer.device.type == "cpu":
+            buffer.data = buffer.data.to("cuda")
+
+
 def _reset_runtime_moe_buffers(model: nn.Module) -> None:
     for module in model.modules():
         if isinstance(module, (MoE, LatentMoE)) and module.tokens_per_expert.device.type != "meta":
@@ -1162,14 +1180,18 @@ def setup_model(
 
     setup_fsdp(model, config, parallel_dims)
 
+    if not possible_to_load_to_meta:
+        _move_buffers_to_cuda(model, config)
+
     # 2. if we can load to meta, we either:
     if possible_to_load_to_meta:
         # - load from checkpoint later if needed
         if loading_from_checkpoint_later:
             logger.warning(
                 "Skipping loading weights. Initializing an empty model on device, loading from checkpoint later."
             )
-            model.to_empty(device="cuda")
+            device = "cpu" if config.fsdp_cpu_offload else "cuda"
+            model.to_empty(device=device)
             torch.distributed.barrier()
             if isinstance(model, PreTrainedModelPrimeRL):
                 model.init_buffers_post_meta()
@@ -1178,6 +1200,8 @@ def setup_model(
                 # Restore weight tying broken by to_empty() for HF models
                 if model.config.tie_word_embeddings:
                     model.tie_weights()
+
+            _move_buffers_to_cuda(model, config)
         # - or load from HF with dcp
         else:
             load_dcp_from_hf(model, config, parallel_dims)

src/prime_rl/trainer/rl/train.py

@@ -108,7 +108,9 @@ def train(config: TrainerConfig):
             health_server.start()
 
     # Set precision
-    setup_torch_distributed(timeout=timedelta(seconds=config.dist_timeout_seconds))
+    setup_torch_distributed(
+        timeout=timedelta(seconds=config.dist_timeout_seconds), enable_gloo=config.model.fsdp_cpu_offload
+    )
     # Configurable to support ROCm/AMD GPUs where reduced precision
     # matmul corrupts softmax over large vocabularies. Override via config
     # (e.g. matmul_precision = "highest") on ROCm.

src/prime_rl/trainer/sft/train.py

@@ -82,7 +82,9 @@ def train(config: SFTConfig):
         heart = Heartbeat(config.heartbeat.url)
 
     # Set precision
-    setup_torch_distributed(timeout=timedelta(seconds=config.dist_timeout_seconds))
+    setup_torch_distributed(
+        timeout=timedelta(seconds=config.dist_timeout_seconds), enable_gloo=config.model.fsdp_cpu_offload
+    )
     # Configurable to support ROCm/AMD GPUs where reduced precision
     # matmul corrupts softmax over large vocabularies. Override via config
     # (e.g. matmul_precision = "highest") on ROCm.

src/prime_rl/trainer/utils.py

@@ -367,10 +367,17 @@ def get_ckpt_disk_metrics(output_dir: Path) -> dict[str, float]:
     }
 
 
-def setup_torch_distributed(timeout: timedelta = DEFAULT_TIMEOUT):
+def setup_torch_distributed(timeout: timedelta = DEFAULT_TIMEOUT, enable_gloo: bool = False):
     device_id = get_world().local_rank
     torch.cuda.set_device(device_id)
-    dist.init_process_group(timeout=timeout, device_id=device_id)
+    # Use Gloo backend for CPU and NCCL for GPU when CPU offloading is enabled
+    # Otherwise use NCCL for better GPU performance
+    backend = None  # by default nccl
+    if enable_gloo:
+        get_logger().info("Using Gloo backend for CPU and NCCL backend for GPU")
+        backend = "cpu:gloo,cuda:nccl"
+
+    dist.init_process_group(backend=backend, timeout=timeout, device_id=device_id)
 
 
 def print_sample(input_ids: list[int], loss_mask: list[bool], tokenizer: PreTrainedTokenizer):