diff --git a/configs/advanced/deepseek-v4-flash/sft.toml b/configs/advanced/deepseek-v4-flash/sft.toml index 69818b7f0d4..928a6c76c16 100644 --- a/configs/advanced/deepseek-v4-flash/sft.toml +++ b/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 diff --git a/configs/ci/integration/reverse-text-sft/full-offload-resume.toml b/configs/ci/integration/reverse-text-sft/full-offload-resume.toml index db2e824c2c6..9dd676b859e 100644 --- a/configs/ci/integration/reverse-text-sft/full-offload-resume.toml +++ b/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 diff --git a/docs/scaling.md b/docs/scaling.md index f80f3751865..fde8f861acd 100644 --- a/docs/scaling.md +++ b/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; `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 diff --git a/examples/advanced/glm-4.5-air/swe-budget.toml b/examples/advanced/glm-4.5-air/swe-budget.toml index 82c67aa022f..a2bf5323630 100644 --- a/examples/advanced/glm-4.5-air/swe-budget.toml +++ b/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] diff --git a/examples/advanced/glm-5.3/sft/h200/16k-12-node.toml b/examples/advanced/glm-5.3/sft/h200/16k-12-node.toml index f72b0b16352..f6c738457e7 100644 --- a/examples/advanced/glm-5.3/sft/h200/16k-12-node.toml +++ b/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] diff --git a/examples/advanced/glm-5.3/sft/h200/base.toml b/examples/advanced/glm-5.3/sft/h200/base.toml index 125e88b4df8..8561b34139c 100644 --- a/examples/advanced/glm-5.3/sft/h200/base.toml +++ b/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] 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 e023ee00087..88c97861b44 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/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).", @@ -621,12 +621,6 @@ def validate_scheduler_steps(self): validate_scheduler(self.scheduler, self.max_steps) 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 - ### Auto-setup and validate shared configs @model_validator(mode="after") 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 9d5d878af50..c2690e8a873 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/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"] 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. ``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. ``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).", @@ -751,12 +707,6 @@ def validate_scheduler_steps(self): validate_scheduler(self.scheduler, self.max_steps) 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_lora_broadcast(self): if self.model.lora is not None and self.weight_broadcast.type in ("nccl", "nixl"): diff --git a/skills/training/start-run/SKILL.md b/skills/training/start-run/SKILL.md index 2b372426383..09c78e803e0 100644 --- a/skills/training/start-run/SKILL.md +++ b/skills/training/start-run/SKILL.md @@ -26,12 +26,11 @@ SLURM launches write generated scripts and coordination files under `/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 diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index a2c33e654f7..9f4d94eaa8b 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -19,7 +19,7 @@ from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointWrapper from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageReader from torch.distributed.checkpoint.state_dict_loader import load as dcp_load -from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, OffloadPolicy, fully_shard +from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, OffloadPolicy, fully_shard from torch.distributed.fsdp._fully_shard._fsdp_common import FSDPMeshInfo, ShardPlacementResult from torch.distributed.fsdp._fully_shard._fsdp_init import _get_mesh_info from torch.distributed.tensor import Shard @@ -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 = 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 = "cuda" model.to_empty(device=device) torch.distributed.barrier() @@ -705,7 +705,6 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: Paral logger = get_logger() if config.debug.random_init: logger.warning("Randomly initializing model. Skipping loading weights from HF.") - _move_buffers_to_cuda(model, config) return if not Path(config.name).exists(): @@ -775,7 +774,6 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: Paral storage_reader=HuggingFaceStorageReader(path=snapshot_path.as_posix()), ) write_back_loaded_packed_parameters(model, state_dict) - _move_buffers_to_cuda(model, config) lora_modules = [m for m in model.modules() if hasattr(m, "_init_lora_parameters")] if lora_modules: @@ -891,15 +889,6 @@ def configure_trainable_parameters(model: nn.Module, config: ModelConfig) -> nn. return frozen_vision_encoder -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) and module.tokens_per_expert.device.type != "meta": @@ -1024,11 +1013,10 @@ 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 = "cuda" model.to_empty(device=device) torch.distributed.barrier() model.init_buffers_post_meta() - _move_buffers_to_cuda(model, config) else: load_dcp_from_hf(model, config, parallel_dims) diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py index 2bfe189405f..a50bd64233b 100644 --- a/src/prime_rl/trainer/optim/__init__.py +++ b/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, ) diff --git a/src/prime_rl/trainer/optim/offload.py b/src/prime_rl/trainer/optim/offload.py index cb5a692c020..f0b698c9939 100644 --- a/src/prime_rl/trainer/optim/offload.py +++ b/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() diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 7819d92a324..6708578a6ba 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/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 == "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)}") diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index bfb070ded93..98727b861b9 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -100,10 +100,10 @@ def train(config: SFTConfig): # 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 == "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. @@ -162,11 +162,10 @@ def train(config: SFTConfig): 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 ), ) diff --git a/src/prime_rl/trainer/utils.py b/src/prime_rl/trainer/utils.py index a17bb373a66..d9fc66d3d0b 100644 --- a/src/prime_rl/trainer/utils.py +++ b/src/prime_rl/trainer/utils.py @@ -21,7 +21,6 @@ from prime_rl.utils.pathing import get_ckpt_dir if TYPE_CHECKING: - from prime_rl.configs.trainer import OptimizerInBackwardOffloadConfig from prime_rl.trainer.optim import GradientOffloadManager DEFAULT_TIMEOUT = timedelta(seconds=600) @@ -176,9 +175,8 @@ def configure_cpu_optimizer_threads() -> None: ) -def setup_full_cpu_optimizer_offload(config: "OptimizerInBackwardOffloadConfig") -> None: - if config.numa_bind: - bind_process_to_gpu_numa_node() +def setup_full_cpu_optimizer_offload() -> None: + bind_process_to_gpu_numa_node() configure_cpu_optimizer_threads() diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 9bcba09496a..f38f21b1678 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -331,8 +331,7 @@ def test_removed_moe_runtime_fields_are_rejected(removed): def test_optimizer_state_offload_keeps_legacy_default(config_cls): config = config_cls.model_validate({}) - assert config.model.optim_cpu_offload is True - assert config.model.full_offload is None + assert config.model.offload == "optimizer" def test_moe_router_dtype_auto_resolves_per_trainer(): @@ -360,7 +359,7 @@ def test_full_optimizer_offload_disables_gradient_clipping(config_cls): with pytest.warns(UserWarning, match="Gradient clipping prevents optimizer-in-backward"): config = config_cls.model_validate( { - "model": {"optim_cpu_offload": False, "full_offload": True}, + "model": {"offload": "full"}, "optim": {"max_norm": 1.0}, } ) @@ -373,7 +372,7 @@ def test_full_optimizer_offload_requires_supported_optimizer(config_cls): with pytest.raises(ValidationError, match="Full optimizer offload only supports AdamW and SignSGD"): config_cls.model_validate( { - "model": {"optim_cpu_offload": False, "full_offload": True}, + "model": {"offload": "full"}, "optim": {"type": "muon", "max_norm": None}, } ) @@ -383,11 +382,11 @@ def test_full_optimizer_offload_requires_supported_optimizer(config_cls): def test_full_optimizer_offload_accepts_sign_sgd(config_cls): config = config_cls.model_validate( { - "model": {"optim_cpu_offload": False, "full_offload": True}, + "model": {"offload": "full"}, "optim": {"type": "sign_sgd", "max_norm": None}, } ) - assert config.model.full_offload is not None + assert config.model.offload == "full" assert config.optim.type == "sign_sgd"