Skip to content

Describe cross-component tied weights - #747

Merged
xiaoyu-work merged 5 commits into
mainfrom
fix/cross-component-tied-weights
Sep 26, 2026
Merged

xiaoyu-work merged 5 commits into
mainfrom
fix/cross-component-tied-weights

Conversation

@xiaoyu-work

@xiaoyu-work xiaoyu-work commented Sep 23, 2026 •

Copy link
Copy Markdown
Member

Summary

  • expose typed cross-component shared-weight metadata through the component manifest and inspect_components()
  • infer tied embedding/LM-head endpoints only when explicit source-path aliases identify a unique pair; for unquantized models the nested text tie flag takes precedence over the parent, while quantization metadata can explicitly restore a tie
  • materialize canonical packed sidecars into split component aliases before model-specific preprocessing

Validation

Companion: microsoft/Olive#2684 handles component quantization and assembly.

Expose canonical and alias parameter endpoints through component inspection so optimizers can coordinate independent builds. Materialize Gemma4 packed token embeddings into the split decoder LM head during ONNX export.

Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
@github-actions

github-actions Bot commented Sep 23, 2026 •

Copy link
Copy Markdown

Performance Comparison

Comparing 783e03b5 → fb6b087c

Model Metric Baseline Current Delta
bert (feature-extraction) model_size_bytes 359 KB 359 KB +0.0% ⚪
bert (feature-extraction) num_nodes 68 68 +0.0% ⚪
falcon model_size_bytes 364 KB 364 KB +0.0% ⚪
falcon num_nodes 66 66 +0.0% ⚪
gemma2 model_size_bytes 428 KB 428 KB +0.0% ⚪
gemma2 num_nodes 105 105 +0.0% ⚪
gpt2 model_size_bytes 324 KB 324 KB +0.0% ⚪
gpt2 num_nodes 54 54 +0.0% ⚪
llama model_size_bytes 425 KB 425 KB +0.0% ⚪
llama num_nodes 60 60 +0.0% ⚪
llama (static-cache) model_size_bytes 425 KB 425 KB +0.0% ⚪
llama (static-cache) num_nodes 56 56 +0.0% ⚪
mamba (ssm-text-generation) model_size_bytes 296 KB 296 KB +0.0% ⚪
mamba (ssm-text-generation) num_nodes 94 94 +0.0% ⚪
phi3 model_size_bytes 421 KB 421 KB +0.0% ⚪
phi3 num_nodes 58 58 +0.0% ⚪
phi3 (static-cache) model_size_bytes 421 KB 421 KB +0.0% ⚪
phi3 (static-cache) num_nodes 54 54 +0.0% ⚪
qwen2 model_size_bytes 425 KB 425 KB +0.0% ⚪
qwen2 num_nodes 60 60 +0.0% ⚪
qwen2 (static-cache) model_size_bytes 425 KB 425 KB +0.0% ⚪
qwen2 (static-cache) num_nodes 56 56 +0.0% ⚪
qwen3_5_moe (hybrid-text-generation) model_size_bytes 506 KB 506 KB +0.0% ⚪
qwen3_5_moe (hybrid-text-generation) num_nodes 265 265 +0.0% ⚪
qwen3_5_text (hybrid-text-generation) model_size_bytes 458 KB 458 KB +0.0% ⚪
qwen3_5_text (hybrid-text-generation) num_nodes 127 127 +0.0% ⚪
qwen3_5_vl (hybrid-qwen-vl) model_size_bytes 977 KB 977 KB +0.0% ⚪
qwen3_5_vl (hybrid-qwen-vl) num_nodes 450 450 +0.0% ⚪
t5 (seq2seq) model_size_bytes 836 KB 836 KB +0.0% ⚪
t5 (seq2seq) num_nodes 176 176 +0.0% ⚪
whisper (speech-to-text) model_size_bytes 1008 KB 1008 KB +0.0% ⚪
whisper (speech-to-text) num_nodes 128 128 +0.0% ⚪

No performance regressions.

@github-actions

github-actions Bot commented Sep 23, 2026 •

Copy link
Copy Markdown

🏗️ Architecture Diff

Comparing 783e03b5 → fb6b087c

Model Sub-model Changes Status
gemma4 (gemma4) decoder 0 ⚪
gemma4 (gemma4) embedding 0 ⚪
gemma4 (gemma4) vision_encoder 0 ⚪

No architecture changes detected. ✅


Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed)

Use Hugging Face tie_word_embeddings metadata together with explicit component source aliases to derive standard embedding/head sharing. Remove the Gemma4-specific shared-weight resolver while retaining its split export materialization.

Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🟡 Changes recommended

Critical typing and standalone component tie-handling issues remain unresolved, with additional unified-model test coverage needed.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1 High severity

Open (1)
What changed in this PR

Adds cross-component shared-weight metadata and Gemma4 tied embedding/LM-head materialization for split exports.

Changes:

  • Adds shared-weight metadata and validation to inspect_components().
  • Declares Gemma4 tied embeddings from configuration.
  • Materializes packed embedding sidecars for decoder LM heads.
  • Adds targeted inspection and quantized-weight tests.
File Summary
src/​mobius/​models/​gemma4.py Gemma4 shared-weight declarations and LM-head materialization
src/​mobius/​models/​gemma4_test.py Quantized tied-weight regression coverage
src/​mobius/​_inspect.py Shared-weight metadata models and validation
src/​mobius/​_inspect_test.py Inspection contract tests
src/​mobius/​__init__.py Public API exports

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread src/mobius/models/gemma4.py Outdated
Carry inferred shared-weight groups on ComponentManifest and expand packed aliases before architecture-specific preprocessing. Remove Gemma4-specific materialization so standard tied component models use the common loading path.

Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🟡 Changes recommended

Quantization-only tied-weight declarations are not detected, preventing packed LM-head sidecar materialization.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 Medium severity

Open (1)
Resolved since last review (1)

Comment thread src/mobius/_component_manifest.py Outdated
@titaiwangms

Copy link
Copy Markdown
Contributor

Cross-repository review of this PR together with microsoft/Olive#2684 (reviewed head 4653707). The Gemma4 shared-weight declaration and packed-alias materialization are in the right layers. Two points need clarification or coverage before treating this as generic tied-weight support:

  1. Tie detection should agree with the export configuration. src/mobius/_component_manifest.py:239-251 returns true if any of the parent/nested configs has tie_word_embeddings=True. ArchitectureConfig resolves the selected text config's explicit value before falling back to the parent (src/mobius/_configs/_base.py:1196-1203). With parent true and text config explicitly false, inspect_components() can declare a tie that the exported model does not use; Olive will subsequently reject it at its actual-parameter-identity check. Please use the same effective tie decision as model construction and test conflicting parent/nested values.
  2. The inferred contract currently depends on explicit module aliases. _aliased_component_endpoint() at src/mobius/_component_manifest.py:254-273 inspects only source_path_aliases, not source_paths. Gemma3 and Qwen3-VL, for example, expose separate embed_tokens and lm_head component source paths but do not declare these aliases, so this inference produces no shared_weights for them. If this is intentionally Gemma4-only for now, please state that limitation; if the intent is general tied-word-embedding support, cover unique source-path endpoints and add an inspect test.

Cross-PR contract: Mobius recognizes an output module named output (_component_manifest.py:32-38), while Olive's new _OUTPUT_HEAD_MODULE_NAMES in hf_component_assembly.py:68-73 omits it. Please align the consumer's name set or narrow the producer's advertised contract; otherwise a future output.weight endpoint can assemble with lm_head=False.

This was a static review of the two PR diffs and immediate consumers, not an end-to-end Olive checkpoint -> Mobius ONNX runtime test. I found no confirmed blocking defect in this PR's tested Gemma4 path; the cross-build failure cases are posted on the Olive PR.

Inspect both mapping and object quantization declarations, including parsed configuration, so a model-level false flag cannot hide a tied packed embedding/head relationship.

Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
@titaiwangms

Copy link
Copy Markdown
Contributor

Follow-up review at head 3c4a8dc (paired with microsoft/Olive#2684): I'm holding approval for now.

The new tests cover ties recorded only in quantization metadata, including the important case where the model-level flag is false but the packed checkpoint records a tie. However, _config_ties_word_embeddings() (src/mobius/_component_manifest.py:239-263) still returns any(...) across model, nested text, and quantization declarations. For an unquantized source with conflicting parent/nested flags, this can report a cross-component tie even when the effective export configuration selects the explicitly false flag. Please establish the precedence for that case and cover it with a conflicting-value test (or explain why such a source cannot occur for the supported architecture). Olive's runtime parameter-identity check protects against silent deduplication, but a false declaration would still cause a valid build to fail.

The endpoint inference still uses explicit source_path_aliases, not component source_paths. That's fine if the intended scope is models with those aliases (currently the tested Gemma4 path); please make that boundary explicit rather than implying automatic support for every split model with tied embeddings.

I have not run a combined Olive assembled-checkpoint -> Mobius ONNX export/runtime validation. This is a follow-up comment, not a formal GitHub changes-requested review.

Keep explicit quantized tying authoritative, but prefer the nested text flag over the parent flag for unquantized models. State that cross-component inference requires unique explicit source aliases and cover Gemma4 and Gemma4-unified conflicts against config parsing.

Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
@xiaoyu-work

Copy link
Copy Markdown
Member Author

@titaiwangms Addressed in fb6b087: for unquantized sources the nested text flag now takes precedence over the parent (matching Gemma4Config parsing); an explicit quantization-level tie still restores sharing. Tests cover conflicting flags for both Gemma4 variants. Inference is scoped to models exposing a unique explicit embedding/head source-path alias pair, now stated in the public inspection contract. The companion microsoft/Olive#2684 includes a real tiny-Gemma4 Olive checkpoint -> Mobius paired ONNX -> ORT CPU full-logit parity test.

@titaiwangms titaiwangms left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed the latest tied-weight inference and explicit alias scope. The earlier tie-precedence concern is resolved, and the PR checks are green. Approved.

xiaoyu-work added a commit to microsoft/Olive that referenced this pull request Sep 26, 2026
## Summary
- preserve one Hugging Face component per build while deferring aliases
only when the canonical table is planned for compatible quantization;
reject incompatible layouts before running builds
- automatically select scoped LM-head, embedding, and vision targets for
KQuant/RTN while retaining explicit opt-outs and whole-model defaults
- fail if required shared-weight assembly is skipped, keep deferred
aliases through follow-up passes, recover missing float heads for
one-sided quantization, and validate shared-weight ownership and
metadata
- recreate packed secondary embedding placeholders on Hugging Face
reload when their sidecars are present

## Validation
- targeted quantizer, assembly, secondary-embedding and cross-repository
suites (187 passing); Ruff and lintrunner PYLINT
- network-free tiny Gemma4: actual Olive two-build assembly -> Mobius
paired embedding/decoder export -> ONNX Runtime CPU inference; full
logits agree with reloaded Hugging Face model (maximum absolute
difference ~4.2e-7)

Depends on onnxruntime/mobius#747 for the shared-weight contract and
generic export materialization.

---------

Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
@xiaoyu-work
xiaoyu-work merged commit 6bd398d into main Sep 26, 2026
25 checks passed
@xiaoyu-work
xiaoyu-work deleted the fix/cross-component-tied-weights branch September 26, 2026 18:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants