Weight Loader v2 PR4: Qwen3.5 / Qwen3-Next language paths - #1
Draft
MartinHua wants to merge 7 commits into
Draft
Conversation
added 7 commits
July 31, 2026 23:53
Adds the Qwen-family dispatch primitives the PR2 auto_loader does not carry (packed GDN stacked mapping, fused-expert fan-out, Qwen3.5 name normalization) and migrates qwen3_next.py behind SGLANG_ENABLE_WEIGHT_LOADER_V2. Qwen3-Next inlines qkv_proj/o_proj onto the attention decoder layer rather than a nested attention module, so that layer owns the stacked dispatch and must re-enter the MoE/dense block's own loader explicitly: the walker hands a whole subtree to the first module defining load_weights, so mlp weights would otherwise never reach expert dispatch. split_submodule_weights makes that hand-off reusable. FusedExpertDispatch raises on a missing runtime target instead of returning it, matching the PR2 ExpertParamsDispatch contract and the plan's "missing targets must fail before execution" invariant. Refs: sgl-project#31051
Adds mapping/dispatch coverage for the PR4 Qwen3-Next path and forwards the loaded-name set from the MTP wrapper so it matches the base loader contract asserted by test_wrapper_base_dispatch. The load-bearing case is the attention-layer hand-off: with a plain stacked dispatch on the decoder layer, every mlp.* key is swallowed by the layer and silently dropped instead of reaching expert dispatch. Verified red against that implementation and green with the explicit hand-off. The GDN case pins the trailing-dot disambiguation that keeps a fused in_proj_qkvz/in_proj_ba checkpoint from being rewritten onto itself with a bogus shard id. Refs: sgl-project#31051
The v2 preprocessing generator rewrote mlp.shared_expert.* to the fused routed slot before the walker ran. That was correct against the earlier PR2 draft, where the block loader did not handle shared experts, but the current Qwen2MoeSparseMoeBlock.load_weights maps shared_expert.gate_proj/down_proj/up_proj straight into slot _ckpt_num_experts itself. Renaming first produces experts.<slot>.* , which matches no entry in the block's expert mapping (make_expert_params_mapping only covers ids < num_experts), so load_moe_sparse_block_weights rejected it: ValueError: No parameter named 'experts.8.gate_proj.weight' in Qwen2MoeSparseMoeBlock. i.e. every shared-expert-fusion checkpoint failed to load under v2. Dropping the remap lets the block's own mapping match. The block also selects an empty dense_stacked when fusion is on, so gate_proj must stay un-rewritten for it to route correctly. The legacy loader keeps its own rename and is unaffected. Refs: sgl-project#31051
Mechanical cleanups from reviewing the three preceding commits; no behavior change to the Qwen3-Next v2 path (fan-out order and the shared-expert passthrough are unchanged and still covered by the existing tests). - Drop QWEN35_GDN_STACKED_MAPPING, QWEN35_STACKED_MAPPING and normalize_qwen35_weight_name. They were ported ahead of qwen3_5.py and had no reference anywhere, not even in tests; they belong in the commit that migrates that model. FusedExpertDispatch stays: it is a named PR4 deliverable and is exercised by unit tests. - FusedExpertDispatch no longer resolves its loader via getattr(param, "weight_loader", default_weight_loader). default_weight_loader accepts only (param, tensor), so the fallback could never satisfy the shard_id/expert_id call and would have raised TypeError instead of degrading. Fused expert params always carry a FusedMoE loader, so access it directly and fail loudly, matching StackedParamsDispatch in the same module. - fan_out_to_experts becomes a protected instance method: it was static yet every caller passed self.num_experts. Its helpers are keyword-only, and the load_with_stacked_dispatch call sites pass mapping by keyword, per the keyword-argument convention in general-code-style. - iter_qwen3_next_checkpoint_weights now requires params_dict. It was optional and defaulting to None silently disabled the unit-scale verification, which is a correctness guard callers should not be able to skip. Refs: sgl-project#31051
Migrates Qwen3_5ForCausalLM behind SGLANG_ENABLE_WEIGHT_LOADER_V2, mirroring the Qwen3-Next structure: the gated-delta-net mixer owns packed GDN dispatch, and the attention decoder layer owns stacked qkv dispatch plus an explicit hand-off of mlp.* to the MLP block's own loader (Qwen3.5 inlines qkv_proj onto the layer, and the walker gives a whole subtree to the first module defining load_weights). iter_qwen3_5_text_checkpoint_weights reuses main's QWEN3_5_KV_SCALE_MAPPER rather than re-implementing the ModelOpt kv-scale remap by hand, and drops the MTP and vision streams, which own their own loaders. Shared-expert keys are passed through unrewritten so the MoE block can map them into the fused slot. Qwen3_5MoeForCausalLM is untouched: it fully overrides load_weights with no super() call, so it stays on the legacy path. Its fused-expert checkpoint format (experts.gate_up_proj stacked across experts) needs routing that the MoE block loader does not yet provide, so it is deliberately left for a follow-up rather than guessed at. Refs: sgl-project#31051
Migrates Qwen3_5ForCausalLMMTP behind SGLANG_ENABLE_WEIGHT_LOADER_V2 by splitting the checkpoint instead of duplicating mapping tables: the wrapper loads its own fc and pre-fc norms, and forwards the rest of the mtp.* branch to the draft body, which is a full Qwen3_5ForCausalLM and already owns stacked / GDN / expert dispatch. That replaces a 200-line copy of those tables with a 33-line split, and means the draft automatically picks up whatever the body's loader supports. Names come back from the body re-prefixed with "model." so they match this module's own parameter names. embed_tokens and lm_head are not checkpoint-loaded on this path: they are installed from the target by set_embed_and_head / set_lm_head_from_target. The draft branch is selected by the "mtp." prefix rather than a bare "mtp" substring, so a malformed key cannot be forwarded to the body with a stale path. The legacy loader is unchanged and still selected when the env gate is off. Fused-expert checkpoints reach the body's MoE block loader, which does not yet handle the stacked experts.gate_up_proj form — that gap is shared with Qwen3_5MoeForCausalLM and is deliberately left for a follow-up. Refs: sgl-project#31051
…ences Comparing PR4 against the two reference points — PR1's state-dict equivalence harness and PR2's qwen3.py dense migration — surfaced coverage that neither PR4 nor the repo had. Shared v2 invariants, none of which were tested anywhere before: - filter_pp_weights drops layers outside [start_layer, end_layer) and passes through keys with no parseable layer index. - The walker rejects an unexpected checkpoint name instead of dropping it, which is the "no silent drops" contract v2 exists for. - ignore_unexpected_suffixes is specific rather than blanket: a declared suffix is ignored, an undeclared one still raises. Qwen3.5 depends on this for .kv_scale. Equivalence cases for both hybrid families, which PR4 had none of. No tiny-random fixture is published for either, so the two harnesses differ deliberately: the manual one resolves cache-only and can name the real 4B/80B checkpoints (it skips unless cached), while the registered e2e resolver downloads, so its defaults are unresolvable placeholders that skip with a message naming the override env var. Also documents why qwen3_next does not list ".kv_scale" in ignore_unexpected_suffixes while qwen3_5 does: the qwen3_next generator already resolves every homeless "*_scale" key, so none reach the walker. Adding the suffix there would be inert, and the asymmetry is a consequence of the two generators rather than an oversight. Refs: sgl-project#31051
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
Weight Loader v2 PR4 from sgl-project#31051 — Qwen3.5 / Qwen3-Next language paths.
Base: stacks on PR2 (sgl-project#32565).
feat/weight-loader-v2-pr2-rebasedis PR2's commit rebased onto currentmain— PR2's own branch is ~209 commits behind and targeting it renders a 1413-file diff. Rebase onto PR2 and drop the base commit once it lands. The dependency is real, not ordering: both models import their MLP fromqwen2_moe.py, and this migration callsself.mlp.load_weights(...), which only exists in PR2.Modifications
Qwen3NextForCausalLMQwen3NextForCausalLMMTPsuper()→ parent gateQwen3_5ForCausalLMQwen3_5ForCausalLMMTPQwen3_5MoeForCausalLMauto_loader.pyaddsQWEN3_NEXT_GDN_STACKED_MAPPING(packed gated-delta-net, shared by both families),FusedExpertDispatch,split_submodule_weights. Qwen3.5 reuses main'sQWEN3_5_KV_SCALE_MAPPERinstead of re-implementing the ModelOpt kv-scale remap.qwen3_5_mtp.pyforwards themtp.branch to the draft body — itself aQwen3_5ForCausalLMthat already owns dispatch — replacing a 200-line copy of the mapping tables with a 33-line split. Legacy paths are untouched and still selected when the gate is off.Two non-obvious points. (a) Both families inline
qkv_proj/o_projonto the decoder layer, so that layer is the only place where the split checkpoint names and the fused runtime param are both visible and must defineload_weights— butAutoWeightsLoader._load_modulegives a whole subtree to the first module defining it, somlp.needs explicit re-entry or every expert key is swallowed there. (b)Qwen2MoeSparseMoeBlock.load_weightsmapsshared_expert.*into the fused slot itself and selects an emptydense_stackedwhen fusion is on; pre-renaming tomlp.experts.<slot>.*yields a key matching no expert mapping, which the block rejects (ValueError: No parameter named 'experts.8.gate_proj.weight') — i.e. every shared-expert-fusion checkpoint fails to load. v2 passes those names through unrewritten.Accuracy Tests
18 CPU tests (
base-b-test-cpu) intest/registered/unit/model_loader/test_pr4_qwen3_{next,5}_v2.py: name remapping; MTP filter/rename both directions and the MTP body/wrapper split; the trailing-dot disambiguation that stops a fusedin_proj_qkvz/in_proj_bacheckpoint being rewritten onto itself with a bogus shard id; the decoder-layer MLP hand-off; shared-expert passthrough; fused-expert fan-out order; raise-on-missing-target. Plus three shared v2 invariants with no prior coverage anywhere in the repo:filter_pp_weightsrange/passthrough, the walker rejecting an unexpected name (the "no silent drops" contract), andignore_unexpected_suffixesbeing specific rather than blanket.Equivalence cases for both families added to PR1's state-dict harness and PR2's e2e harness as
hybrid_gdn_dense/hybrid_mamba. No tiny-random fixture exists for either, so defaults differ by resolver:test/manual/…_equiv.pyis cache-only and names the real 4B/80B checkpoints (skips unless cached); the registered e2e resolver downloads, so its defaults are unresolvable placeholders that skip with a message naming the override env var.Not verified — no GPU available to the author:
SGLANG_WEIGHT_LOADER_V2_HYBRID_GDN_DENSE_MODEL=Qwen/Qwen3.5-4B pytest test/manual/test_weight_loader_v2_equiv.py -k hybrid_gdn_denseFusedExpertDispatchnow raises rather than returning a missing target; PP is filtered upstream and EP still materializesw13_weightper rank, so this should be safe but wants one EP/PP loadSpeed Tests and Profiling
N/A — loading only, behind an env gate that is off by default.
Checklist
black26.1.0,isort7.0.0),py_compilecleanOpen questions
Qwen3_5MoeForCausalLMfused-expert routing. PR2's block loader knows only the per-expert form, but Qwen3.5 MoE checkpoints can stack all experts into oneexperts.gate_up_proj(w1/w3 packed on dim −2). EitherFusedExpertDispatchat decoder-layer level (faithful mirror of the legacyload_fused_expert_weights, but needs loaders on both decoder layer types plus the block's_ckpt_num_experts), or un-fusing in the weight stream (cleaner, but the key spelling can't be confirmed without a real fused checkpoint).qwen3_5_mtp.pyinherits the gap through the body, so one fix covers both. The class is insulated meanwhile — it overridesload_weightswith nosuper()call.qwen3_5_text.pyis affected but absent from PR4's file list.EntryClass = [Qwen3_5MoeForCausalLM, Qwen3_5ForCausalLM], delegating viaself.model.load_weights(body_weights)(line 164) — a bound delegate, so it auto-enters v2; same for MiniCPM-V 4.6 (minicpmv.py:1592). Per the plan's "parent migration includes affected downstream paths" invariant these want env-off/env-on routing smoke. Should the file and the 187 denominator be updated?qwen3_5/qwen3_nextshould leavePROTECTED_MODEL_PREFIXESintest_weight_loader_v2_model_contracts.py.Tracking: sgl-project#31051