refactor(configs)!: one model.offload field for CPU offloading

configs/advanced/deepseek-v4-flash/sft.toml

@@ -9,10 +9,9 @@
 # `dist_timeout_seconds = 14400` is needed on any run that has to build the HF -> prime-rl
 # conversion cache, which takes far longer than the 3600s default.
 #
-# To buy memory back for quicker testing on fewer nodes, set `full_offload = true` under `[model]`,
-# which keeps the fp32 masters, moments and gradients in CPU RAM and needs `optim_cpu_offload =
-# false` and `optim.max_norm = "None"` set alongside it; the lighter `optim_cpu_offload = true` is
-# already the default here.
+# To buy memory back for quicker testing on fewer nodes, set `offload = "full"` under `[model]`,
+# which keeps the fp32 masters, moments and gradients in CPU RAM and disables gradient clipping;
+# the lighter `offload = "optimizer"` is already the default here.
 #
 # There is no `[ckpt]` block. The trainer saves unconditionally at the last step, and a 43-layer
 # checkpoint (fp32 master weights plus two fp32 AdamW moments) is ~3 TiB. Add one with

configs/ci/integration/reverse-text-sft/full-offload-resume.toml

@@ -9,8 +9,7 @@ step = 5
 
 [model]
 name = "PrimeIntellect/Qwen3-0.6B"
-optim_cpu_offload = false
-full_offload = true
+offload = "full"
 
 [optim]
 lr = 0.0

docs/scaling.md

@@ -86,9 +86,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 optimizer state to CPU between steps. Enabled by default. |
-| `trainer.model.full_offload` | Offload gradients, FP32 masters, and optimizer state and run the optimizer (AdamW or SignSGD) on CPU during backward. Disabled by default. |
+| `trainer.model.offload` | CPU offloading: `optimizer` (default) offloads optimizer state between steps; `full` offloads gradients, FP32 masters, and optimizer state and runs the optimizer (AdamW or SignSGD) on CPU during backward; `fsdp` offloads params + grads + optimizer state through FSDP (big memory win, large throughput hit); `none` disables offloading. |
 
 ### Expert Parallelism
 
@@ -143,7 +141,7 @@ Activation offloading still applies to tensors saved by autograd, but tensors re
 
 ### Optimizer Offloading
 
-State-only optimizer offload remains enabled by default with `model.optim_cpu_offload = true`. For full offload, set `model.optim_cpu_offload = false` and `model.full_offload = true`; this keeps BF16 compute weights on GPU and runs CPU optimizer chunks as gradients become ready during backward. Full offload only supports AdamW and SignSGD (`optim.type = "sign_sgd"`) and disables gradient clipping. SignSGD is stateless, so it halves the host RAM footprint versus AdamW (8 instead of 16 bytes per parameter: FP32 master + FP32 accumulated gradient, no moments).
+State-only optimizer offload is enabled by default (`model.offload = "optimizer"`). For full offload, set `model.offload = "full"`; this keeps BF16 compute weights on GPU and runs CPU optimizer chunks as gradients become ready during backward. Full offload only supports AdamW and SignSGD (`optim.type = "sign_sgd"`) and disables gradient clipping. SignSGD is stateless, so it halves the host RAM footprint versus AdamW (8 instead of 16 bytes per parameter: FP32 master + FP32 accumulated gradient, no moments).
 
 ### LM Head Chunking
 

examples/advanced/glm-4.5-air/swe-budget.toml

@@ -57,8 +57,7 @@ name = "zai-org/GLM-4.5-Air"
 attn = "flash_attention_3"
 cp = 4
 cp_style = "ulysses"
-optim_cpu_offload = false
-full_offload = true
+offload = "full"
 fused_lm_head_token_chunk_size = 1024
 
 [trainer.model.ac]

examples/advanced/glm-5.3/sft/h200/16k-12-node.toml

@@ -4,7 +4,7 @@
 
 [model]
 cp = 1
-optim_cpu_offload = false
+offload = "none"
 ac_offloading = "None"
 
 [model.compile]

examples/advanced/glm-5.3/sft/h200/base.toml

@@ -8,7 +8,7 @@ name = "zai-org/GLM-5.3-BF16"
 attn = "flash_attention_2"
 cp = 8
 ep = 8
-optim_cpu_offload = true
+offload = "optimizer"
 dsa_backend = "cudnn_flashmla"
 
 [model.ac]

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

@@ -388,13 +388,13 @@ def deepep_disables_grad_clipping(self):
 
     @model_validator(mode="after")
     def full_optimizer_offload_requires_supported_optimizer(self):
-        if self.model.full_offload and self.optim.type not in ("adamw", "sign_sgd"):
+        if self.model.offload == "full" and self.optim.type not in ("adamw", "sign_sgd"):
             raise ValueError("Full optimizer offload only supports AdamW and SignSGD")
         return self
 
     @model_validator(mode="after")
     def full_optimizer_offload_disables_grad_clipping(self):
-        if self.model.full_offload and self.optim.max_norm is not None:
+        if self.model.offload == "full" and self.optim.max_norm is not None:
             warnings.warn(
                 "Gradient clipping prevents optimizer-in-backward overlap with CPU optimizer offload. "
                 "Automatically setting optim.max_norm to None (disabled).",
@@ -623,7 +623,7 @@ def validate_scheduler_steps(self):
 
     @model_validator(mode="after")
     def validate_opt_and_fsdp_offload(self):
-        if self.optim.type == "muon" and self.model.fsdp_cpu_offload:
+        if self.optim.type == "muon" and self.model.offload == "fsdp":
             raise ValueError("Muon optimizer does not support FSDP CPU offload")
         return self
 

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

@@ -3,7 +3,7 @@
 from pathlib import Path
 from typing import Annotated, Any, Literal, TypeAlias
 
-from pydantic import BeforeValidator, Field, field_validator, model_validator
+from pydantic import Field, field_validator, model_validator
 
 from prime_rl.configs.monitors import MonitorsConfig
 from prime_rl.configs.shared import (
@@ -49,34 +49,7 @@ class ActivationOffloadingConfig(BaseConfig):
     """Max activations kept in flight while offloading. More activations smooth overlap at the cost of GPU memory."""
 
 
-class OptimizerInBackwardOffloadConfig(BaseConfig):
-    """Full CPU optimizer offload: FP32 masters, optimizer state (AdamW moments; SignSGD is
-    stateless), and accumulated gradients live in CPU RAM, each optimizer chunk runs on CPU as
-    soon as its last gradient arrives, and the refreshed BF16 weights stream back while backward
-    is still executing.
-
-    Gradient numerics: gradients are reduced across ranks in FP32 (``reduce_dtype``) but FSDP2
-    materializes them in the sharded parameter's dtype, which is BF16 for the offload compute
-    model — so each gradient is rounded to BF16 once before the FP32 CPU update. Masters,
-    moments, accumulation, and optimizer arithmetic remain FP32. For gradient numerics
-    bit-faithful to that path, disable offloading.
-    """
-
-    numa_bind: bool = True
-    """Pin each rank's CPUs to its GPU's NUMA node. Disable when the launcher already manages CPU affinity or GPU sysfs topology is unavailable."""
-
-
-def _normalize_optimizer_in_backward_offload(value: Any) -> Any:
-    if value is True:
-        return {}
-    if value is False:
-        return None
-    return value
-
-
-OptimizerInBackwardOffload = Annotated[
-    OptimizerInBackwardOffloadConfig | None, BeforeValidator(_normalize_optimizer_in_backward_offload)
-]
+OffloadMode: TypeAlias = Literal["none", "optimizer", "full", "fsdp"]
 
 
 class CompileConfig(BaseConfig):
@@ -288,14 +261,8 @@ 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."""
-
-    full_offload: OptimizerInBackwardOffload = None
-    """Full CPU optimizer offload: FP32 masters, moments, and gradients live in CPU RAM and the optimizer runs on CPU, overlapped with backward. Enable with ``true`` or a ``[model.full_offload]`` section; disabled by default."""
+    offload: OffloadMode = "optimizer"
+    """CPU offloading. ``optimizer`` keeps optimizer states (momentum, variance) in CPU RAM and weights on GPU, avoiding the H2D all-gather overhead of ``fsdp`` while still saving GPU memory. ``full`` keeps FP32 masters, optimizer states, and accumulated gradients in CPU RAM and runs each optimizer chunk on CPU as soon as its last gradient arrives, overlapped with backward; it supports AdamW and SignSGD only, disables gradient clipping, and rounds each gradient to BF16 once before the FP32 CPU update. ``fsdp`` offloads parameters, gradients, and optimizer states through FSDP with pinned memory. ``none`` keeps everything on GPU."""
 
     reshard_after_forward: bool = True
     """Reshard the model after each forward pass."""
@@ -363,17 +330,6 @@ def ac_offloading_requires_ac(self):
             self.ac = ActivationCheckpointConfig()
         return self
 
-    @model_validator(mode="after")
-    def cpu_offload_mutual_exclusion(self):
-        if self.fsdp_cpu_offload and (self.optim_cpu_offload or self.full_offload):
-            raise ValueError("Cannot combine fsdp_cpu_offload with optimizer CPU offloading.")
-        if self.optim_cpu_offload and self.full_offload:
-            raise ValueError(
-                "Cannot enable both optim_cpu_offload and full_offload. "
-                "Set optim_cpu_offload=false when enabling full optimizer offload."
-            )
-        return self
-
     @model_validator(mode="after")
     def validate_moe_runtime(self):
         if self.ep == 1:
@@ -711,13 +667,13 @@ def deepep_disables_grad_clipping(self):
 
     @model_validator(mode="after")
     def full_optimizer_offload_requires_supported_optimizer(self):
-        if self.model.full_offload and self.optim.type not in ("adamw", "sign_sgd"):
+        if self.model.offload == "full" and self.optim.type not in ("adamw", "sign_sgd"):
             raise ValueError("Full optimizer offload only supports AdamW and SignSGD")
         return self
 
     @model_validator(mode="after")
     def full_optimizer_offload_disables_grad_clipping(self):
-        if self.model.full_offload and self.optim.max_norm is not None:
+        if self.model.offload == "full" and self.optim.max_norm is not None:
             warnings.warn(
                 "Gradient clipping prevents optimizer-in-backward overlap with CPU optimizer offload. "
                 "Automatically setting optim.max_norm to None (disabled).",
@@ -753,7 +709,7 @@ def validate_scheduler_steps(self):
 
     @model_validator(mode="after")
     def validate_opt_and_fsdp_offload(self):
-        if self.optim.type == "muon" and self.model.fsdp_cpu_offload:
+        if self.optim.type == "muon" and self.model.offload == "fsdp":
             raise ValueError("Muon optimizer does not support FSDP CPU offload")
         return self
 

skills/training/start-run/SKILL.md

@@ -26,12 +26,11 @@ SLURM launches write generated scripts and coordination files under `<run_dir>/l
 - Validation aliases let renamed fields keep working; legacy keys can be remapped in a `model_validator(mode="before")`.
 - Auto-generated `--help` panels from `Field(description=...)` or PEP 224 docstrings.
 - Friendly errors: required-field boxes, validator errors point at the offending flag, unknown flags get a "did you mean" hint.
-- State-only optimizer offload remains enabled by default with `model.optim_cpu_offload = true`.
+- State-only optimizer offload is enabled by default (`model.offload = "optimizer"`).
 - For gradients, FP32 masters, optimizer state, and optimizer-in-backward CPU execution, set
-  `model.optim_cpu_offload = false` and `model.full_offload = true`. This mode uses the native
-  CPU optimizer kernel, only supports AdamW and SignSGD (SignSGD is stateless and
-  halves the host RAM footprint), and disables gradient clipping. Use a
-  `[model.full_offload]` table only to disable NUMA binding.
+  `model.offload = "full"`. This mode uses the native CPU optimizer kernel, only supports
+  AdamW and SignSGD (SignSGD is stateless and halves the host RAM footprint), disables
+  gradient clipping, and pins each rank's CPUs to its GPU's NUMA node.
 
 ## `rl` — RL training
 

src/prime_rl/trainer/model.py

@@ -509,7 +509,7 @@ def placement_fn(parameter: nn.Parameter) -> ShardPlacementResult | Shard | None
 
 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()
+    offload_policy: OffloadPolicy = CPUOffloadPolicy(pin_memory=True) if config.offload == "fsdp" else OffloadPolicy()
 
     fused_shard_placement_fn = get_fsdp_shard_placement_fn(model) if config.fusions.shard_fused_on_dim1 else None
     hsdp_mesh = parallel_dims.get_mesh("hsdp")
@@ -695,7 +695,7 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None:
 
 
 def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDims):
-    device = "cpu" if config.fsdp_cpu_offload else "cuda"
+    device = "cpu" if config.offload == "fsdp" else "cuda"
     model.to_empty(device=device)
     torch.distributed.barrier()
 
@@ -893,7 +893,7 @@ def configure_trainable_parameters(model: nn.Module, config: ModelConfig) -> nn.
 
 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:
+    if config.offload != "fsdp":
         return
     for _, buffer in model.named_buffers():
         if buffer.device.type == "cpu":
@@ -1024,7 +1024,7 @@ def setup_model(
         logger.warning(
             "Skipping loading weights. Initializing an empty model on device, loading from checkpoint later."
         )
-        device = "cpu" if config.fsdp_cpu_offload else "cuda"
+        device = "cpu" if config.offload == "fsdp" else "cuda"
         model.to_empty(device=device)
         torch.distributed.barrier()
         model.init_buffers_post_meta()

src/prime_rl/trainer/optim/__init__.py

@@ -5,7 +5,7 @@
 from torch.distributed.device_mesh import DeviceMesh
 from torch.optim import AdamW, Optimizer
 
-from prime_rl.configs.trainer import OptimizerConfig, OptimizerInBackwardOffloadConfig
+from prime_rl.configs.trainer import OffloadMode, OptimizerConfig
 from prime_rl.trainer.models.fusions import get_model_packed_parameters
 from prime_rl.trainer.optim.base import OffloadOptimizer as OffloadOptimizer
 from prime_rl.trainer.optim.base import OptimizerLike
@@ -51,21 +51,20 @@ def setup_optimizer(
     config: OptimizerConfig,
     named_params: list[tuple[str, nn.Parameter]],
     parallel_dims: ParallelDims,
-    cpu_offload: bool = False,
-    full_offload_config: OptimizerInBackwardOffloadConfig | None = None,
+    offload: OffloadMode = "none",
     model: nn.Module | None = None,
     full_offload_dtype_policy: dict[int, tuple[torch.dtype, torch.dtype]] | None = None,
 ) -> tuple[OptimizerLike, GradientOffloadManager | None]:
-    if cpu_offload and full_offload_config is not None:
-        raise ValueError("State-only and full optimizer CPU offload cannot both be enabled")
-    if full_offload_config is not None and config.type not in ("adamw", "sign_sgd"):
+    cpu_offload = offload == "optimizer"
+    full_offload = offload == "full"
+    if full_offload and config.type not in ("adamw", "sign_sgd"):
         raise ValueError("Full optimizer offload only supports AdamW and SignSGD")
-    if full_offload_config is not None and config.max_norm is not None:
+    if full_offload and config.max_norm is not None:
         get_logger().warning("Disabling gradient clipping because CPU optimizer offload updates during backward")
         config.max_norm = None
     optimizer_named_params = named_params
     master_weights = None
-    if full_offload_config is not None:
+    if full_offload:
         if model is None:
             raise ValueError("CPU optimizer offload requires the model")
         if full_offload_dtype_policy is None:
@@ -82,12 +81,11 @@ def setup_optimizer(
         model=model,
     )
 
-    if full_offload_config is not None:
+    if full_offload:
         assert master_weights is not None
         get_logger().info("Using CPU offload for gradients and the optimizer step")
         optimizer = FullCPUOffloadOptimizer(
             optimizer,
-            offload_config=full_offload_config,
             master_weights=master_weights,
             dp_replicate=parallel_dims.dp_replicate,
         )

src/prime_rl/trainer/optim/offload.py

@@ -14,7 +14,6 @@
 from torch.distributed.tensor import DTensor
 from torch.optim import AdamW, Optimizer
 
-from prime_rl.configs.trainer import OptimizerInBackwardOffloadConfig
 from prime_rl.trainer.optim.base import OffloadOptimizer
 from prime_rl.trainer.optim.cpu_adam import adamw_step as native_cpu_adamw_step
 from prime_rl.trainer.optim.cpu_adam import copy_or_add_bfloat16_multi_ as native_copy_or_add_bfloat16_multi_
@@ -714,12 +713,10 @@ class FullCPUOffloadOptimizer(OffloadOptimizer):
     def __init__(
         self,
         optimizer: Optimizer,
-        offload_config: OptimizerInBackwardOffloadConfig,
         master_weights: dict[int, _MasterWeight],
         dp_replicate: int = 1,
     ):
         self.optimizer = optimizer
-        self.offload_config = offload_config
         self._initialized = False
         self._master_weights = master_weights
         self._chunks = self._build_chunks()

src/prime_rl/trainer/rl/train.py

@@ -106,10 +106,10 @@ def train(config: TrainerConfig):
     # Set precision
     setup_torch_distributed(
         timeout=timedelta(seconds=config.dist_timeout_seconds),
-        enable_gloo=config.model.fsdp_cpu_offload or config.model.full_offload is not None,
+        enable_gloo=config.model.offload in ("fsdp", "full"),
     )
-    if config.model.full_offload is not None:
-        setup_full_cpu_optimizer_offload(config.model.full_offload)
+    if config.model.offload == "full":
+        setup_full_cpu_optimizer_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.
@@ -165,11 +165,10 @@ def train(config: TrainerConfig):
         config.optim,
         list(model.named_parameters()),
         parallel_dims,
-        cpu_offload=config.model.optim_cpu_offload,
-        full_offload_config=config.model.full_offload,
+        offload=config.model.offload,
         model=model,
         full_offload_dtype_policy=(
-            get_full_offload_dtype_policy(model, config.model) if config.model.full_offload is not None else None
+            get_full_offload_dtype_policy(model, config.model) if config.model.offload == "full" else None
         ),
     )
     logger.debug(f"Initialized optimizer in {format_time(time.perf_counter() - t0)}")