fix(weight-update): sync and check drafts only when the trainer trains MTP by guapisolo · Pull Request #4043 · radixark/miles

guapisolo · 2026-10-13T04:42:16Z

Depends on: #4041
Stacked on the parent PR's branch. Review / merge the parent first.

Summary

Weight syncs and the weight checker covered speculative drafts the trainer never updates

Symptom & Reproduction

Symptom: under speculative decoding, FSDP/torchtitan open every update session with "all", so sglang restores, then postprocesses, a draft that gets no weights; p2p connect raises NotImplementedError for EAGLE3, STANDALONE, separate draft checkpoints or multi-layer EAGLE (p2p: update speculative drafts from trained MTP weights #3985 (comment)).
Symptom (checker): with --check-weight-update-equal the server cell scrambles every runner at engine init, so a draft the sync leaves out stays random for the run.
Reproduction: test_the_checker_snapshots_and_resets_only_the_runners_it_compares fails on the parent commit, where snapshot plus reset_tensors get "all" while the checker compares "target"; it passes here.

Root Cause

MegatronHfWeightIteratorBase.__init__ alone narrowed weight_update_selector to "target".
FSDPHfWeightIterator and TitanHfWeightIterator kept the base default "all".
ServerCell._tick_when_initializing hard-coded selector="all" for snapshot and reset_tensors.
--check-weight-update-selector defaulted to "all", independent of the sync.

Fix

One rule, weight_update_selector(args) in weight_update/session.py: "target" when a speculative algorithm is set without --enable-mtp-training, else "all". The parent PR gates Megatron's MTP layers on that intent and rejects it on FSDP and torchtitan, so the rule holds on every backend without a per-backend hook.

The updater evaluates it on each trainer's args for p2p connect and begin_weight_update.
--check-weight-update-selector defaults to it at argument validation; the server cell snapshots and resets that selector, so reset covers exactly what compare checks.
Removed: the iterator attribute, Megatron's model-based copy of the rule, the manual selectors in the GLM-4.7-Flash/Qwen3.5 e2e harnesses.

Megatron syncs are unchanged, since the parent PR builds MTP layers exactly when --enable-mtp-training is set, including trainer overrides and Bridge LoRA. An explicit --check-weight-update-selector overrides the default.

Verification

tests/fast/backends, tests/fast/utils/test_arguments.py, tests/fast/ray/rollout, tests/fast/launch_scripts, tests/fast/examples, tests/fast/e2e, tests/fast/test_train*.py: 4440 passed on the radixark/miles:dev image (CPU); the 10 failures (a secret-argument check against the image's newer sglang router flags, a missing tap module, 8 helm template cases missing the miles-common chart dependency) fail identically on main.
test_weight_update_selector.py (moved from megatron_utils/, same stage-b-cpu registration) covers trained, frozen and external drafts.
The test_megatron_config.py trainer-override tests and test_an_explicit_zero_becomes_megatrons_no_mtp_value cover the parent PR's rule that keeps this selector true per trainer.
GPU e2e was not run.

Review Focus

server_cell.py: snapshot and reset_tensors now follow --check-weight-update-selector; a runner outside it is neither scrambled nor checked.
Multi-policy: the checker default comes from the base args, while each trainer's sync uses its own args.

Generated with Claude Code

…s MTP Only the Megatron HF weight iterator narrowed weight_update_selector to "target"; FSDP and torchtitan kept "all". Every update session from them restored and postprocessed a speculative draft that received no weights, and p2p connect raised NotImplementedError for EAGLE3, STANDALONE, separate draft checkpoints and multi-layer EAGLE drafts. The weight checker snapshotted and reset_tensors every runner with a hard-coded "all" at engine init, so a draft the sync leaves out stayed scrambled for the run unless --check-weight-update-selector target was passed by hand. weight_update_selector(args) now derives the selector from --enable-mtp-training, the MTP training intent the parent commit gates Megatron's MTP layers on and rejects on FSDP and torchtitan, for every backend and for each trainer's args. --check-weight-update-selector defaults to it, and the server cell snapshots and resets exactly the runners the checker compares. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

claude

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

Brings in 045f983 (MTP layers built only under --enable-mtp-training), which the selector rule here already follows. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

guapisolo · 2026-10-13T05:51:28Z

@claude review always

…tor-from-mtp-intent Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

claude

I reviewed this PR and didn't find any bugs. Because it reworks weight-sync selector logic shared across backends (Megatron/FSDP/torchtitan), the updater's connect/update_weights paths, and the server cell's equality-check gating, a human look would still be worthwhile.

What was reviewed: the new weight_update_selector(args) in session.py and its call sites in updater.py; the removal of the per-iterator weight_update_selector attribute from hf_weight_iterator/init.py and megatron_utils/hf_weight_iterator.py; server_cell.py's snapshot/reset_tensors now using args.check_weight_update_selector instead of a hardcoded "all"; and the arguments.py defaulting logic (including import-cycle risk, which is absent). Also checked the multi-policy case where a trainer's own enable_mtp_training could diverge from the base args used to default check_weight_update_selector — this is the PR's documented, accepted behavior (checker uses base args, each trainer's sync uses its own args), not a new bug.

Extended reasoning...

The diff relocates the weight-sync "which engine runners to cover" decision from a per-backend class attribute into a single session-level function keyed on --enable-mtp-training, threads it through WeightUpdater.connect/update_weights, and makes server_cell's post-update equality check (snapshot/reset_tensors) follow the same selector instead of a hardcoded "all"; arguments.py now defaults --check-weight-update-selector to that function's result. No injection/auth/data-exposure surface is touched; the risk is purely behavioral correctness of weight sync and the equality checker across Megatron/FSDP/torchtitan and multi-policy setups. Decided to defer rather than approve because this is cross-backend, behavior-changing infra code with a documented but non-trivial multi-policy edge case, even though no bug was found.

…tor-from-mtp-intent Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…trained MTP layers Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

guapisolo · 2026-10-13T06:45:19Z

/rerun-test tests/e2e/megatron/test_qwen3_5_35B_A3B/test_spec_nor3_fp8_sgl_tp4_meg_tp2cp2.py

guapisolo · 2026-10-13T06:45:20Z

/rerun-test tests/e2e/megatron/test_glm47_flash/test_spec_mtptrain_nor3_bf16_sgl_tp2_meg_tp1pp2_p2p.py

guapisolo · 2026-10-13T06:45:21Z

/rerun-test tests/e2e/megatron/test_glm47_flash/test_nospec_nor3_fp8_deepep_sgl_dpattn1x4_meg_tp2cp1_bshd.py

guapisolo · 2026-10-13T06:45:23Z

/rerun-test tests/e2e/megatron/test_mimo_7B_mtp_only_grad.py

guapisolo · 2026-10-13T06:45:24Z

/rerun-test tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py

github-actions · 2026-10-13T06:45:51Z

✅ tests/e2e/megatron/test_qwen3_5_35B_A3B/test_spec_nor3_fp8_sgl_tp4_meg_tp2cp2.py passed on stage-c-4-gpu-h200 in 1h23m04s — workflow run

github-actions · 2026-10-13T06:46:03Z

✅ tests/e2e/megatron/test_glm47_flash/test_nospec_nor3_fp8_deepep_sgl_dpattn1x4_meg_tp2cp1_bshd.py passed on stage-c-4-gpu-h200 in 1h40m14s — workflow run

github-actions · 2026-10-13T06:46:15Z

✅ tests/e2e/megatron/test_mimo_7B_mtp_only_grad.py passed on stage-c-4-gpu-h200 in 1h34m31s — workflow run

github-actions · 2026-10-13T06:46:29Z

✅ tests/e2e/megatron/test_glm47_flash/test_spec_mtptrain_nor3_bf16_sgl_tp2_meg_tp1pp2_p2p.py passed on stage-c-8-gpu-h200 in 12m41s — workflow run

github-actions · 2026-10-13T06:46:39Z

✅ tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py passed on stage-c-8-gpu-h200 in 12m58s — workflow run

claude

Code review found no issues

No high-confidence issues detected in this change.

Picks up main's fix for the miles.utils.arguments <-> weight_update.session import cycle (#3376 + #4043) that failed every stage-a shard at conftest import. No conflicts; model.py merges main's restored_trained_iteration start-rollout logic around this branch's native LoRA adapter load.

guapisolo requested review from Shi-Dong, Zhichenzzz, fzyzcjy, jybsuper, maocheng23, yueming-yuan and yushengsu-thu as code owners October 13, 2026 04:42

claude Bot reviewed Oct 13, 2026

View reviewed changes

Merge PR-A post-hook into the selector branch

640144d

Brings in 045f983 (MTP layers built only under --enable-mtp-training), which the selector rule here already follows. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

guapisolo added this pull request to stack #4045 October 13, 2026 05:08

guapisolo and others added 2 commits October 12, 2026 23:07

Merge jiajun/mtp-train-by-layer-count into jiajun/weight-update-selec…

e011945

…tor-from-mtp-intent Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

test(megatron-config): cover a trainer override of mtp_num_layers

aca6d53

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

claude Bot reviewed Oct 13, 2026

View reviewed changes

guapisolo and others added 2 commits October 12, 2026 23:21

Merge jiajun/mtp-train-by-layer-count into jiajun/weight-update-selec…

cb0530a

…tor-from-mtp-intent Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

test(mtp): cover an explicit zero layer count and a trainer dropping …

05adf84

…trained MTP layers Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

guapisolo added the run-ci label Oct 13, 2026

claude Bot reviewed Oct 13, 2026

View reviewed changes

yueming-yuan approved these changes Oct 13, 2026

View reviewed changes

guapisolo merged commit f240fcc into main Oct 13, 2026
45 checks passed

guapisolo deleted the jiajun/weight-update-selector-from-mtp-intent branch October 13, 2026 21:24

bingxche mentioned this pull request Oct 24, 2026

[CI Monitor] Daily Report - 2026-10-24 bingxche/sglang-ci-bot#243

Open

Conversation

guapisolo commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Symptom & Reproduction

Root Cause

Fix

Verification

Review Focus

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

guapisolo commented Oct 13, 2026

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

guapisolo commented Oct 13, 2026

Uh oh!

guapisolo commented Oct 13, 2026

Uh oh!

guapisolo commented Oct 13, 2026

Uh oh!

guapisolo commented Oct 13, 2026

Uh oh!

guapisolo commented Oct 13, 2026

Uh oh!

github-actions Bot commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

guapisolo commented Oct 13, 2026 •

edited

Loading

github-actions Bot commented Oct 13, 2026 •

edited

Loading

github-actions Bot commented Oct 13, 2026 •

edited

Loading

github-actions Bot commented Oct 13, 2026 •

edited

Loading

github-actions Bot commented Oct 13, 2026 •

edited

Loading

github-actions Bot commented Oct 13, 2026 •

edited

Loading