[mikasenghaas]

Supersedes #3045 — rather than fixing fsdp_cpu_offload for NCCL weight broadcast, remove the knob entirely.

Why

• Broken with NCCL weight broadcast (RL): gathered state-dict tensors stay on CPU and broadcast_state_dict passes them straight to the NCCL communicator → AssertionError: this nccl communicator is created to work on cuda:0, but the input tensor is on cpu (hit during bring-up of a 196k Nemotron Super run).
• Muon rejects it (both trainer + SFT validators existed just to guard this).
• Mutually exclusive with optim_cpu_offload, which remains the supported CPU-offload knob.
• No config in the repo uses it.

What's removed

• ModelConfig.fsdp_cpu_offload + the mutual-exclusion validator + both muon validators
• CPUOffloadPolicy selection in setup_fsdp (behavior unchanged: fully_shard's own default is OffloadPolicy())
• _move_buffers_to_cuda helper + 4 call sites; CPU to_empty device forks
• enable_gloo plumbing in setup_torch_distributed (existed solely for CPU offload)
• Docs (scaling.md)

Mechanical removal; default path byte-identical. Note for the record: FSDP CPU offload was worth ~19–25 GiB/GPU of param baseline in fp32 long-context SFT measurements — if that headroom is ever needed again, the NCCL fix from #3045 is the alternative.

🤖 Generated with Claude Code

---

Note

Medium Risk
Touches core distributed init and FSDP sharding defaults; behavior changes only if configs still set the removed flag, but that path was broken with NCCL weight broadcast.

Overview
Removes the fsdp_cpu_offload training knob end-to-end so FSDP always uses the default in-GPU offload policy and optim_cpu_offload remains the only supported CPU memory lever.

Config & validation: Drops ModelConfig.fsdp_cpu_offload and related validators (mutual exclusion with optim_cpu_offload, Muon + FSDP CPU offload checks in trainer and SFT configs). Docs in scaling.md no longer document FSDP CPU offload.

Trainer runtime: setup_fsdp no longer passes CPUOffloadPolicy; weight init paths always use to_empty(device="cuda") and the _move_buffers_to_cuda helper is deleted. RL and SFT call setup_torch_distributed without enable_gloo (Gloo was only for CPU-offload collectives).

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

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 65e5d3a. Configure here.}

[mikasenghaas]

Good catch — confirmed real: three fully_shard call sites passed offload_policy=offload_policy directly (not via fsdp_config) and would NameError in setup_fsdp. Fixed by dropping the kwargs (fully_shard defaults to OffloadPolicy(), behavior-identical) and verified by importing the trainer modules + checking setup_fsdp's symbol table.