From b1cecf5f798b15093d4aa52dddd40eecb4f25fa8 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 16 Jul 2026 03:24:40 +0000 Subject: [PATCH] Revert "chore: remove fsdp_cpu_offload (#3049)" This reverts commit c9acaa3c9070f445c4a167cb8569cc482ec5b523. --- docs/scaling.md | 3 +- .../src/prime_rl/configs/sft.py | 6 ++++ .../src/prime_rl/configs/trainer.py | 15 ++++++++++ src/prime_rl/trainer/model.py | 30 +++++++++++++++++-- src/prime_rl/trainer/rl/train.py | 4 ++- src/prime_rl/trainer/sft/train.py | 4 ++- src/prime_rl/trainer/utils.py | 11 +++++-- 7 files changed, 65 insertions(+), 8 deletions(-) diff --git a/docs/scaling.md b/docs/scaling.md index e716cce3fff..6b145560e64 100644 --- a/docs/scaling.md +++ b/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 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 1f457c135ec..1760ce8acac 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/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" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 43214b7db9f..8f790c24e55 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/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: diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 6c7188d5d49..a37aca177cf 100644 --- a/src/prime_rl/trainer/model.py +++ b/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,6 +1180,9 @@ 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 @@ -1169,7 +1190,8 @@ def setup_model( 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) diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index a501c293937..15c901a67a7 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/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. diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 899f75a367d..692b5488eae 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/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. diff --git a/src/prime_rl/trainer/utils.py b/src/prime_rl/trainer/utils.py index 8919f8457ba..bc0f2e8c93e 100644 --- a/src/prime_rl/trainer/utils.py +++ b/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):