Skip to content

Revert #27749 and re-impl mtp update weight from distributed - #28575

Merged
yueming-yuan merged 2 commits into
sgl-project:sglang-milesfrom
guapisolo:fix/draft_weight_upd_pr
Jun 19, 2026
Merged

Revert #27749 and re-impl mtp update weight from distributed#28575
yueming-yuan merged 2 commits into
sgl-project:sglang-milesfrom
guapisolo:fix/draft_weight_upd_pr

Conversation

@guapisolo

@guapisolo guapisolo commented Jun 17, 2026 •
edited by github-actions Bot
Loading

Copy link
Copy Markdown
Collaborator

Summary

Centralize speculative weight-update fan-out in the scheduler behind a selector field.

Motivation

sglang-miles carries #27749 (a0ff6904) as the first cut of distributed weight-update for speculative draft worker(s). That version fans out through a disable_draft_model flag, per-worker update_weights_from_distributed / update_weights_from_tensor methods, and ModelRunner.update_weights_from_distributed_to_model_runners([...]). Each spec worker reimplements its own runner-list assembly, and disable_draft_model only expresses all-or-nothing. The observable capability — pushing updated weights to the target and draft runners over the existing _model_update_group — is unchanged.

Before / After

  • Before: each spec worker builds its own runner list inside update_weights_from_distributed.
  • After: SchedulerWeightUpdaterManager fans out once for every worker.
  • The target receives the broadcast a single time via receive_weights_from_distributed.
  • The received weights load into each runner from get_model_runners(selector).
  • Per-worker update_weights_from_{distributed,tensor} are removed.
  • Each worker now exposes only iter_draft_runners().
  • ModelRunner.update_weights_from_distributed splits into receive_weights_from_distributed plus load_weights.
  • disable_draft_model: Optional[bool] becomes selector: {target, draft, all} (default all).

Behavior Preservation

  • How we know: test/srt/test_distributed_weight_update_spec_worker.py asserts receive-once-on-target / load-into-each-runner.
  • The same test asserts selector="target" never enumerates the draft worker.

Verification

  • Existing test suite: test/srt/test_distributed_weight_update_spec_worker.py — 2 passed locally.
  • All 10 touched modules import cleanly.

Review Focus

  • Scrutinize SchedulerWeightUpdaterManager.update_weights_from_distributed: one weights object loads into every runner from get_model_runners(selector).
  • Scrutinize dflash_worker.iter_draft_runners() returns the real draft, not the target via __getattr__.
  • Scrutinize ngram_worker.iter_draft_runners() returns [] since it shares the target runner.

CI States

Latest PR Test (Base): ❌ Run #27793336996
Latest PR Test (Extra): ❌ Run #27793336669

@gemini-code-assist

gemini-code-assist Bot commented Jun 17, 2026

Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@guapisolo
guapisolo force-pushed the fix/draft_weight_upd_pr branch from 9c3eaa5 to e4b245b Compare June 18, 2026 22:13
@guapisolo
guapisolo changed the base branch from sglang-miles-v0.5.13 to sglang-miles June 18, 2026 22:13
@guapisolo @claude
feat: re-impl mtp update weight from distributed
b29c5e4 
Re-implement distributed weight-update for speculative draft worker(s) on the
v2 worker layout, on top of the preceding revert of sgl-project#27749's earlier design.

Centralize the weight-update fan-out in SchedulerWeightUpdaterManager: an
UpdateWeightsFrom{Distributed,Tensor}ReqInput `selector` ({target, draft,
both}) resolves to the (role, ModelRunner) pairs to update. For distributed
updates the target model receives the broadcast once and the weights are
loaded into every selected runner locally; tensor updates deserialize once and
fan out the same way. Draft workers expose their runner(s) via
iter_draft_runners(); the per-worker update_weights_from_{distributed,tensor}
methods and the `disable_draft_model` flag are gone.

Combines sglang-miles sgl-project#27749 and sgl-project#28548.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@guapisolo
guapisolo force-pushed the fix/draft_weight_upd_pr branch from e4b245b to b29c5e4 Compare June 18, 2026 22:31
@yueming-yuan yueming-yuan mentioned this pull request Jun 19, 2026
Closed
@yueming-yuan
yueming-yuan merged commit 22b6497 into sgl-project:sglang-miles Jun 19, 2026
37 of 42 checks passed
nanjiangwill pushed a commit to nanjiangwill/sglang that referenced this pull request Jul 7, 2026
@guapisolo @yueming-yuan
[sglang-miles] Update weights from distributed for speculative draft …
bc2dd18 
…worker(s) (sgl-project#27749, sgl-project#28575)

Squash of sgl-project#27749 and its revert/re-impl sgl-project#28575: draft workers expose
iter_draft_runners() and the target runner receives the broadcast once
via receive_weights_from_distributed, then loads into each draft runner
locally. v1 dflash/frozen-kv-mtp workers were removed in v0.5.14; their
hunks are dropped (spec v2 is the only path).
yueming-yuan pushed a commit that referenced this pull request Jul 14, 2026
@guapisolo @yueming-yuan
[17/24] [sglang-miles] Update weights from distributed for speculativ…
e523ebc 
…e draft worker(s) (#27749, #28575)

Squash of #27749 and its revert/re-impl #28575: draft workers expose
iter_draft_runners() and the target runner receives the broadcast once
via receive_weights_from_distributed, then loads into each draft runner
locally. v1 dflash/frozen-kv-mtp workers were removed in v0.5.14; their
hunks are dropped (spec v2 is the only path).

Reconciled distributed-update return onto v0.5.15's kw_only ReqOutput
(error logged once in the except path).
yueming-yuan pushed a commit that referenced this pull request Jul 14, 2026
@guapisolo @yueming-yuan
[17/24] [sglang-miles] Update weights from distributed for speculativ…
414c2b8 
…e draft worker(s) (#27749, #28575)

Squash of #27749 and its revert/re-impl #28575: draft workers expose
iter_draft_runners() and the target runner receives the broadcast once
via receive_weights_from_distributed, then loads into each draft runner
locally. v1 dflash/frozen-kv-mtp workers were removed in v0.5.14; their
hunks are dropped (spec v2 is the only path).

Reconciled distributed-update return onto v0.5.15's kw_only ReqOutput
(error logged once in the except path).
yueming-yuan pushed a commit that referenced this pull request Jul 14, 2026
@guapisolo @yueming-yuan
[17/24] [sglang-miles] Update weights from distributed for speculativ…
2e696f3 
…e draft worker(s) (#27749, #28575)

Squash of #27749 and its revert/re-impl #28575: draft workers expose
iter_draft_runners() and the target runner receives the broadcast once
via receive_weights_from_distributed, then loads into each draft runner
locally. v1 dflash/frozen-kv-mtp workers were removed in v0.5.14; their
hunks are dropped (spec v2 is the only path).

Reconciled distributed-update return onto v0.5.15's kw_only ReqOutput
(error logged once in the except path).
yueming-yuan added a commit that referenced this pull request Jul 25, 2026
@guapisolo @yueming-yuan
@JD-ETH @maocheng23
[16/23] [sglang-miles] RL weight-update sessions + distributed update…
ca22702 
… for spec draft worker(s) (#27749, #28575, #18565, #22663, #28001, #29675, #27750)

Squash of the spec-draft distributed-update work and the weight-processing
session work: the latter rewrites iter_draft_runners()/get_model_runners() the
former introduces, so they cannot be applied independently.

  - Draft runners never join the update group, so the target receives the
    broadcast once and the weights are loaded into every selected runner.
  - Begin/EndWeightUpdate session chain: engine / http_server /
    tokenizer_control_mixin (pause-aware locking) -> scheduler ->
    SchedulerWeightUpdaterManager, with {target,draft,all} runner selectors.
  - loader.py: public post_load_weights plus restore_weight/postprocess_weight
    via _apply_quant_method_hook (skips LoRA wrappers).
  - weight checker: skip_tensor_list + role-prefixed overall_checksum over
    get_model_runners(selector), replacing _get_draft_model_runner.

Reimplemented onto v0.5.16, which had already extracted weight updating into
ModelRunner.weight_updater (WeightUpdater): the receive/load split lands there
as receive_weights_from_distributed() + load_weights() rather than on
ModelRunner, and the scheduler fan-out drives runner.weight_updater.*. The
worker-level update_weights_from_{distributed,tensor} entry points on
TpModelWorker / EAGLEWorkerV2 / NGRAMWorker are dropped: the scheduler now owns
the fan-out, and leaving them would be a second path that updates one runner
only. iter_runners() replaces the ad-hoc draft-runner discovery on every spec v2
worker (incl. DFlashWorkerV2). The unit test moves to test/registered/rl/, since
v0.5.16 retired test/srt/ from CI collection.

Co-authored-by: Yueming Yuan <yym022502@gmail.com>
Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
yueming-yuan added a commit that referenced this pull request Jul 25, 2026
@guapisolo @yueming-yuan
@JD-ETH @maocheng23
[16/24] [sglang-miles] RL weight-update sessions + distributed update…
1f9d84a 
… for spec draft worker(s) (#27749, #28575, #18565, #22663, #28001, #29675, #27750)

Squash of the spec-draft distributed-update work and the weight-processing
session work: the latter rewrites iter_draft_runners()/get_model_runners() the
former introduces, so they cannot be applied independently.

  - Draft runners never join the update group, so the target receives the
    broadcast once and the weights are loaded into every selected runner.
  - Begin/EndWeightUpdate session chain: engine / http_server /
    tokenizer_control_mixin (pause-aware locking) -> scheduler ->
    SchedulerWeightUpdaterManager, with {target,draft,all} runner selectors.
  - loader.py: public post_load_weights plus restore_weight/postprocess_weight
    via _apply_quant_method_hook (skips LoRA wrappers).
  - weight checker: skip_tensor_list + role-prefixed overall_checksum over
    get_model_runners(selector), replacing _get_draft_model_runner.

Reimplemented onto v0.5.16, which had already extracted weight updating into
ModelRunner.weight_updater (WeightUpdater): the receive/load split lands there
as receive_weights_from_distributed() + load_weights() rather than on
ModelRunner, and the scheduler fan-out drives runner.weight_updater.*. The
worker-level update_weights_from_{distributed,tensor} entry points on
TpModelWorker / EAGLEWorkerV2 / NGRAMWorker are dropped: the scheduler now owns
the fan-out, and leaving them would be a second path that updates one runner
only. iter_runners() replaces the ad-hoc draft-runner discovery on every spec v2
worker (incl. DFlashWorkerV2). The unit test moves to test/registered/rl/, since
v0.5.16 retired test/srt/ from CI collection.

Co-authored-by: Yueming Yuan <yym022502@gmail.com>
Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
yueming-yuan added a commit that referenced this pull request Jul 25, 2026
@guapisolo @yueming-yuan
@JD-ETH @maocheng23
[16/25] [sglang-miles] RL weight-update sessions + distributed update…
57ef29a 
… for spec draft worker(s) (#27749, #28575, #18565, #22663, #28001, #29675, #27750)

Squash of the spec-draft distributed-update work and the weight-processing
session work: the latter rewrites iter_draft_runners()/get_model_runners() the
former introduces, so they cannot be applied independently.

  - Draft runners never join the update group, so the target receives the
    broadcast once and the weights are loaded into every selected runner.
  - Begin/EndWeightUpdate session chain: engine / http_server /
    tokenizer_control_mixin (pause-aware locking) -> scheduler ->
    SchedulerWeightUpdaterManager, with {target,draft,all} runner selectors.
  - loader.py: public post_load_weights plus restore_weight/postprocess_weight
    via _apply_quant_method_hook (skips LoRA wrappers).
  - weight checker: skip_tensor_list + role-prefixed overall_checksum over
    get_model_runners(selector), replacing _get_draft_model_runner.

Reimplemented onto v0.5.16, which had already extracted weight updating into
ModelRunner.weight_updater (WeightUpdater): the receive/load split lands there
as receive_weights_from_distributed() + load_weights() rather than on
ModelRunner, and the scheduler fan-out drives runner.weight_updater.*. The
worker-level update_weights_from_{distributed,tensor} entry points on
TpModelWorker / EAGLEWorkerV2 / NGRAMWorker are dropped: the scheduler now owns
the fan-out, and leaving them would be a second path that updates one runner
only. iter_runners() replaces the ad-hoc draft-runner discovery on every spec v2
worker (incl. DFlashWorkerV2). The unit test moves to test/registered/rl/, since
v0.5.16 retired test/srt/ from CI collection.

Co-authored-by: Yueming Yuan <yym022502@gmail.com>
Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
yueming-yuan added a commit that referenced this pull request Jul 25, 2026
@guapisolo @yueming-yuan
@JD-ETH @maocheng23
[16/27] [sglang-miles] RL weight-update sessions + distributed update…
bb17e99 
… for spec draft worker(s) (#27749, #28575, #18565, #22663, #28001, #29675, #27750)

Squash of the spec-draft distributed-update work and the weight-processing
session work: the latter rewrites iter_draft_runners()/get_model_runners() the
former introduces, so they cannot be applied independently.

  - Draft runners never join the update group, so the target receives the
    broadcast once and the weights are loaded into every selected runner.
  - Begin/EndWeightUpdate session chain: engine / http_server /
    tokenizer_control_mixin (pause-aware locking) -> scheduler ->
    SchedulerWeightUpdaterManager, with {target,draft,all} runner selectors.
  - loader.py: public post_load_weights plus restore_weight/postprocess_weight
    via _apply_quant_method_hook (skips LoRA wrappers).
  - weight checker: skip_tensor_list + role-prefixed overall_checksum over
    get_model_runners(selector), replacing _get_draft_model_runner.

Reimplemented onto v0.5.16, which had already extracted weight updating into
ModelRunner.weight_updater (WeightUpdater): the receive/load split lands there
as receive_weights_from_distributed() + load_weights() rather than on
ModelRunner, and the scheduler fan-out drives runner.weight_updater.*. The
worker-level update_weights_from_{distributed,tensor} entry points on
TpModelWorker / EAGLEWorkerV2 / NGRAMWorker are dropped: the scheduler now owns
the fan-out, and leaving them would be a second path that updates one runner
only. iter_runners() replaces the ad-hoc draft-runner discovery on every spec v2
worker (incl. DFlashWorkerV2). The unit test moves to test/registered/rl/, since
v0.5.16 retired test/srt/ from CI collection.

Co-authored-by: Yueming Yuan <yym022502@gmail.com>
Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@yueming-yuan yueming-yuan yueming-yuan approved these changes
@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 is a code owner
@Qiaolin-Yu Qiaolin-Yu Awaiting requested review from Qiaolin-Yu Qiaolin-Yu is a code owner
@kpham-sgl kpham-sgl Awaiting requested review from kpham-sgl kpham-sgl is a code owner
@pyc96 pyc96 Awaiting requested review from pyc96 pyc96 is a code owner
@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner
@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner
@Fridge003 Fridge003 Awaiting requested review from Fridge003 Fridge003 is a code owner
@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner
@xiezhq-hermann xiezhq-hermann Awaiting requested review from xiezhq-hermann xiezhq-hermann is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@guapisolo @yueming-yuan