[RL] Add /pull_weights: engine-side pull of published weights into a host-local checkpoint - #30367
Open
nanjiangwill wants to merge 1 commit into
Open
[RL] Add /pull_weights: engine-side pull of published weights into a host-local checkpoint#30367nanjiangwill wants to merge 1 commit into
nanjiangwill wants to merge 1 commit into
Conversation
nanjiangwill
requested review from
CatherineSue,
Fridge003,
JustinTong0323,
Ying1123,
hnyls2002,
ispobock,
merrymercy,
slin1237 and
xiezhq-hermann
as code owners
July 7, 2026 08:17
github-actions
Bot
added
the
dependencies
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a new /pull_weights endpoint to enable host-local pulling of published weight deltas into local checkpoints, adding dependencies on xxhash and zstandard. The review feedback highlights critical compatibility issues, including a potential crash on non-POSIX systems (macOS/Windows) due to the use of os.posix_fadvise and another crash in non-distributed environments when torch.distributed is uninitialized. Additionally, several performance improvements were suggested to eliminate redundant memory copies when handling compressed data and byte buffers.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
nanjiangwill
force-pushed
the
pull-weights-main
branch
3 times, most recently
from
July 7, 2026 18:22
003bb57 to
9577941
Compare
nanjiangwill
added a commit
to modal-projects/sglang
that referenced
this pull request
Jul 13, 2026
…-local checkpoint
Add a disaggregated weight-sync receiver so an inference engine can bring its own
host-local checkpoint up to a published version while it keeps serving, decoupled
from the trainer.
- POST /pull_weights {local_checkpoint_dir, source_dir, target_version}: each host
walks the published weight_v{N:06d}/ chain back to the nearest full anchor and
replays the deltas forward (xor + zstd, xxh3-128 checksum) into a host-local
checkpoint that /update_weights_from_disk then reloads. New
weight_sync/local_checkpoint.py plus wiring through io_struct, scheduler,
weight_updater, tokenizer_control_mixin, http_server, and server_args
(custom_pull_weights_pre_read_hook lets the deployment refresh its volume view
before a read).
- Robust to an eventually-consistent volume mount whose view can change under a
running engine: read each delta blob whole-file into memory and size-verify it
against its own safetensors header before the xor, so the apply is immune to the
mount shifting mid-read; a short read (source not yet materialized) raises
rather than applying a partial delta, and the caller reloads + retries. One
reload per host (a per-host file lock collapses co-located ranks; only the lock
winner reloads). Reseed from the pristine boot checkpoint captured on the first
pull, never the mutated model_path.
Includes unit tests for torn/truncated/missing-blob sources.
Upstreaming: sgl-project#30367
nanjiangwill
added a commit
to modal-projects/sglang
that referenced
this pull request
Jul 13, 2026
…-local checkpoint
Add a disaggregated weight-sync receiver so an inference engine can bring its own
host-local checkpoint up to a published version while it keeps serving, decoupled
from the trainer.
- POST /pull_weights {local_checkpoint_dir, source_dir, target_version}: each host
walks the published weight_v{N:06d}/ chain back to the nearest full anchor and
replays the deltas forward (xor + zstd, xxh3-128 checksum) into a host-local
checkpoint that /update_weights_from_disk then reloads. New
weight_sync/local_checkpoint.py plus wiring through io_struct, scheduler,
weight_updater, tokenizer_control_mixin, http_server, and server_args
(custom_pull_weights_pre_read_hook lets the deployment refresh its volume view
before a read).
- Robust to an eventually-consistent volume mount whose view can change under a
running engine: read each delta blob whole-file into memory and size-verify it
against its own safetensors header before the xor, so the apply is immune to the
mount shifting mid-read; a short read (source not yet materialized) raises
rather than applying a partial delta, and the caller reloads + retries. One
reload per host (a per-host file lock collapses co-located ranks; only the lock
winner reloads). Reseed from the pristine boot checkpoint captured on the first
pull, never the mutated model_path.
Includes unit tests for torn/truncated/missing-blob sources.
Upstreaming: sgl-project#30367
nanjiangwill
added a commit
to modal-projects/sglang
that referenced
this pull request
Jul 20, 2026
…-local checkpoint
Add a disaggregated weight-sync receiver so an inference engine can bring its own
host-local checkpoint up to a published version while it keeps serving, decoupled
from the trainer.
- POST /pull_weights {local_checkpoint_dir, source_dir, target_version}: each host
walks the published weight_v{N:06d}/ chain back to the nearest full anchor and
replays the deltas forward (xor + zstd, xxh3-128 checksum) into a host-local
checkpoint that /update_weights_from_disk then reloads. New
weight_sync/local_checkpoint.py plus wiring through io_struct, scheduler,
weight_updater, tokenizer_control_mixin, http_server, and server_args
(custom_pull_weights_pre_read_hook lets the deployment refresh its volume view
before a read).
- Robust to an eventually-consistent volume mount whose view can change under a
running engine: read each delta blob whole-file into memory and size-verify it
against its own safetensors header before the xor, so the apply is immune to the
mount shifting mid-read; a short read (source not yet materialized) raises
rather than applying a partial delta, and the caller reloads + retries. One
reload per host (a per-host file lock collapses co-located ranks; only the lock
winner reloads). Reseed from the pristine boot checkpoint captured on the first
pull, never the mutated model_path.
Includes unit tests for torn/truncated/missing-blob sources.
Upstreaming: sgl-project#30367
nanjiangwill
added a commit
to modal-projects/sglang
that referenced
this pull request
Jul 20, 2026
…-local checkpoint
Add a disaggregated weight-sync receiver so an inference engine can bring its own
host-local checkpoint up to a published version while it keeps serving, decoupled
from the trainer.
- POST /pull_weights {local_checkpoint_dir, source_dir, target_version}: each host
walks the published weight_v{N:06d}/ chain back to the nearest full anchor and
replays the deltas forward (xor + zstd, xxh3-128 checksum) into a host-local
checkpoint that /update_weights_from_disk then reloads. New
weight_sync/local_checkpoint.py plus wiring through io_struct, scheduler,
weight_updater, tokenizer_control_mixin, http_server, and server_args
(custom_pull_weights_pre_read_hook lets the deployment refresh its volume view
before a read).
- Robust to an eventually-consistent volume mount whose view can change under a
running engine: read each delta blob whole-file into memory and size-verify it
against its own safetensors header before the xor, so the apply is immune to the
mount shifting mid-read; a short read (source not yet materialized) raises
rather than applying a partial delta, and the caller reloads + retries. One
reload per host (a per-host file lock collapses co-located ranks; only the lock
winner reloads). Reseed from the pristine boot checkpoint captured on the first
pull, never the mutated model_path.
Includes unit tests for torn/truncated/missing-blob sources.
Upstreaming: sgl-project#30367
nanjiangwill
force-pushed
the
pull-weights-main
branch
from
July 20, 2026 19:20
9577941 to
1c7f269
Compare
…host-local checkpoint
Add the disaggregated receiver used by Stitch: POST /pull_weights fans out across the deployment, walks the published weight_v{N} chain from the nearest full anchor, and replays zstd-compressed XOR deltas with per-tensor checksums into a host-local checkpoint for update_weights_from_disk.
Harden eventual-consistency and recovery semantics by reading and size-verifying each source blob in memory before mutation, distinguishing incomplete sources from corrupt local state, syncing dirty pages before advancing the version marker, deduplicating work under a per-host lock, and reseeding from the latched pristine boot checkpoint after a torn apply. Seed shards copy in parallel with 16 MiB streaming reads and progress logging.
Ported from the Stitch production fork. Upstreaming keeps the engine-side receiver independent of trainer framework and omits the separate partial-reload optimization tier.
nanjiangwill
force-pushed
the
pull-weights-main
branch
from
July 20, 2026 19:59
1c7f269 to
c1dd508
Compare
nanjiangwill
added a commit
to modal-projects/sglang
that referenced
this pull request
Jul 20, 2026
…-local checkpoint
Add a disaggregated weight-sync receiver so an inference engine can bring its own
host-local checkpoint up to a published version while it keeps serving, decoupled
from the trainer.
- POST /pull_weights {local_checkpoint_dir, source_dir, target_version}: each host
walks the published weight_v{N:06d}/ chain back to the nearest full anchor and
replays the deltas forward (xor + zstd, xxh3-128 checksum) into a host-local
checkpoint that /update_weights_from_disk then reloads. New
weight_sync/local_checkpoint.py plus wiring through io_struct, scheduler,
weight_updater, tokenizer_control_mixin, http_server, and server_args
(custom_pull_weights_pre_read_hook lets the deployment refresh its volume view
before a read).
- Robust to an eventually-consistent volume mount whose view can change under a
running engine: read each delta blob whole-file into memory and size-verify it
against its own safetensors header before the xor, so the apply is immune to the
mount shifting mid-read; a short read (source not yet materialized) raises
rather than applying a partial delta, and the caller reloads + retries. One
reload per host (a per-host file lock collapses co-located ranks; only the lock
winner reloads). Reseed from the pristine boot checkpoint captured on the first
pull, never the mutated model_path.
Includes unit tests for torn/truncated/missing-blob sources.
Upstreaming: sgl-project#30367
nanjiangwill
added a commit
to modal-projects/sglang
that referenced
this pull request
Jul 21, 2026
…-local checkpoint
Add a disaggregated weight-sync receiver so an inference engine can bring its own
host-local checkpoint up to a published version while it keeps serving, decoupled
from the trainer.
- POST /pull_weights {local_checkpoint_dir, source_dir, target_version}: each host
walks the published weight_v{N:06d}/ chain back to the nearest full anchor and
replays the deltas forward (xor + zstd, xxh3-128 checksum) into a host-local
checkpoint that /update_weights_from_disk then reloads. New
weight_sync/local_checkpoint.py plus wiring through io_struct, scheduler,
weight_updater, tokenizer_control_mixin, http_server, and server_args
(custom_pull_weights_pre_read_hook lets the deployment refresh its volume view
before a read).
- Catch-up is chosen per host, all in _pull_locked: a steady-state host applies the
one new delta in place (mmap, only the delta's dirty pages, O(delta)); a host
several versions behind folds the whole range in one pass per tensor — one buffered
read + XOR the range in RAM + one write — instead of N mmap read-modify-writes, each
faulting the whole tensor in; and a fresh mid-run joiner, when an aggregated delta
aggregate/ (base xor vM as one delta, base_version 0) has been published, applies it
once on the base seed to reach vM and then folds only the d(M+1)..target tail. The
aggregate is a separate writer's single in-place dir on an eventually-consistent
source, so a reader can observe a partial (mid-update) state; every shard stamps its
version and a reader that sees a disagreement (a torn read whose per-shard checksums
still self-verify) skips it and folds from base.
- Robust to an eventually-consistent volume mount whose view can change under a
running engine: read each delta blob whole-file into memory and size-verify it
against its own safetensors header before the xor, so the apply is immune to the
mount shifting mid-read; a short read (source not yet materialized) raises
rather than applying a partial delta, and the caller reloads + retries. One
reload per host (a per-host file lock collapses co-located ranks; only the lock
winner reloads). Reseed from the pristine boot checkpoint captured on the first
pull, never the mutated model_path.
Includes unit tests for torn/truncated/missing-blob sources, the multi-delta fold,
and the aggregate fast-forward.
Upstreaming: sgl-project#30367
nanjiangwill
added a commit
to modal-projects/sglang
that referenced
this pull request
Jul 21, 2026
…-local checkpoint
Add a disaggregated weight-sync receiver so an inference engine can bring its own
host-local checkpoint up to a published version while it keeps serving, decoupled
from the trainer.
- POST /pull_weights {local_checkpoint_dir, source_dir, target_version}: each host
walks the published weight_v{N:06d}/ chain back to the nearest full anchor and
replays the deltas forward (xor + zstd, xxh3-128 checksum) into a host-local
checkpoint that /update_weights_from_disk then reloads. New
weight_sync/local_checkpoint.py plus wiring through io_struct, scheduler,
weight_updater, tokenizer_control_mixin, http_server, and server_args
(custom_pull_weights_pre_read_hook lets the deployment refresh its volume view
before a read).
- Catch-up is chosen per host, all in _pull_locked: a steady-state host applies the
one new delta in place (mmap, only the delta's dirty pages, O(delta)); a host
several versions behind folds the whole range in one pass per tensor — one buffered
read + XOR the range in RAM + one write — instead of N mmap read-modify-writes, each
faulting the whole tensor in. Pure xor/zstd ranges only; a non-xor version in the
range falls back to the per-version apply.
- Robust to an eventually-consistent volume mount whose view can change under a
running engine: read each delta blob whole-file into memory and size-verify it
against its own safetensors header before the xor, so the apply is immune to the
mount shifting mid-read; a short read (source not yet materialized) raises
rather than applying a partial delta, and the caller reloads + retries. One
reload per host (a per-host file lock collapses co-located ranks; only the lock
winner reloads). Reseed from the pristine boot checkpoint captured on the first
pull, never the mutated model_path.
Includes unit tests for torn/truncated/missing-blob sources and the multi-delta fold.
Upstreaming: sgl-project#30367
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
RL trainers with disaggregated rollout publish each weight sync as a version dir
weight_v{N:06d}/on a shared filesystem — a full HF checkpoint, or zstd-compressed per-tensor byte deltas (xor/overwrite) with per-tensor checksums, packaged as a canonical HF checkpoint dir. This adds the engine-side receiver:POST /pull_weightsbrings a host-local checkpoint up to a target version on every host the deployment spans, so the trainer talks to one endpoint per engine.Trainer PR: radixark/miles#1235 ·
sglang-milesvariant: #30366 · Refs: THUDM/slime#2181Behavior
POST /pull_weights {local_checkpoint_dir, source_dir, target_version}fans out to every scheduler rank on every node. Each host seeds from the newest full version ≤ target (or the server's ownmodel_path— version 0 is the engine's base), then applies the delta chain in place via mmap, parallelized across tensors. A checksum mismatch or out-of-order apply raises — never serve bad weights; a per-host flock + applied-version marker collapse co-located ranks to one pull. Success is gathered across the TP group, so the reply covers every host. The trainer then reloads via the ordinary/update_weights_from_disk— weight loading never sees the delta format.--custom-pull-weights-pre-read-hook <import.path>: refresh hook for object-store-backed mounts without cross-host read-after-write consistency (POSIX shared FS needs none).Notes
xxhashadded to pyproject (zstandardalready present on main);blake3stays an optional lazy import. The module is imported only inside the handler.PullWeightsReqInput/Outputfollow main's msgspecBaseReq; the server arg uses the annotation-driven declaration.CI States
Latest PR Test (Base): ❌ Run #29774219418
Latest PR Test (Extra): ❌ Run #29774219280