[faresobeid]

Stacked on #4085 (sparse filesystem weight broadcast). With model.offload = "full" (AdamW, sign_sgd), the RL trainer takes the sparse delta from the optimizer step itself, instead of diffing its weights against a snapshot of the previous broadcast.

How it works

• Hook. FullCPUOffloadOptimizer gets an optional weight_hook(local_param, source), called on its H2D stream in place of local_param.copy_(source). At that point the local shard still holds the previous broadcast's weights, and source is the chunk's new values in the pinned output slot.
• Diff. The sparse sender sets the hook before each step's final backward (wait_for_weight_reads). For every trainer tensor stored in a parameter, the hook copies the old and new values into one of two GPU slots, writes the new values to the parameter, and queues the slot to a diff thread. The diff thread runs feat(weights): sparse XOR deltas for the filesystem weight broadcast #4085's diff code on the slot: bf16 raw-bit differences, or for fp8 engines the AOTI quantizer on old and new plus _scale_inv records, then zigzag, index gaps and checksums of the new values. The GPU parts are held until the step ends. A tensor may be a strided view of a parameter (for example gate_proj/up_proj of the packed gate_up_proj).
• Publish. broadcast() hands the parts to the writer (D2H, part files, manifest, markers). The delta format, the engine apply path and the per-step checksums are unchanged.
• Guard. The encode recomputes a position-weighted checksum of the live weights of every hooked tensor and compares it with the values the step wrote. If a weight changed outside the optimizer step, or a tensor was never written by it, the broadcast raises before publishing. Engines verify against the checksums of the written values.
• Anchors. At anchor steps the hook pwrites the new values from the pinned slot into the anchor files at precomputed offsets. Slices go round-robin to the ~1 GiB files, and each parameter's slices are written by 8 threads. One file is written at ~0.5 GiB/s on /home; 4 files at once reach 1.9 GiB/s. The files are written next to the step dir and moved in when the step is published, so an anchor holds exactly one version. There is no host buffer and no GPU-to-host anchor copy.
• Residual tensors. A small snapshot (feat(weights): sparse XOR deltas for the filesystem weight broadcast #4085's path) is kept for tensors that the hook cannot diff:
  • state-dict entries that are not stored in an optimizer parameter: buffers, and copies such as q/k/v_proj and router.gate on Qwen3-30B (0.94 GiB per rank);
  • fp8 tensors whose 128-row blocks straddle FSDP shards, which need all-gathers (1.07 GiB per rank on 30B fp8, including those).
• What goes away for full offload:
  • the 2 B/param snapshot and the host anchor buffer (10.5 GiB per rank each on 30B; 21.6 GiB per rank on GLM-5.3);
  • the startup prefault and the anchor-copy skip;
  • the v1 read-back of the boot checkpoint for hooked tensors (v1 is diffed against the live v0 weights).
• Other offload modes (optimizer, none, Muon) and SFT keep feat(weights): sparse XOR deltas for the filesystem weight broadcast #4085's snapshot path unchanged. Hooked deltas need a broadcast after every optimizer step: a broadcast for a step whose pass was not hooked raises.

Results

Qwen3-30B-A3B-Instruct-2507, H200, 1 node (6 trainer GPUs with offload = "full", engine TP2), reverse-text, 20 steps, anchors every 3 steps. Each pair ran on the same node. Medians over steps 3-20. Every delta was verified by the engines in every run (40/40, no divergence).

	bf16 engine, AdamW: #4085	this PR	fp8 engine (DeepGEMM), sign_sgd: #4085	this PR
trainer GPU peak (max over steps)	40.3 GiB	31.7 GiB	38.0 GiB	33.0 GiB
node host memory used (max)	622 GB	576 GB	404 GB	352 GB
v1: trainer wait before the final backward	8.32 s	0.03 s	9.48 s	0.25 s
fwd/bwd, normal steps	8.68 s	9.16 s	7.08 s	7.88 s
fwd/bwd, anchor steps	8.32 s	18.55 s	6.79 s	16.31 s
optimizer drain after backward	2.65 s	3.6 s	1.26 s	2.51 s
step	12.8 s	10.79 s	12.4 s	8.49 s
delta / step	335 MiB	332 MiB	134 MiB	134 MiB
delta committed after the broadcast call		2.03 s		2.67 s
orchestrator prefetch / update	1.39 / 3.15 s	2.3 / 3.95 s	0.45 / 2.6 s	1.05 / 3.95 s
engine apply / verify	0.23 / 0.3 s	0.19 / 0.33 s	0.27 / 0.31 s	0.24 / 0.34 s

• Normal steps. The work moves from the encode after the step into the optimizer pass. This adds 0.5-0.8 s to fwd/bwd on normal steps: the diff of the last chunks finishes after the backward.
• Orchestrator update. It is 0.8-1.4 s longer: the whole delta is copied to the host and written after the step, while feat(weights): sparse XOR deltas for the filesystem weight broadcast #4085 writes parts as its encode produces them. Step time is rollout-bound and noisy here.
• Anchor steps. Writing the anchor during the pass is bound by storage (9.5 GiB per rank, 57 GiB per node). It adds 8.5-10 s to fwd/bwd per anchor step. feat(weights): sparse XOR deltas for the filesystem weight broadcast #4085 instead writes anchors from a host copy in the background in 25-29 s.
• Fresh-engine checks (on this PR's anchors, both engines):
  • A fresh engine joins at a non-anchor step: it loads the step_18 anchor written by the optimizer pass, then applies and verifies v19.
  • Fail-then-retry heals.
  • A double-applied delta and a delta shifted by one element both fail the checksum.
• Guard. The guard caught one real case during development: the state-dict entries of a packed parameter are strided views, and the hook first read them as contiguous. They are now mapped through their strides.