Revert "chore: remove fsdp_cpu_offload"

[mikasenghaas]

Reverts #3049 — decision reversed after measuring the 2-node ceilings for Nemotron Super at cp=8 without it:

config	fsdp_cpu_offload	verdict
SFT 131k	on	79.8/139.8 GiB (57%) — comfortable
SFT 196k	on	114.8 GiB (82%), 37.5% MFU
SFT 131k	off	OOM (demand ~142+ GiB; fp32 param+grad baseline stays on GPU)
RL 131k	off	OOM by 15–20 MiB (with expandable_segments)

Without FSDP param offload, fp32 master weights+grads (~60 GiB/GPU on 2 nodes) cap trainable context at ~96–120k; with it, 131k–196k fit. Restoring the knob; validation runs at 131k with it re-enabled are in flight (results to follow).

Note: combining fsdp_cpu_offload with NCCL weight broadcast still needs the staging fix from #3045 (closed; can be reopened) — trainer-only / filesystem-broadcast paths are unaffected.

🤖 Generated with Claude Code

---

Note

Medium Risk
Touches core FSDP sharding, weight loading, and distributed backends; misconfiguration is guarded but throughput and NCCL weight broadcast interactions remain sensitive.

Overview
Restores the trainer.model.fsdp_cpu_offload option (reverting its removal) so FSDP can offload parameters, gradients, and optimizer state to CPU with pinned memory—needed for long-context SFT/RL on tight GPU memory.

Trainer wiring: setup_fsdp applies CPUOffloadPolicy across FSDP units; weight load / to_empty uses CPU when offload is on, with _move_buffers_to_cuda because FSDP does not offload buffers. RL and SFT call setup_torch_distributed(..., enable_gloo=True) when offload is enabled (cpu:gloo,cuda:nccl).

Config/docs: fsdp_cpu_offload is mutually exclusive with optim_cpu_offload; Muon rejects FSDP CPU offload. docs/scaling.md documents the knob and tradeoffs.

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

[mikasenghaas]

Validation complete: with fsdp_cpu_offload restored, Nemotron Super @131k/cp=8 on 2×H200 runs at 79.8/139.8 GiB (57%) peak — SFT 22.8% MFU, RL trainer 30.8% MFU, 4/4 steps each (research-prod PR #155 debug configs). Without it the same configs OOM (15–336 MiB short at 131k; ~140+ GiB demand at 96–120k).