[faresobeid]

Replaces the three switches model.fsdp_cpu_offload, model.optim_cpu_offload and model.full_offload with one field:

[model]
offload = "optimizer"   # default; also "full", "none"

This applies to both the RL trainer and SFT. The fsdp switch is not migrated: the mode is removed outright (rationale below). The [model.full_offload] numa_bind knob is dropped: no recipe, deployment, template or doc sets it. Full offload now always pins each rank's CPUs to its GPU's NUMA node. That was already the default. Ranks whose GPU reports no NUMA node still skip binding with a warning. On hosts without GPU sysfs topology (/sys/bus/pci/devices/<bus>/numa_node), binding is no longer optional: full offload fails there instead of being switched off.

Why: the three switches could only be used one at a time, enforced by cpu_offload_mutual_exclusion. Because optim_cpu_offload defaulted to true, enabling full offload also meant writing optim_cpu_offload = false, otherwise validation failed (see the old configs/advanced/deepseek-v4-flash/sft.toml comment). With one value the conflicting combinations cannot be expressed, so that validator is deleted. The default behaviour does not change: optimizer-state offload stays on.

Why drop fsdp: the mode was the stock torch FSDP2 CPUOffloadPolicy with no in-tree config, example, or test user; Muon rejected it and it broke NCCL weight broadcast. It was kept only for long-context runs whose fp32 masters/moments no longer fit on GPU (the reason it was restored in #3052); the full mode added in #3234 keeps FP32 masters, moments, and gradients in host RAM, covering that memory pressure instead. Gone with it: the CPUOffloadPolicy selection in setup_fsdp, CPU-side model init in setup_model/load_dcp_from_hf, _move_buffers_to_cuda, and the fsdp leg of the Gloo special case.

Changes:

• ModelConfig.offload: Literal["none", "optimizer", "full"] = "optimizer" replaces the three fields. OptimizerInBackwardOffloadConfig (and with it numa_bind) and its bool-normalizing BeforeValidator are gone.
• Validators in TrainerConfig and SFTConfig:
  • full offload accepts only AdamW or SignSGD and disables clipping, now reading offload == "full";
  • the Muon optimizer does not support FSDP CPU offload validators are deleted with the mode.
• Runtime:
  • setup_optimizer(..., offload=...) replaces cpu_offload + full_offload_config, and its "both enabled" check is gone.
  • FullCPUOffloadOptimizer drops its unused offload_config argument.
  • setup_full_cpu_optimizer_offload() takes no arguments and always binds.
  • model.py: FSDP always uses the default OffloadPolicy(); CPU-side model init in setup_model/load_dcp_from_hf and _move_buffers_to_cuda are gone.
• Tomls: glm-4.5-air/swe-budget.toml and ci/integration/reverse-text-sft/full-offload-resume.toml → offload = "full"; glm-5.3/sft/h200/16k-12-node.toml → offload = "none"; glm-5.3/sft/h200/base.toml → offload = "optimizer". Also updated docs/scaling.md, the deepseek-v4-flash/sft.toml comment, and skills/training/start-run/SKILL.md, which documented the removed fields.

Breaking: the old keys now fail with extra_forbidden, and offload = "fsdp" fails the Literal. There is no shim.

Interaction with #3849 (one trainer-common config base): #3849 moves the full-offload and Muon/FSDP validators from TrainerConfig/SFTConfig onto BaseTrainerConfig. This PR edits the full-offload validator's condition and deletes the Muon/FSDP one, so the conflict is textual only: keep #3849's single copy of the full-offload validators with self.model.offload == "full" and drop its Muon/FSDP copy. Whichever lands second takes a small rebase.

Other open PRs that set or read the old keys need a rebase:

• tomls: perf(dsa): cuDNN indexer + radix top-k for dsa_backend=cudnn_flashmla #3883, feat(trainer): DeepSeek-V4.1 Flash (engram, shared compressed attention, CP) #3879, docs: performance.md with the Qwen3-235B numbers vs torchtitan and the max-MFU benchmark configs #3766, chore(configs): add the Qwen3-235B SFT benchmark config matched to torchtitan #3765, feat: integrate DSpark draft training, weight sync, and checkpoints #3729, feat(dynamo): support python based dynamo vLLM workers administration #3629, feat(deepseek_v4): fp8 RL #3627, feat(inference): configure KV offload per P/D role #3616, feat(deepseek_v4): SFT #3544
• code: perf(trainer): DeepSeek-V4.1 FP8 experts, engram all-to-all EP, pipeline parallelism with one-sided NCCL transport #3885 (optim_cpu_offload), feat(sft): packed multiple-choice SFT with a pluggable loss and choice evals #3739 (full_offload validators), feat(trainer): remove CP token padding and duplicated RL microbatches #3724, fix(optim): stream Muon gradients and state for 8-node GLM-5.3 SFT #3719, Add PPO and SAO value-model training #3699, feat(multimodal): enable Nemotron-H Omni VLM training #3610, feat: add ModelExpress weight updates #3487 (fsdp_cpu_offload / enable_gloo)

Test plan

• pytest tests/unit/orchestrator tests/unit/test_configs.py tests/unit/train/test_cpu_adam.py tests/unit/train/test_state_offload.py -m "not gpu": 336 passed, 1 skipped (Qwen3-VL snapshot not cached on this host), 4 deselected (GPU; this host has no GPU). The existing offload config tests are updated to the new field. test_load_configs covers the 4 edited tomls.
• cli(TrainerConfig|SFTConfig, ...): --model.offload full|none and the default (optimizer) resolve as expected; optim_cpu_offload and offload = "fsdp" are rejected.
• import prime_rl.trainer.rl.train, prime_rl.trainer.sft.train works.
• ruff check and ruff format --check on the changed Python files: clean.
• Not run: GPU or multi-node training with offload = "full".

---

Note

Medium Risk
Breaking TOML/API surface touches distributed training memory paths (FSDP offload, optimizer-in-backward, Gloo); behavior should match prior modes when mapped correctly, but mis-migrated configs will fail at parse time.

Overview
Breaking config change: model.fsdp_cpu_offload, model.optim_cpu_offload, and model.full_offload are replaced by a single model.offload setting ("optimizer" default, or "full", "fsdp", "none"). RL and SFT share the same field; example TOMLs and scaling docs are updated accordingly.

The mutual-exclusion validator between the old flags is removed because only one mode can be expressed. [model.full_offload] / numa_bind is dropped — "full" always runs NUMA binding and CPU optimizer thread setup via parameterless setup_full_cpu_optimizer_offload(). Runtime paths (setup_optimizer, FSDP CPU offload in model.py, Gloo when offload is "full" or "fsdp") key off the new enum; FullCPUOffloadOptimizer no longer takes an offload config object.

Default training behavior is unchanged: optimizer-state CPU offload remains the default ("optimizer"). Old TOML keys are not aliased and should fail validation.

^{Reviewed by Cursor Bugbot for commit 35da73a. Bugbot is set up for automated code reviews on this repo. Configure here.}

[mikasenghaas]

nice ya always found our naming confusing before. one thing: couldn't we compose fsdp+optim offload before, which would not be possible anymore no?