refactor weight processing in RL weight update by yueming-yuan · Pull Request #28001 · sgl-project/sglang

yueming-yuan · 2026-06-12T05:32:06Z

Summary

Refactors the online weight-update post-processing on the SGLang side. Training (miles) no longer drives the restore/quantize ordering via a 3-flag post_process_weights RPC. Instead the engine owns a weight-update session:

begin_weight_update — restores in-place-packed weights to a loadable state (no-op for schemes that don't repack, e.g. plain fp8). Sets _weight_update_in_progress.
update_weights_from_* — assert a session is open (catches silent misuse); record that model.load_weights() ran this session.
end_weight_update — quant finalize (process_weights_after_loading) on the full model, plus model.post_load_weights() only when load_weights was bypassed this session (i.e. P2P/RDMA). The engine decides this itself via _weight_update_loaded; no flag crosses the API.

post_process_weights (RPC + the restore/quantize/post_load boolean dispatcher) is deleted.

Why

The restore/quantize/post-load transforms depend on engine-side settings (kernel backend, hardware, quant scheme) — they belong on the engine, not the trainer.
A begin/end bracket gives "restore once / finalize once" naturally and works uniformly for tensor / broadcast / P2P (P2P has no per-bucket update RPC, so explicit begin/end is required).
The open-session assert turns a previously-silent misuse (updating weights without the surrounding prep) into a loud error.

Changes

io_struct.py: BeginWeightUpdate* / EndWeightUpdate* req/resp replace PostProcessWeights*.
model_loader/loader.py: extracted restore_weights_before_loading_all / process_weights_after_loading_all as free functions (single canonical definition, alongside the existing post-load helper); renamed _post_load_weights → post_load_weights.
scheduler / tp_worker / tokenizer_control_mixin / http_server / engine: begin/end wiring; removed post_process_weights.

Paired change

Requires the matching miles-side PR (begin/end calls replace post_process_weights).

Notes

post_load_weights is still embedded in model.load_weights() for tensor/broadcast (runs per-bucket); pulling it out so all transports run it once at end_weight_update is a deliberate follow-up. _weight_update_loaded is the seam for it.
Lint: isort + black applied to the changed files manually (full pre-commit suite not runnable in this partial checkout).

CI States

Latest PR Test (Base): ❌ Run #28138539341
Latest PR Test (Extra): 🚫 Run #28138539302

gemini-code-assist · 2026-06-12T05:32:09Z

Warning

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

Replaces the post_process_weights RPC calls with SGLang's new begin_weight_update / end_weight_update session API. begin opens the session (restore packed weights) right after pause/flush; end closes it (quant finalize + post_load) before resuming generation. Drops the int4-only restore conditional and the per-transport post_load_weights flag -- the engine now decides post_load itself based on whether model.load_weights ran this session (P2P bypasses it). Requires the paired sglang change (sgl-project/sglang#28001). Co-authored-by: Nan Jiang <59716405+nanjiangwill@users.noreply.github.com>

nanjiangwill · 2026-06-12T17:53:26Z

+        restore_weight(self.model_runner.model, torch.device(self.device))
+        return True, "Success"
+
+    def end_weight_update(self, run_post_load: bool):


do we need run_post_load: bool? cuz i notice miles side just do end_weight_update(self.rollout_engines). also for correctness i think we need to do post_load_weights all the time?

post_load_weight is coupled with model.load_weights(), so any update weight method triggering load_weight() has already do this.
This run_post_load method is used to support the p2p weight update, where load_weight() is not triggered, and the tensors are directly replaced by p2p.

guapisolo · 2026-06-12T20:05:34Z

+    def begin_weight_update(self, recv_req: BeginWeightUpdateReqInput):
+        """Begin a new weight update session: restore packed weights to a loadable state."""
+        self._quiesce_for_weight_update()
+        success, message = self.tp_worker.begin_weight_update()


Whether draft worker need begin weight update?

let me push some commit in this PR to fix...

Training no longer drives restore/quantize ordering. The engine opens a weight-update session via begin_weight_update (restores in-place-packed weights to a loadable state) and closes it via end_weight_update (quant finalize on the full model, plus model.post_load_weights only when load_weights was bypassed this session, e.g. P2P/RDMA). The engine tracks the session itself; update_weights_from_* now assert a session is open to avoid silent misuse. - io_struct: BeginWeightUpdate/EndWeightUpdate req/resp (replaces PostProcessWeights) - loader: extract restore_weight / postprocess_weight free functions; rename _post_load_weights -> post_load_weights - scheduler/tp_worker/tokenizer_control/http_server/engine: begin/end wiring; remove post_process_weights and the 3-flag dispatcher On v0.5.13 the scheduler weight-update logic lives in SchedulerWeightUpdaterManager (scheduler_components/weight_updater.py), so the session state + begin/end methods are hosted there. Co-authored-by: Nan Jiang <59716405+nanjiangwill@users.noreply.github.com>

…_refactor # Conflicts: # python/sglang/srt/managers/scheduler_components/weight_updater.py

…28710) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

…low_quant_error (#28082)

…sgl-project#18565, sgl-project#22663, sgl-project#28001, sgl-project#29675) Combined weight-processing machinery, folding the general post-process plumbing from the INT4 QAT work (sgl-project#18565), the P2P post_load_weights extension (sgl-project#22663), the weight-processing refactor (sgl-project#28001), and the pause-aware weight locking fix (sgl-project#29675) into its final form: - Begin/EndWeightUpdate request chain: engine/http_server/ tokenizer_control_mixin (pause-aware locking) -> scheduler -> weight_updater sessions with {target,draft,all} runner selectors - loader.py: public post_load_weights, restore_weight/postprocess_weight via _apply_quant_method_hook (skips LoRA wrappers) - model_runner begin/end_weight_update + receive_weights_from_distributed - iter_runners() on tp_worker and all spec v2 workers (incl. port to DFlashWorkerV2; v1 dflash/frozen-kv workers no longer exist) - CompressedTensorsFusedMoEMethod.restore_weights_before_loading dispatch - weight_checker ComparableWeight refactor + weight_checker_comparator Co-authored-by: JD-ETH <jaedon.guo@gmail.com> Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>

…sessions + spec-draft weight check (#18565, #22663, #28001, #29675, #27750) Combined weight-processing machinery, folding the general post-process plumbing from the INT4 QAT work (#18565), the P2P post_load_weights extension (#22663), the weight-processing refactor (#28001), and the pause-aware weight locking fix (#29675) into its final form, plus the extension of the weight checker to speculative draft worker(s) (#27750): - Begin/EndWeightUpdate request chain: engine/http_server/ tokenizer_control_mixin (pause-aware locking) -> scheduler -> weight_updater sessions with {target,draft,all} runner selectors - loader.py: public post_load_weights, restore_weight/postprocess_weight via _apply_quant_method_hook (skips LoRA wrappers) - model_runner begin/end_weight_update + receive_weights_from_distributed - iter_runners() on tp_worker and all spec v2 workers (incl. DFlashWorkerV2) - CompressedTensorsFusedMoEMethod.restore_weights_before_loading dispatch - weight checker: selector + skip_tensor_list + role-prefixed overall_checksum over get_model_runners(selector), replacing _get_draft_model_runner Kept v0.5.15's upstreamed ComparableWeight weight_checker_comparator (incl. the #29623 ue8m0 scale-unpack fix); adapted the weight_checker caller to its public compare_weights/CHUNK_NUMEL API. Co-authored-by: JD-ETH <jaedon.guo@gmail.com> Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com> Co-authored-by: Jiajun Li <jiajun.li@radixark.ai>

… 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>

… for spec draft worker(s) (sgl-project#27749, sgl-project#28575, sgl-project#18565, sgl-project#22663, sgl-project#28001, sgl-project#29675, sgl-project#27750)

… for spec draft worker(s) (sgl-project#27749, sgl-project#28575, sgl-project#18565, sgl-project#22663, sgl-project#28001, sgl-project#29675, sgl-project#27750) (cherry picked from commit 4290cf4)

… for spec draft worker(s) (#27749, #28575, #18565, #22663, #28001, #29675, #27750) (cherry picked from commit 4290cf4) (cherry picked from commit 3c0c432)

… for spec draft worker(s) (sgl-project#27749, sgl-project#28575, sgl-project#18565, sgl-project#22663, sgl-project#28001, sgl-project#29675, sgl-project#27750) (cherry picked from commit 4290cf4)

yueming-yuan requested review from CatherineSue, Fridge003, JustinTong0323, Ying1123, hnyls2002, ispobock, merrymercy, slin1237 and xiezhq-hermann as code owners June 12, 2026 05:32

yueming-yuan force-pushed the weight_update_refactor branch from f0f0a14 to c556cb1 Compare June 12, 2026 05:33

yueming-yuan mentioned this pull request Jun 12, 2026

[refactor] refactor weight update post-process and weight-checker radixark/miles#1329

Merged

yueming-yuan force-pushed the weight_update_refactor branch from c556cb1 to f9e4415 Compare June 12, 2026 05:46

yueming-yuan changed the title ~~Replace post_process_weights with begin/end_weight_update session API~~ refactor weight processing in RL weight update Jun 12, 2026

nanjiangwill reviewed Jun 12, 2026

View reviewed changes

guapisolo reviewed Jun 12, 2026

View reviewed changes

This was referenced Jun 12, 2026

[sglang-miles] weight checker: ULP-based quant error tolerance via allow_quant_error #28082

Merged

[refactor] weight update refactor radixark/miles#1360

Open

yueming-yuan force-pushed the sglang-miles branch from 1dbbf3a to 6a4ffcc Compare June 18, 2026 21:02

yueming-yuan requested review from Qiaolin-Yu, ch-wan, fzyzcjy, hanming-lu, key4ng, kpham-sgl, mickqian, pyc96, yhyang201 and yuan-luo as code owners June 18, 2026 21:02

yueming-yuan requested review from ByronHsu, Kangyan-Zhou, b8zhong, bingxche, hebiao064 and sogalin as code owners June 18, 2026 21:02

yueming-yuan force-pushed the weight_update_refactor branch from fe5b459 to 33f9d0d Compare June 18, 2026 23:44

yueming-yuan and others added 3 commits June 18, 2026 19:33

Merge remote-tracking branch 'origin/sglang-miles' into weight_update…

055cb1f

…_refactor # Conflicts: # python/sglang/srt/managers/scheduler_components/weight_updater.py

refactor: unify RL weight-update fanout across target and draft/MTP (#…

d5243e8

…28710) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

[sglang-miles] weight checker: ULP-based quant error tolerance via al…

f3914d0

…low_quant_error (#28082)

yueming-yuan merged commit 64b0059 into sglang-miles Jun 25, 2026
5 of 6 checks passed

yueming-yuan deleted the weight_update_refactor branch June 25, 2026 00:26

xyuzh mentioned this pull request Jun 29, 2026

[sglang-miles] RDT/NIXL weight sync support for Ray scheduler actors #27723

Merged

yushengsu-thu mentioned this pull request Jul 17, 2026

[LoRA] Skip LoRA wrapper modules in post-load quant weight processing #31524

Open

5 tasks

guapisolo mentioned this pull request Sep 14, 2026

[Tracking] Upstream remaining sglang-miles changes to main #39467

Open

46 tasks

Navigation Menu

refactor weight processing in RL weight update - #28001

refactor weight processing in RL weight update #28001
yueming-yuan merged 4 commits into
sglang-milesfrom
weight_update_refactor

yueming-yuan commented Jun 12, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot commented Jun 12, 2026

Uh oh!

nanjiangwill Jun 12, 2026

Uh oh!

yueming-yuan Jun 12, 2026

Uh oh!

guapisolo Jun 12, 2026

Uh oh!

guapisolo Jun 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

yueming-yuan commented Jun 12, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Why

Changes

Paired change

Notes

CI States

Uh oh!

gemini-code-assist Bot commented Jun 12, 2026

Uh oh!

nanjiangwill Jun 12, 2026

Choose a reason for hiding this comment

Uh oh!

yueming-yuan Jun 12, 2026

Choose a reason for hiding this comment

Uh oh!

guapisolo Jun 12, 2026

Choose a reason for hiding this comment

Uh oh!

guapisolo Jun 16, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yueming-yuan commented Jun 12, 2026 •

edited by github-actions Bot

Loading