Skip to content

[2/n] Disaggregated rollout: disk-level delta weight sync - #2089

Merged
zhuzilin merged 10 commits into
THUDM:mainfrom
modal-projects:disk-delta-weight-sync
Jul 2, 2026
Merged

[2/n] Disaggregated rollout: disk-level delta weight sync#2089
zhuzilin merged 10 commits into
THUDM:mainfrom
modal-projects:disk-delta-weight-sync

Conversation

@nanjiangwill

@nanjiangwill nanjiangwill commented Jun 16, 2026

Copy link
Copy Markdown
Collaborator

Series:

  1. [1/n] Disaggregated rollout: delta weight sync (disk + NCCL transports) #1806
  2. [2/n] Disaggregated rollout: disk-level delta weight sync #2089
  3. [3/n] Disaggregated rollout: engine-side /pull_weights #2181

Ship only the changed bytes between weight syncs instead of a full checkpoint, for non-colocated training/inference across clusters. The trainer publishes a per-tensor delta as a canonical HF checkpoint directory; each rollout host applies it in place and reloads via the ordinary update_weights_from_disk path — no delta-specific engine code. Replaces the NCCL delta transport from #1806.

The delta is computed on raw tensor bytes and the engine just reloads a standard checkpoint, so the rollout side stays decoupled from the trainer and is free to use:

  • any low precision — int4, nvfp4, mxfp8, fp8-block
  • any attention/moe backend
  • any parallelism scheme

Future work

  • Apply the delta during the engine's weight load instead of patching a host-local checkpoint first (minor sglang change).
  • Read only the changed tensors via the overwrite encoding instead of reloading the full checkpoint (larger sglang refactor).

Ship only the changed bytes between weight syncs as a canonical HF delta
checkpoint; rollout hosts apply it into a host-local checkpoint and reload via
the vanilla update_weights_from_disk path. Replaces the NCCL delta transport
from THUDM#1806 with a disk-only path that needs no engine-side delta support.
sync_local_checkpoint (was sync_weights) materializes the base lazily via the
idempotent init_local_checkpoint instead of a background thread; record per-sync
update time in update_weight_metrics; state the pre-read/pre-push hooks' purpose
(non-POSIX filesystem coherence).
The actor's update_weights is already @timer-wrapped (perf/update_weights_time),
so the per-sync total/publish/reload breakdown was duplicate instrumentation.
Keep only the delta-specific metrics (density, wire bytes).
The delta scaffold reworked the update-weight args: delta requires
--update-weight-transport=disk (was nccl-or-disk), needs
--update-weight-local-checkpoint-dir, and the --update-weight-delta-dir
compatibility alias is gone (the directory belongs to the transport, not the
encoding). Drop the alias resolve/backfill/conflict tests, point the transport
and colocate tests at the disk path, and cover the local-checkpoint requirement.
With the delta-dir alias gone, _resolve_update_weight_disk_dir no longer
normalizes anything — it's a single transport-level check, so fold it into
_validate_update_weight_args.
slime_validate_args validates everything else inline; the extracted
_validate_update_weight_args was the lone exception. Fold it in and test it
the same way as the other slime_validate_args checks (make_slime_validate_args).
Materialize the host-local checkpoint in a daemon thread at engine init so the
one-time base copy overlaps sglang launch and the first rollout (which serves
from init-loaded weights) instead of blocking the first delta reload. The first
sync_local_checkpoint's init_local_checkpoint is idempotent and flock-guarded,
so it either finds the copy done or blocks on the same lock — no join needed.
@nanjiangwill
nanjiangwill force-pushed the disk-delta-weight-sync branch from fd7c00d to a0b4b09 Compare June 17, 2026 18:17
@nanjiangwill
nanjiangwill requested a review from zhuzilin July 1, 2026 04:49
@zhuzilin
zhuzilin merged commit af3d7ff into THUDM:main Jul 2, 2026
40 checks passed
@nanjiangwill nanjiangwill changed the title Disk-level delta weight sync [2/n] Disaggregated rollout: disk-level delta weight sync Jul 6, 2026
nanjiangwill added a commit to modal-projects/slime that referenced this pull request Jul 6, 2026
…_weights)

Follow-up to THUDM#2089. That PR fanned the host-local delta apply out from slime via Ray (all_engine_actors, one actor per host of each multi-node engine) because only node 0 of a multi-node engine has an HTTP server. That leaks engine topology into slime and cannot work for external rollout engines, where slime only has an endpoint.

A new /pull_weights endpoint (shipped as the standalone docker/patch/latest/sglang-pull_weights.patch) has each engine pull the published weights onto every host it spans, riding the existing control-request broadcast; a per-host flock collapses co-located ranks to one pull, and the reply is all-gathered across the TP group so success means every host holds a verified checkpoint. The pull is artifact-driven: a published version is either a full HF checkpoint (copied as-is, resetting the chain) or a delta against its predecessor (patched in place, per-tensor checksums), detected from the index metadata; a fresh host seeds from the newest full version at or below the target, or from the engine's own model path for a pure-delta stream. The shared-filesystem refresh hook follows sglang's custom-weight-loader convention (--custom-pull-weights-pre-read-hook), forwarded through the existing --sglang-* passthrough.

slime now only talks to one endpoint per engine: all_engine_actors and SGLangEngine.sync_local_checkpoint are removed, the delta updater calls pull_weights(v) before the reload, and pull_weights(0) during baseline capture replaces the init-thread base materialization (still overlapped with the snapshot gather). Full-mode disk sync pulls to local disk too when --update-weight-local-checkpoint-dir is set. The post-write hook is renamed --custom-update-weight-post-write-path since it now serves full and delta alike, and two pre-existing full-disk bugs on non-POSIX shared filesystems are fixed: every writing rank creates the version dir and renames its own shards, and the post-write hook runs on every rank.

Validated end to end on GLM-4.7-Flash non-colocated (2 trainer nodes + one 2-node tp16 rollout engine, object-store-backed shared filesystem): delta mode 4/4 pulls with 0 checksum mismatches (~23s steady-state sync); full mode 4/4 pulls, one pull per host per version, engines reloading from the pulled local checkpoint.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants