[faresobeid]

Builds on #4008 (this branch is #4008's head plus two commits, merged with main). With model.offload = "optimizer", AdamW and model.offload_overlap_forward = true, the chunked state copies of the optimizer step run during the next forward, so the step no longer adds its full copy time to every training step.

Opt-in for now: model.offload_overlap_forward defaults to false, which keeps #4008's synchronous chunked step. It stays off until an RL run has exercised the wait before the weight broadcast and it has been tested with sparse weight deltas (#4085).

Why the step was still slow

On 8xH200, the AdamW states are 28.4 GiB per rank for Qwen3-30B-A3B. Each step moves them to the GPU and back, 56.9 GiB per rank. With all 8 GPUs copying at once, each GPU gets 26 GiB/s in one direction or 38.8 GiB/s with both directions at once, and only if its pinned buffers sit on its own socket (the NUMA fix is #4426). Each pair of GPUs shares one PCIe Gen5 x16 uplink (52 GiB/s per direction per pair), and each socket tops out near 155 GiB/s with both directions running. That puts a floor of 1.47 s on the copies, while the GPU math takes about 75 ms. #4008 reaches that floor (1.48 s with #4426), but the GPU still does nothing else for those 1.5 s of every ~3.9 s step.

Change

• optimizer.step() returns as soon as it has queued the first 8 chunks. The rest are queued as the next forward reaches them, staying 4 chunks ahead of it. Queuing all ~33 chunks at once filled the CUDA launch queue and blocked the CPU for about 1 s, which stopped any overlap.
• Chunks update on a third stream. Each FSDP parameter group waits for the update event of the last chunk holding its parameters, on the all-gather streams in unshard and on the current stream in wait_for_unshard. This wraps FSDPParamGroup.unshard / wait_for_unshard per group (FSDP2 internals) and covers each mesh's group, e.g. EP experts. If a parameter is not managed by FSDP, the step logs a warning and runs synchronously as in perf(optim): overlap optimizer state offload copies with the AdamW step #4008.
• The step's hyperparameters and gradients are captured at step time. A chunk queued later sees neither the scheduler's next LR nor zero_grad. Gradients are record_streamed onto the update stream so their memory is not reused early.
• model.offload_overlap_forward (default false, opt-in; AdamW only). wait_for_optimizer_step(optimizer) drains a step in flight. It runs before the weight broadcasts in the SFT and RL loops. Checkpoint save/load and state_dict drain through _move_states.

Numbers

Rows labelled "this PR" are with offload_overlap_forward = true; with the default (false) the step runs as in #4008.

Qwen3-30B-A3B, 1 node 8xH200, FSDP 8 / EP 8, 16k seq, micro batch 1, global batch 8, full AC, compile, fake data. tok/s is the mean of the last 10 of 20 steps; step time is the mean of steps 6-20. "NUMA" means bound to the GPU's socket as in #4426.

offload = "optimizer", no activation offloading:

	tok/s	step	peak GPU
no offload	57.8k	2.26 s	81.3 GiB
main	26.8k	4.90 s	63.2 GiB
#4008	31.8k	4.11 s	63.2 GiB
#4008 + #4154	31.3k	4.19 s	52.8 GiB
#4008, NUMA	33.8k	3.87 s	63.2 GiB
this PR	32.8k	4.00 s	63.6 GiB
this PR, NUMA	37.3k	3.52 s	63.6 GiB
this PR, NUMA, offload_overlap_forward = false (default)	33.8k	3.87 s	63.2 GiB

offload = "optimizer" plus activation offloading (the default config otherwise):

	tok/s	step
main	23.7k	5.53 s
#4008	30.3k	4.33 s
this PR	31.9k	4.11 s
this PR, NUMA	35.0k	3.74 s

Without NUMA binding, results depend on which socket each rank's buffers happened to land on (3 to 5 of 8 ranks were remote in these runs). Host RSS is 41 GiB per rank (45 GiB with activation offload) in every arm. Peak GPU memory rises by 0.4 GiB, because gradients live into the next forward until their chunk has updated. The #4154 row shows that #4154's bounded first step lowers the peak recorded in steady state. That is independent of this change.

Trace of the NUMA run (rank 0, one step): the copies start as soon as step() returns. They run at about 95% of the link while the trainer finishes the step's bookkeeping, then at 70-80% under the forward. The forward stretches from 0.64 s to 1.56 s because it waits for its chunks. Backward is unchanged at 1.54 s.

Validation

• FSDP2 on 2 GPUs, 6 sharded blocks plus root, 1 MiB chunks (one or two parameters each), grad clipping, AdamW with weight decay, 6 steps. Parameters are bitwise equal to plain AdamW for both the synchronous chunked path and the overlapped one.
• Trainer runs above: loss and grad norm per step match perf(optim): overlap optimizer state offload copies with the AdamW step #4008 within run-to-run noise. Runs with identical optimizer math differ by up to 1.5e-3 in loss and 2.7e-2 in grad norm (MoE kernels are not deterministic); this PR vs perf(optim): overlap optimizer state offload copies with the AdamW step #4008 differs by 1.9e-3 and 2.1e-2.
• pytest -m "not gpu" tests/unit/train tests/unit/test_configs.py passes, except test_cpu_adam.py, which needs ninja on the machine I used. ruff check and ruff format --check pass. The RL loop's new wait before its weight broadcast was not exercised at runtime.

🤖 Generated with Claude Code