fix: provide lm_head.weight for tied VLM decoders in preprocess_weights - #295
Conversation
When tie_word_embeddings=True, the Qwen VL 3-model split decoder models have lm_head.weight and embed_tokens.weight aliased to the same nn.Parameter. The ONNX graph creates initializer nodes for both names, and the weight loader requires data for all initializers. Previously, preprocess_weights discarded lm_head.weight with the comment 'tied at graph level; no separate entry needed'. This was incorrect: while the graph correctly shares the tensor, the weight loading validation (in _check_weights) requires every initializer to have data. Fix: copy embed_tokens.weight to lm_head.weight in preprocess_weights when tie_word_embeddings=True, matching the pattern used by CausalLMModel and Gemma4's VLM decoder. Affected classes: - Qwen25VLCausalLMModel (Qwen2.5-VL 3-model composite) - Qwen25VLDecoderModel (Qwen2.5-VL standalone decoder) - Qwen3VL3ModelCausalLMModel (Qwen3-VL 3-model composite) - Qwen3VLDecoderModel (Qwen3-VL standalone decoder) Tested: - L1: all Qwen VL graph build tests pass (91 passed, 0 VL failures) - L5: Qwen3-VL-4B-Instruct f16 ort-genai export succeeds (7.6G decoder + 742M embedding + 638M vision) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Pull request overview
This pull request fixes weight loading for Qwen2.5-VL and Qwen3-VL 3-model (ORT GenAI) builds when tie_word_embeddings=True by ensuring the decoder’s lm_head.weight initializer receives data even when the checkpoint only provides embed_tokens.weight.
Changes:
- In the 3-model composite preprocessors, copy the token embedding weight into
decoder.lm_head.weightwhen embeddings are tied. - In standalone decoder preprocessors, stop discarding
lm_head.weightfor tied embeddings and instead synthesize it fromembed_tokens.weightwhen missing.
|
I don't think we should copy them. The weights should be shared |
Replace inline lm_head.weight handling with the tie_word_embeddings() utility from _weight_utils.py. This ensures the same tensor object is referenced by both embed_tokens.weight and lm_head.weight keys in the state dict, allowing apply_weights' id()-based deduplication to unify them into a single ONNX initializer. For 3-model split composites (Qwen25VLCausalLMModel, Qwen3VL3Model), the embed/lm_head keys use 'decoder.' prefix so we pass custom key names to tie_word_embeddings(). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
7eaaa98 to
57447ec
Compare
Document how the Python object identity created by tie_word_embeddings flows through apply_weights (data_ptr()-based dedup) to produce a single ONNX initializer shared by both Gather and MatMul nodes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
titaiwangms
left a comment
titaiwangms
left a comment
There was a problem hiding this comment.
Review summary
Direction is right and existing tests pass cleanly (2736 unit tests in 32s, 115 qwen tests), but adversarial trace turned up one path that's still silently broken, plus a coverage gap that lets it slip through the suite.
Major
-
Qwen3VLDecoderModel.preprocess_weights()is silently still broken. The rename loop (qwen_vl.py:843-852) strips bothmodel.andlanguage_model.unconditionally, so for an HF key likemodel.language_model.embed_tokens.weightthe dict ends up withembed_tokens.weight(nomodel.prefix). The new calltie_word_embeddings(renamed)uses defaultsembed_key="model.embed_tokens.weight",head_key="lm_head.weight"— neither is present, so the helper silently no-ops andlm_head.weightis never created. The standalone Qwen3-VL decoder build with tied weights will hit the sameValueError: Component 'decoder' has 1 initializer(s) without weights: 'decoder.lm_head.weight'this PR is trying to fix.
Either:tie_word_embeddings(renamed, embed_key="embed_tokens.weight", head_key="lm_head.weight")
or rework the rename so
model.embed_tokens.weightlands inrenamed(so it matches the standalone ONNX initializer name and keeps default-args semantics consistent withQwen25VLDecoderModel). -
No regression test for any of the four fixed paths. The full suite passes, but:
- No test invokes
preprocess_weightsonQwen25VLCausalLMModel,Qwen25VLDecoderModel,Qwen3VL3ModelCausalLMModel, orQwen3VLDecoderModel. - All qwen-VL graph-build tests use
tie_word_embeddings=False. - Graph-build tests don't call
preprocess_weightsat all — the failure surfaces atapply_weightstime. - The only qwen-VL
preprocess_weightstest (test_qwen35_vl_preprocess_weights_model_prefix) is for a different class inqwen35.pyand doesn't set tying.
Add ~4 cheap co-located tests insrc/mobius/models/qwen_vl_test.pywith a tiny fake state dict andtie_word_embeddings=True, assertingdecoder.lm_head.weight in renamedanddata_ptr()identity with the embed. These would have caught both the original bug and bug #1 above.
- No test invokes
-
Missing "why" comment at the composite tie sites. The deleted comment was wrong; nothing replaced it. The non-obvious reason —
onnxscriptqualifies parameter names by module path, so the__init__-time aliasdecoder.lm_head.weight = decoder.model.embed_tokens.weightdoesn't cross composite module boundaries — is the entire reason this PR exists. Without a comment, the next person will re-delete the block. Suggested:# onnxscript qualifies params by module path, so the in-tree alias set in # Qwen25VLDecoderModel.__init__ does not cross composite boundaries. Establish # the identity here so apply_weights sees a single data_ptr() across both # initializers. if self.config.tie_word_embeddings: tie_word_embeddings(renamed, embed_key=..., head_key=...)
Minor
-
_weight_utils.tie_word_embeddingssilent no-op masks routing bugs. Whentie=Trueand both configured keys are absent, the helper does nothing — which is exactly what surfaces as bug #1. Recommend: raise (or at least log a clear warning naming both keys) when neither key is present and the caller saidtie=True. Tightly coupled to this PR; worth rolling in. -
_weight_utils.tie_word_embeddingsdocstring says "copies"; the implementation is identity assignment.state_dict[head_key] = state_dict[embed_key]is a Python reference assignment, same object, samedata_ptr(). The PR's correctness depends on identity, not copy. The next maintainer reading "copies" will reasonably "fix" it toclone()and silently break the tying. Update the docstring to say something like "assigning it to the same Python tensor object so downstream consumers see onedata_ptr()." (Note: the diff/branch does not currently change_weight_utils.py— please ship the docstring fix with this PR.) -
Composite preprocess drops
lm_head.*before the helper runs. InQwen25VLCausalLMModel:126-128andQwen3VL3ModelCausalLMModel:788-790, whentie=Truethe loop skipslm_head.*keys entirely. If a (non-standard but legal) checkpoint stored onlylm_head.weightand not the embed, that value is dropped beforetie_word_embeddingscould backfill the embed. Defense-in-depth: keeplm_head.*inrenamedregardless and let the helper normalize. -
Qwen2VLCausalLMModel.preprocess_weights(line 414) not migrated. After this PR the file has three different idioms for the same concept:- inline assignment (Qwen2VL)
- post-loop helper with explicit keys (composite Qwen2.5/3-VL)
- post-loop helper with defaults (standalone Qwen2.5/3-VL decoders)
Plus the legitimately-different.pop()in single-modelQwen3VLCausalLMModel. Migrate Qwen2VL to the helper now, or add a# TODO: align with Qwen2.5/Qwen3 patterncomment so the asymmetry is flagged rather than invisible.
-
Stale comment at
Qwen3VLCausalLMModel:718. The wording# lm_head.weight is tied at graph level; discard any separate checkpoint entry.is word-for-word the comment this PR deletes as wrong elsewhere. For this single-model class it's accurate (in-tree alias survives onnxscript), but the wording reads ambiguously after this PR. Tighten to make the single-model-vs-composite distinction explicit, e.g. "Single-model build: onnxscript preserves the in-tree alias, so thelm_headcheckpoint entry is redundant — pop it."
Nit
embed_key/head_keystrings ("decoder.model.embed_tokens.weight"/"decoder.lm_head.weight") duplicated in two composites — module-level constants if a third composite lands.Qwen3VLDecoderModel.preprocess_weightsdocstring is one line; sisterQwen25VLDecoderModelhas a four-line docstring covering HF key format and standalone-vs-composite distinction. Match the level.# Handle weight tyingatqwen_vl.py:191replaces a wrong-but-explicit comment with a content-free one. One line: "HF checkpoint omits lm_head.weight when tied; fill it in for the initializer."
Praise
- Right call using the existing helper instead of growing a new pattern.
- Both composite paths and both standalone-decoder paths get the same treatment in one PR — good symmetry (modulo bug #1).
- Commit messages name the wrong assumption, the failure mode, and the L1/L5 verification — exactly the right context for a non-obvious weight-loading fix.
QA / test run
python -m pytest tests/build_graph_test.py tests/cli_test.py src/ -q -k "not phi4mm and not apply_weights_unknown" --tb=short -n auto→ 2736 passed, 41 skipped, 0 failed in 32.4s.- Qwen-only subset → 115 passed, 24 skipped, 0 failed in 12.6s.
- Lint cannot run in the test env:
ruff 0.12.12installed,pyproject.toml:133referencesRUF067requiringruff 0.15.11. Pre-existing environment mismatch, not caused by this PR.
Net: the fix's intent is correct and the existing 3-model composite paths trace cleanly, but bug #1 (Qwen3VLDecoderModel) and the absent regression coverage mean this still slips the same class of failure. Worth one more pass.
- Fix Qwen3VLDecoderModel tie_word_embeddings key mismatch: after stripping language_model., keys are embed_tokens.weight not model.embed_tokens.weight. Pass explicit embed_key/head_key. - Add 8 regression tests (qwen_vl_test.py) verifying lm_head.weight presence and data_ptr() identity for all 4 model classes. - Add warning to tie_word_embeddings when both keys are absent. - Add 'why' comments explaining onnxscript module boundary limitation at all tie_word_embeddings call sites. - Tighten Qwen3VLCausalLMModel comment to explain single-model-vs- composite distinction. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Fail explicitly instead of silently doing nothing when both embed_key and head_key are absent from state_dict. This catches key name mismatches early (e.g. after prefix stripping). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Summary
Fix VLM decoder weight loading for models with
tie_word_embeddings=Trueby using the sharedtie_word_embeddings()utility.Problem
When a VLM model (e.g. Qwen-VL, InternVL) has
tie_word_embeddings=True, HuggingFace stores only the embedding weight in the checkpoint —lm_head.weightis absent. The previous VLM decoderpreprocess_weightscode discarded any weight not under thelanguage_model.prefix, solm_head.weightwas never populated. The ONNX model then had an uninitialized LM head, producing garbage logits.Fix
Use
tie_word_embeddings()from_weight_utils.pyafter stripping thelanguage_model.prefix. This copies the embedding tensor reference to thelm_head.weightkey when it is missing.How ONNX Initializer Sharing Works End-to-End
The weight-sharing mechanism spans three layers:
tie_word_embeddings(state_dict)— Whenlm_head.weightis missing, assignsstate_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"]. Both keys now point to the same Python tensor object (samedata_ptr()).apply_weights(model, state_dict)in_weight_loading.py— Iterates over the state dict and assigns tensors to ONNX initializers. It trackstensor.data_ptr()to detect shared storage. When it encounterslm_head.weightand sees the samedata_ptr()as the already-assignedmodel.embed_tokens.weight, it:initializer.replace_all_uses_with(canonical)to redirect all graph usesmodel.graph.initializersResult — The saved ONNX file contains a single copy of the embedding table, used by both the Gather (embedding lookup) and MatMul (LM head projection) nodes. No duplication.
Changes
src/mobius/models/qwen_vl.py: Usetie_word_embeddings()inpreprocess_weightssrc/mobius/_weight_utils.py: Enhanced docstring explaining the full sharing mechanism