feat: overlap CPU optimizer with backward

[samsja]

Summary

TL;DR — the trade-off: full optimizer offload cuts peak GPU memory to a third–half (Qwen3-30B at 64K: 90 → 45 GiB) and is what makes models trainable when optimizer state exceeds HBM (GLM-5-class: 257 GiB/GPU of state). The price is throughput at short sequences: vs a true GPU-resident baseline it keeps 44% of TPS at 8K per-GPU tokens, 90% at 16K, and is parity within run noise from ~32K up. At production sequence lengths the offload is effectively free; below ~16K you're trading real speed for memory.

Supersedes #3212 (closed; this PR targets main and carries the combined diff). Turns CPU optimizer offload into an optimizer-in-backward wavefront: each CPU optimizer chunk runs as soon as its last gradient reaches host memory, and the refreshed BF16 weights stream back while backward is still executing. Gradients and weight refreshes cross PCIe in BF16; masters, moments, and all optimizer arithmetic stay FP32. A native multi-tensor CPU AdamW kernel (AVX-512/AVX2/ATen dispatch) replaces fused torch.optim.AdamW by default (cpu_optimizer_backend = "torch" retained for parity checks). Pinned memory is bounded by size-classed transfer rings instead of model-sized slabs.

Breaking simplification: optim_cpu_offload now means exactly one thing — full offload — and is disabled by default (the previous default silently ran state-only offload, costing ~45% throughput at seq 8K vs a GPU-resident baseline; that mode and gradients mode are removed along with their streaming machinery, ~250 lines). Enable with model.optim_cpu_offload = true. numa_bind defaults on. Validated: config suite, DCP checkpoint save + resume with exact loss parity, native and torch backends, and the no-offload default.

Also in this PR: per-step pipeline timing diagnostics (debug level), an automatic per-rank CPU thread budget at startup (#3248, merged — the launcher default OMP_NUM_THREADS=1 would serialize the CPU optimizer), and random-init debug runs no longer pre-download weights. Launch configs in benchmarks/offload/; scaling docs updated in docs/scaling.md.

Gradient numerics

Full offload is not gradient-bit-identical to no-offload, by construction: gradients reduce across ranks in FP32, but FSDP2 materializes .grad in the sharded parameter's dtype — BF16 for the full-offload compute model — so each reduced gradient is rounded to BF16 once before the FP32 CPU update (masters, moments, accumulation, and Adam arithmetic stay FP32; the BF16 weight refresh is exact). An FP32-transport toggle is structurally blocked (FSDP2 has no grad-dtype override on BF16 sharded params); optimizer-state-only offload is the gradient-bit-faithful mode. Documented in the config, the training skill, and handover.md. Paired runs matched losses to display precision.

Results (one 8xH200 node, fake data, random init, losses identical in every comparison)

Qwen3-30B-A3B, paired same-session runs, baseline = true no-offload (optimizer states resident on GPU, --model.optim-cpu-offload None; note the repo default silently enables state-only offload, a mode this PR has since removed):

Seq	No offload (GPU-resident)	Full offload (native)	TPS retained	Peak HBM
8K	1.75 s	3.96 s	44%	63.1 → 17.5 GiB
16K	5.17 s	5.73 s	90%	67.0 → 21.5 GiB
32K	16–21 s (noisy)	15.4–16.7 s	~100%	73.7 → 29.0 GiB
64K (ac)	38.7 s	34.1 s	~parity¹	90.3 → 44.9 GiB

¹ Single unprofiled runs; run-to-run variance on this node is ±15–20% at these lengths, and paired profiler traces at 64K show identical backward GPU work and sub-second tails for both modes — treat 32K–64K as parity, not a win for either side.

Against the real baseline, full offload costs throughput below ~16K per-GPU sequence and reaches parity from roughly 16–32K upward — at a third to half of the resident-state HBM footprint throughout, and it remains the only mode that runs when optimizer state cannot fit on the GPU.

GLM-5 truncated to 8 layers (53.2B) — the largest size where resident optimizer state fits an H200 at all (pre-#3248/#3249 numbers; both arms carry the respective overheads):

Seq	No offload	Full native	TPS retained	Peak HBM
16K	2.1 s / 28.0k TPS	6.0 s / 10.7k TPS	38%	122.9 → 53.9 GiB
32K	7.7 s / 15.9k TPS	7.1 s / 18.6k TPS	117%	134.1 → 67.0 GiB
64K	OOM	10.7 s / 24.5k TPS	—	91.6 GiB

GLM-5 20-layer proxy (171B, CP2/EP8, no-offload impossible at 257 GiB/GPU of state; 16K remeasured with #3248):

Seq	Step med	TPS	Peak HBM	Exposed drain
16K	19.4 s	3.4k	102 GiB	12.6 s
32K	23.1 s	5.8k	115 GiB	11.8 s
64K	37.2 s	7.3k	134 GiB	6.9 s

The durable wins: roughly half the peak HBM at every point, and training models whose optimizer state cannot fit on GPU at all. For a 3.8B-param rank shard the GPU optimizer step itself measures ~32 ms; the no-offload post-backward tail is bookkeeping (~1.3-2.1 s in these runs, of which the metric removed in #3249 accounts for ~0.3 s at 8K). Where full offload wins (64K), the gain is split between a faster backward — gradients leave the GPU as BF16 instead of accumulating as FP32 shards, and the allocator runs at 45 instead of 59 GiB — and a smaller tail.

Validation

ruff clean; native-kernel unit test passes on AVX-512 hosts; five-step native-vs-torch parity, BF16-transport bitwise equality, DCP save/resume, accumulation-8 stress, and synchronous-validation smokes all pass (details in prior revisions and handover.md).

🤖 Generated with Claude Code