Skip to content

Support mixed-width QMoE export - #744

Open
titaiwangms wants to merge 7 commits into
mainfrom
titaiwang/mixed-qmoe-schema
Open

titaiwangms wants to merge 7 commits into
mainfrom
titaiwang/mixed-qmoe-schema

Conversation

@titaiwangms

Copy link
Copy Markdown
Contributor

Summary

  • follow the projection-specific QMoE bit-width contract added by Add mixed-width weight contract for QMoE microsoft/onnxruntime#32697
  • ingest Olive fused expert checkpoints with INT2 gate/up (FC1/FC3) and INT4 down (FC2) banks
  • allocate and validate FC1/FC2 packed weights, scales, and zero points independently
  • preserve legacy uniform INT4 QMoE emission and prevent mixed geometry from being incorrectly rewrite-fused
  • support authoritative Qwen3, Qwen3.5, and Qwen3.5-VL checkpoint prefixes

Contract

Mixed nodes emit the legacy INT4 fallback together with:

  • fc1_expert_weight_bits=2
  • fc2_expert_weight_bits=4
  • fc3_expert_weight_bits=2 for fused SwiGLU

Checkpoint preprocessing fails closed for missing, partial, stale, malformed, or unsupported expert sidecars. Plain and regex overrides follow Olive producer semantics.

Validation

  • python -m pytest src/mobius/models/moe_test.py src/mobius/models/qwen35_test.py src/mobius/components/_moe_test.py src/mobius/rewrite_rules/_qmoe_fusion_test.py src/mobius/_weight_utils_test.py -q --tb=short — 241 passed
  • touched-file formatting and lint passed
  • python -m mypy src/mobius/_weight_utils.py passed
  • git diff --check passed

Runtime boundary

microsoft/onnxruntime#32697 provides the schema and validation contract, but mixed-width execution currently returns NOT_IMPLEMENTED. This PR qualifies graph emission, initializer binding, failure behavior, and package roundtrip; it does not claim runtime numerical support.

Addresses #735.

Follow the ORT QMoE projection-specific bit-width contract for Olive checkpoints with INT2 gate/up experts and INT4 down experts. Preserve legacy uniform INT4 emission and add strict checkpoint, package, and rewrite regression coverage.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: titaiwang <titaiwang@microsoft.com>
@titaiwangms
titaiwangms requested review from a team and a lite review from Copilot September 21, 2026 23:11
@github-actions

github-actions Bot commented Sep 21, 2026

Copy link
Copy Markdown

Performance Comparison

Comparing 783e03b5a0a05d08

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.

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

Mixed and legacy per-layer QMoE plans can diverge between graph construction and weight preprocessing, causing valid layer combinations or dense fallbacks to fail weight binding.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 2 High severity

Open (2)
What changed in this PR

Adds mixed-width Olive QMoE export support, including projection-specific INT2/INT4 layouts, sidecar validation, authoritative Qwen prefixes, and legacy fusion safeguards.

Changes:

  • Adds per-projection QMoE shapes, zero-point handling, and graph attributes.
  • Wires layer-specific quantization resolution through MoE and Qwen3.5 models.
  • Adds validation, roundtrip, and rewrite-fallback tests plus design documentation.
File Description
src/​mobius/​_weight_utils.py Resolves and validates mixed QMoE layouts and sidecars.
src/​mobius/​components/​_moe.py Emits projection-specific QMoE parameters and attributes.
src/​mobius/​models/​moe.py Adds per-layer QMoE source paths and preprocessing metadata.
src/​mobius/​models/​qwen35.py Supports Qwen3.5 text/VL source aliases and module plans.
src/​mobius/​rewrite_rules/​_qmoe_fusion.py Prevents fusion of mixed geometries.
src/​mobius/​_configs/​_base.py Adds QMoE source-path configuration.
src/​mobius/​_configs/​_quantization.py Documents producer-specific override semantics.
src/​mobius/​models/​moe_test.py Tests mixed layouts, validation, binding, and roundtrip.
src/​mobius/​models/​qwen35_test.py Tests Qwen3.5 text/VL mixed export behavior.
src/​mobius/​components/​_moe_test.py Updates QMoE predicate tests.
src/​mobius/​rewrite_rules/​_qmoe_fusion_test.py Tests dense fallback for mixed projections.
docs/​design/​per-component-quantization-loading.md Documents mixed QMoE loading and runtime boundaries.

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread src/mobius/_weight_utils.py Outdated
Comment thread src/mobius/_weight_utils.py
@github-actions

github-actions Bot commented Sep 21, 2026

Copy link
Copy Markdown

🏗️ Architecture Diff

Comparing 783e03b5a0a05d08

Model Sub-model Changes Status
bert (feature-extraction) model 0
falcon model 0
gemma2 model 0
gemma4 (gemma4) decoder 0
gemma4 (gemma4) embedding 0
gemma4 (gemma4) vision_encoder 0
gemma4_text model 0
gpt2 model 0
llama model 0
llama (static-cache) model 0
mamba (ssm-text-generation) model 0
phi3 model 0
phi3 (static-cache) model 0
qwen model 0
qwen (static-cache) model 0
qwen2 model 0
qwen2 (static-cache) model 0
qwen2_moe model 0
qwen2_moe (static-cache) model 0
qwen3 model 0
qwen3 (static-cache) model 0
qwen3_5_moe (hybrid-text-generation) model 0
qwen3_5_text (hybrid-text-generation) model 0
qwen3_5_vl (hybrid-qwen-vl) decoder 0
qwen3_5_vl (hybrid-qwen-vl) embedding 0
qwen3_5_vl (hybrid-qwen-vl) vision_encoder 0
qwen3_moe model 0
qwen3_moe (static-cache) model 0
qwen3_next (hybrid-text-generation) model 0
t5 (seq2seq) decoder 0
t5 (seq2seq) encoder 0
whisper (speech-to-text) decoder 0
whisper (speech-to-text) encoder 0

No architecture changes detected.


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

Resolve every routed layer against its effective Olive layout so mixed and legacy uniform QMoE layers can coexist. Fail closed when packed sidecars target dense fallbacks and honor projection-specific packing geometry.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: titaiwang <titaiwang@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

Non-Olive overrides, legacy-layer validation, and excluded expert fallbacks can produce graph-to-checkpoint layout mismatches.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 4 High severity

Open (4)
Resolved since last review (2)

Comment thread src/mobius/_weight_utils.py Outdated
Comment thread src/mobius/_weight_utils.py Outdated
Comment thread src/mobius/models/moe.py
Comment thread src/mobius/models/qwen35.py
titaiwangms and others added 2 commits September 22, 2026 12:49
Reject routed-expert exclusions and projection-specific GPTQ/AWQ layouts that cannot bind the current graph, while preserving redundant dense-fallback overrides. Validate every routed layer's Olive sidecars in mixed plans.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: titaiwang <titaiwang@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

Fallback-equivalent layers in mixed plans still bypass strict malformed-tensor validation.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 High severity

Open (1)
Resolved since last review (4)

Comment thread src/mobius/_weight_utils.py Outdated
Apply strict fused-expert tensor validation to fallback-equivalent layers in mixed plans while preserving their legacy Olive conversion path. Allow uniform INT4 FC1 and FC2 projections to select zero-point inputs independently, and add regression coverage for malformed sibling layers.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: titaiwang <titaiwang@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

🔵 Needs a closer look

The required multi-layer test does not verify emitted per-layer attributes or final initializer bindings.

Review effort: Balanced
Findings: None

Resolved since last review (1)
Previously missed (1)

In code that hasn't changed since last review

Medium severity Two-layer QMoE test misses per-layer attributes and initializer bindings

src/​mobius/​models/​moe_test.py:526

This two-layer regression validates only the preprocessed dictionary, so the new per-layer qmoe_source_paths graph wiring could assign mixed-width attributes to the wrong layer while this test still passes. The linked acceptance criterion also requires the multi-layer case to inspect node attributes and final initializer bindings. Please build the two-layer package, associate each QMoE node by its FC1 initializer input, assert that only mixed_layer has the FC-specific bit attributes, apply the processed weights, and compare both layers' bound initializer bytes.

Strengthen the two-layer regression to associate each QMoE node with its FC1 initializer, validate projection-specific bit attributes on only the mixed layer, and compare both layers' bound FC1 and FC2 initializer bytes.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: titaiwang <titaiwang@microsoft.com>
@titaiwangms

Copy link
Copy Markdown
Contributor Author

Addressed the previously missed multi-layer coverage in 5fdadb6. The two-layer regression now associates each QMoE node with its layer-specific FC1 initializer through the emitted producer chain, asserts that only the selected mixed layer carries fc1=2/fc2=4/fc3=2 attributes, applies the processed weights, and byte-compares both layers’ bound FC1/FC2 initializers. The focused QMoE suite passes (156 tests).

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

🔵 Needs a closer look

GPTQ/AWQ parent overrides are resolved using incompatible Olive matching semantics.

Review effort: Balanced
Findings: None

Previously missed (1)

In code that hasn't changed since last review

Medium severity Use generic subtree matching for GPTQ and AWQ overrides

src/​mobius/​_weight_utils.py:207

This applies Olive's literal override/substring-exclusion rules to GPTQ and AWQ too. Those formats use QuantizationConfig.for_module subtree matching, so a valid parent override such as model.layers.0.mlp is currently ignored for both expert projections; the graph can retain the global geometry and later mismatch the packed checkpoint instead of failing closed. Delegate non-Olive resolution to the generic matcher before applying Olive-specific semantics.

Delegate GPTQ and AWQ module rules to the generic subtree matcher while preserving Olive's literal override and substring exclusion semantics. Add parent, near-match, and exclusion regressions and document the producer-specific contract.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: titaiwang <titaiwang@microsoft.com>
@titaiwangms

Copy link
Copy Markdown
Contributor Author

Addressed review 5284088183 in a0a05d0. Non-Olive QMoE resolution now delegates to QuantizationConfig.for_module, restoring generic plain-path subtree matching for GPTQ/AWQ while leaving Olive literal overrides and substring exclusions unchanged. Added GPTQ/AWQ regressions for a matching parent override, a substring-only near miss, and a parent exclusion; updated the design contract. Validation: 162 QMoE tests passed, mypy and lint passed.

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

🔵 Needs a closer look

The cross-cutting quantization and weight-binding changes are well tested, but mixed-width runtime execution remains intentionally unavailable.

Review effort: Balanced
Findings: None

This branch has not been deployed

No deployments
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.

2 participants