[refactor] refactor weight update post-process and weight-checker - #1329
Merged
Merged
Conversation
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>
yueming-yuan
requested review from
Shi-Dong,
Zhichenzzz,
fzyzcjy,
guapisolo,
jybsuper,
maocheng23 and
yushengsu-thu
as code owners
June 12, 2026 05:35
yueming-yuan
changed the title
![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 refactors the weight post-processing mechanism by replacing the post_process_weights function and method with a pair of explicit session lifecycle methods: begin_weight_update and end_weight_update. This encapsulates the restoration of packed weights and the post-load/quantization post-processing steps within a structured session, allowing the engine to internally decide when post-loading is necessary. There are no review comments, and I have no feedback to provide.
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.
This was referenced Jun 12, 2026
11 tasks
# Conflicts: # miles/backends/megatron_utils/update_weight/common.py # miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py # tests/fast/backends/megatron_utils/test_lora_weight_sync_validation.py
guapisolo
added a commit
that referenced
this pull request
Jun 19, 2026
…ion for sglang #28575 Brings the post_process_weights -> begin/end_weight_update migration onto the CI branch so the weight-update suite runs against sglang #28575/#28710 with the miles-side adaptation in place. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Zhichenzzz
approved these changes
Jun 24, 2026
yueming-yuan
changed the title
fzyzcjy
approved these changes
Jun 24, 2026
mouad-hpc
pushed a commit
to Osmosis-AI/miles
that referenced
this pull request
Jul 9, 2026
Sync osmosis/dev with upstream radixark/miles main (254091f, 2026-07-06). 51 upstream commits; 212 files (+10066/-1205). 8 content conflicts resolved. Conflict resolutions (strategy: preserve osmosis behavior, port to upstream radixark#1329 API): - session/sessions.py: took upstream thin-adapter refactor (radixark#1510/radixark#1518/radixark#1563/radixark#1569); osmosis custom features (routing/indexer replay, R3 stripping) already live in core.py. - ray/rollout/train_data_conversion.py: kept BOTH osmosis adapter_slots + upstream opd_reverse_kl. - update_weight/common.py: dropped duplicate _check_weight_sync_results + orphaned post_process_weights tail (removed upstream by radixark#1329). - update_weight/update_weight_from_tensor.py: adopted begin/end_weight_update import; ported update_multi_lora_weights post_process_weights() -> end_weight_update(). - update_weight/update_weight_from_distributed/{broadcast,mixin,p2p}.py: kept osmosis multi-LoRA, never-sync-base, NCCL contiguity fix, and p2p LoRA support. - tests/.../test_lora_weight_sync_validation.py: took upstream signature (mock_common_ray for the new common.ray patch). Local static checks: py_compile + ruff (update_weight/ + session/) pass; all common imports resolve. Dynamic pytest tests/fast requires the miles container (torch/ray) — run in CI.
fzyzcjy
added a commit
that referenced
this pull request
Jul 10, 2026
Two mismatches on the real_rollout per-rollout weight-checksum path, both masked while the FT e2e tests were disabled in CI: 1. Nesting. RolloutManager.check_weights returns the single updatable server's result (server_group -> engine_body); #1329 established this (and test_check_weights_targets_only_updatable_model pins it). But flatten_inference_engine_checksums iterated one extra level, treating dict keys as engine bodies -> AttributeError: 'str' object has no attribute 'get'. Flatten the actual server_group -> engine_body shape. 2. Per-role parallelism_info. sglang's /weights_checker now returns each rank's parallelism_info as a per-role list ([{role: target, ...}, {role: draft, ...}]) for target + draft models, but the merge indexed it as a bare {rank: int} dict -> TypeError: list indices must be integers or slices, not str. Collapse the per-role list to the single shared GPU rank (asserting the roles agree). Unit tests updated to the real check_weights output and sglang schema, with regression coverage for the multi-role and rank-disagreement cases.
fzyzcjy
added a commit
that referenced
this pull request
Jul 10, 2026
Two mismatches on the real_rollout per-rollout weight-checksum path, both masked while the FT e2e tests were disabled in CI: 1. Nesting. RolloutManager.check_weights returns the single updatable server's result (server_group -> engine_body); #1329 established this (and test_check_weights_targets_only_updatable_model pins it). But flatten_inference_engine_checksums iterated one extra level, treating dict keys as engine bodies -> AttributeError: 'str' object has no attribute 'get'. Flatten the actual server_group -> engine_body shape. 2. Per-role parallelism_info. sglang's /weights_checker now returns each rank's parallelism_info as a per-role list ([{role: target, ...}, {role: draft, ...}]) for target + draft models, but the merge indexed it as a bare {rank: int} dict -> TypeError: list indices must be integers or slices, not str. Collapse the per-role list to the single shared GPU rank (asserting the roles agree). Update the checksum unit tests and the test_group mock to the real check_weights output and sglang schema, with regression coverage for the multi-role and rank-disagreement cases.
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
4 participants
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.
ci-sglang-pr: #28082