[faresobeid]

Lands after #4427: this branch merges it (both add a model.offload_* field next to each other) and current main.

With model.offload = "full", every backward hands each parameter's BF16 gradient to GradientOffloadManager, which copies it to the CPU through 4 pinned slots and adds it into an FP32 accumulator on a CPU thread. With N micro steps per step that is N full gradient copies over PCIe and N-1 CPU adds. The copies share host bandwidth with activation offloading and weight copies; when they fall behind, gradients wait on the GPU for a free slot.

model.offload_grad_accumulation = "gpu" (opt-in, default "cpu") sums the gradients of the non-final micro steps in FP32 on the GPU, then sends the sum through the same transfer path once per step.

• Non-final backwards: the first gradient is cast to an FP32 GPU tensor, later ones are added to it; param.grad is released as before.
• Final backward: the gradient is added and the FP32 sum is queued (the manager also keeps FP32 transfer slots). A parameter with no gradient in the final backward has its accumulator queued in finish_backward. The CPU accumulator ends up with the same FP32 sum, in the same order, as the CPU path.
• Cost: 4 bytes of GPU memory per local trainable parameter, from the first backward until the final copy; nothing is kept between steps.

Validation (branch merged with main and #4427, f829bc9)

Numerics

• 2 GPUs, FSDP2 + AdamW, 3 steps x 3 micro steps, including a parameter that gets a gradient only in the first micro step: the CPU masters and GPU parameters are bitwise equal to cpu accumulation.
• SFT reverse-text, Qwen3-0.6B, 1 node 8xH200, offload = "full", AdamW, 4 micro steps per rank, 10 steps, run through the ported SFT loop (trainer/train.py). The loss is identical at every step: cpu 6.229171 → 5.484061 in both of its runs, gpu 6.229171 → 5.484061. Full offload turns off gradient clipping, so there is no grad norm to compare. Steady step time 1.35 / 1.30 s → 1.20 s; peak memory 4.79 → 5.07 GiB (+4 bytes per local parameter).

Performance

Measured on the merged branch (one 8×H200 node, full offload, steady steps 2–5):

micro steps	accumulation	step time	throughput	peak memory
8	CPU (default)	52.1 s	20.4k tok/s	122.7 GiB
8	GPU (offload_grad_accumulation)	47.7 s (−8.5%)	22.2k tok/s	98.1 GiB
4	CPU (default)	27.2 s	19.6k tok/s	65.5 GiB
4	GPU	26.8 s (−1.5%)	19.9k tok/s	98.1 GiB

• Per-step loss and mismatch KL are identical between CPU and GPU accumulation in every arm.

• The 2-GPU bitwise check passes: master weights and parameters are bitwise equal on both ranks.

• GPU accumulation pins one FP32 accumulator per parameter, so peak memory is flat at 98.1 GiB. CPU accumulation's peak grows with micro steps and is the larger of the two at 8.

• The gain shows at 8 micro steps (−8.5% step time and less peak memory) and is small at 4.

• tests/unit/train, tests/unit/utils, tests/unit/test_configs.py (-m "not gpu"): all pass except test_cpu_adam, which also fails on main in this environment (no ninja on PATH). ruff is clean.

---

Note

High Risk
Changes gradient pipelines, FSDP unshard synchronization, and optimizer stepping timing—areas that affect training correctness and can race with weight broadcast if waits are wrong.

Overview
Adds two optimizer offload tuning knobs on ModelConfig: offload_grad_accumulation ("cpu" default, "gpu" only with **offload = "full") and offload_overlap_forward (default false, AdamW + offload = "optimizer").

For full offload, GradientOffloadManager can sum micro-batch gradients in FP32 on GPU and copy once per step, trading a little GPU memory for less host PCIe traffic. For state-only offload, CPUOffloadOptimizer uses a new _ChunkedStep path that streams AdamW moment states through staged GPU buffers instead of moving all state for each step; with overlap_model, chunks update on a side stream while the next forward runs, with FSDP unshard hooks waiting only for chunks that own each module’s parameters. Trainers call wait_for_optimizer_step before weight broadcasts so overlapped steps finish before empty_cache / gathers.

Docs in scaling.md describe chunked AdamW and offload_overlap_forward. A unit test checks late-first-gradient parameters match the non-offloaded AdamW reference.

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

🤖 Generated with Claude Code