feat: disk-delta weight sync for non-colocated rollout engines

[nanjiangwill]

ci-sglang-pr: #30366

Adds --update-weight-transfer-mode disk-delta: weight sync for non-colocated rollout engines that ships only the bytes changed between syncs — for cross-cluster/datacenter disaggregation where moving the full model every sync dominates.

Engine-side /pull_weights: sgl-project/sglang#30366 (sglang-miles), sgl-project/sglang#30367 (main). Refs: THUDM/slime#1806, THUDM/slime#2089, THUDM/slime#2181.

Design

• First update_weights() publishes nothing: it seeds a CPU snapshot from --hf-checkpoint's safetensors bytes (= the engines' actual base, robust to non-byte-exact megatron→HF round-trips) and issues pull_weights(0) so each host warms its local checkpoint.
• Each sync: source ranks diff gathered HF tensors against the snapshot (pinned-buffer GPU→CPU copies pipelined into a diff+zstd thread pool) and publish weight_v{N:06d}/ under --update-weight-disk-dir as a canonical HF checkpoint dir — compressed-diff shards + an index carrying version/encoding/checksum metadata, coordinated over gloo rather than the filesystem.
• Engines pull_weights(N) before pausing (disk-only, overlaps rollout): per-host flock, in-place apply, per-tensor checksum verify, fail-loud. Then pause → flush → vanilla update_weights_from_disk(local_dir, weight_version=N) → continue.
• PP-replicated params (word embedding on an MTP-hosting or tied-embedding last stage) are deduped at publish time — name→checksum maps gathered, divergent duplicates raise, lowest rank keeps — since the XOR apply is an involution (applied twice = reverted).
• Built on DistBucketedWeightUpdateMixin's callback gather; no NCCL groups or engine lock; density/wire metrics drain via pop_metrics() into the step log.

Args

--update-weight-transfer-mode disk-delta
--update-weight-disk-dir ...              # shared FS: trainer publishes, hosts read
--update-weight-local-checkpoint-dir ...  # host-local full checkpoint, patched in place
--update-weight-delta-encoding xor        # or: overwrite (idempotent, larger)
--update-weight-delta-checksum xxh3-128   # or: blake3, adler32
--custom-update-weight-post-write-path    # non-POSIX FS publish hook

Read-side counterpart: --sglang-custom-pull-weights-pre-read-hook. New deps: xxhash, zstandard, blake3.

Validation

GLM-4.7-Flash GRPO on 4× H200: 2-node actor (TP2/PP2/CP2/EP8, MTP on → exercises the dedup) + one tp16/dp16 engine spanning 2 rollout nodes (exercises the multi-host pull fan-out). Disk-delta publishes to an object-store-backed shared volume via the post-write/pre-read hooks (xor + xxh3-128); the baseline is the same config with the default --update-weight-transfer-mode broadcast (full-checkpoint NCCL each sync).

One sync per training step: vN is the model after optimizer step N, published and applied before rollout N+1. "baseline" is the startup update_weights() before the first rollout: disk-delta ships nothing (the trainer captures its diff-base snapshot while each host copies the engine's base checkpoint to local disk); broadcast pushes the full weights.

disk-delta vs full-checkpoint NCCL broadcast

	broadcast (default)	disk-delta
payload per sync	62.4 GB (full bf16)	0.69–0.83 GB (−98.8%)
update_weights time	5.9 s first, 3.5 s steady	54.8 s first, 23.2 s steady
generation-paused window	the whole sync (3.5–5.9 s; engines pause before the broadcast)	3–5 s — just the update_weights_from_disk reload; every other phase runs while the engine serves
effective transfer rate	10.7–17.7 GB/s through gather→convert→NCCL (RDMA)	≥2.1 GB/s per host volume read (62.4 GB base seed in ≤30 s, hosts in parallel); per-sync wire is ~1% of full
connectivity required	NCCL reachability trainer↔every engine GPU	a shared filesystem / object store only

Broadcast wins on raw sync latency inside one datacenter; disk-delta moves ~80× fewer bytes with a comparable generation-paused window and runs where NCCL can't (cross-cluster/datacenter, object-store-only links).

disk-delta per-sync detail

Sync	Density	Wire	update_weights time
baseline (pre-rollout 1)	—	—	48.7 s (snapshot; 62.4 GB host seed overlapped)
v1 (after step 1)	0.37%	0.69 GB	54.8 s
v2 (after step 2)	0.44%	0.79 GB	23.2 s
v3 (after step 3)	0.44%	0.78 GB	— (run ends after final sync)

Steady-state breakdown of the 23.2 s (from engine request timestamps):

Phase	Time	Engine
trainer: diff 62.4 GB vs snapshot + zstd + publish shards to shared volume	~19 s (combined with next row)	up, generating
engine hosts: pull, decompress, checksum-verify, apply into local checkpoint	″ (completes before pause)	up, generating
pause + KV flush	~1 s	paused
update_weights_from_disk: local disk → GPU	3–5 s	paused
resume	<1 s	—

The only phase that needs the engine stopped is the final update_weights_from_disk reload — in a fully async rollout loop the engine keeps serving through diff/publish/pull, so the effective generation pause per sync is just the 3–5 s reload. All pulls checksum-verified on both engine hosts; the index holds each tensor exactly once.

Caveats

• The baseline call publishes nothing, so a run resumed mid-training serves base weights for its first rollout (engines assumed to start from the trainer's HF checkpoint).
• CI: compatible with --check-weight-update-equal (when the checker is armed, the baseline also reloads the pulled v0 checkpoint — the checker resets engine tensors and expects the first sync to rewrite everything). tests/e2e/megatron/test_qwen3_4B_disk_delta.py is registered under run-ci-weight-update, disabled until the CI image ships sglang /pull_weights ([RL] Add /pull_weights: engine-side pull of published weights into a host-local checkpoint (sglang-miles) sgl-project/sglang#30366).
• LoRA and PD-disaggregation unsupported (validated).

[gemini-code-assist]

Code Review

This pull request introduces Delta Weight Sync, a sparse, non-colocated weight synchronization mode that sends only changed weight positions and values instead of broadcasting full parameters at each step. It supports both NCCL and disk-based transports, along with multiple encoding strategies (indices, deltas, and zstd compression). The code review identified three critical or high-severity issues in the delta sync implementation: a runtime AttributeError due to the use of a non-existent torch.hash_tensor API, a potential distributed deadlock in _finalize_sync caused by conditionally calling the collective _publish_batch operation, and a potential RuntimeError when calling .view() on non-contiguous tensors in _bytewise_diff_mask.

[yueming-yuan]

qq, is this compatible with --check-weight-update-equal?

[nanjiangwill]

qq, is this compatible with --check-weight-update-equal?

Not yet. --check-weight-update-equal (snapshot → reset_tensors → compare) assumes the first sync rewrites every engine tensor - but disk-delta's first update_weights() only seeds the trainer-side baseline (engines untouched), and later syncs write only the changed bytes, so the compare would fail.

Integrity is covered differently here: every sync, each host verifies a per-tensor xxh3-128 checksum of the applied bytes against the trainer's export and fails loud on mismatch.

If we want CI compatibility: when the checker is on, make the baseline also full-reload the pulled v0 checkpoint via update_weights_from_disk

[yueming-yuan]

Note for future maintenance: (1) unify the checksum mechanism in miles repo (2) when refactor the weight update logic, unify some logic here (e.g. lifecycle, pop metrics...)

[nanjiangwill]

Update: --check-weight-update-equal is now supported — when the checker is armed, the baseline also reloads the pulled v0 checkpoint (the checker resets engine tensors and expects the first sync to rewrite everything). Added test_qwen3_4B_disk_delta.py under run-ci-weight-update, disabled until the CI image ships sglang /pull_weights (sgl-project/sglang#30366).