Support per-layer Olive mixed-precision QMoE export

[titaiwangms]

Motivation

microsoft/Olive#2645 can materialize fused-MoE checkpoints whose routed expert projections use parameter-level precision overrides. For example, the checkpoint default can be INT4 while selected decoder layers use a higher precision.

onnxruntime/mobius#733 adds the generic infrastructure needed to consume such metadata:

• per-component quantization layouts
• exact/regex per-module overrides and exclusions
• projection-specific graph construction
• projection-specific packed-weight normalization and binding

That support covers ordinary quantized linear projections, including different layouts within one component. The remaining gap is the fused com.microsoft::QMoE path: MoELayer still selects one decoder-level quantization config and uses it for every routed-expert layer.

Each QMoE node exposes one expert_weight_bits and one block_size shared by FC1 (gate_up_proj) and FC2 (down_proj). Mobius therefore needs layer-local expert resolution in addition to the generic per-projection support in #733.

Observed failure

A local tiny Qwen3-MoE deployment probe used:

• global/default routed-expert precision: 4-bit, group size 16
• layers 0 and 2 experts.down_proj: 8-bit
• the corresponding experts.gate_up_proj: 4-bit

The Olive checkpoint saved and reloaded correctly. On the pre-#733 export path, Mobius constructed a global 4-bit FC2 initializer with expected shape [2, 32, 8], while the selected layer supplied the 8-bit packed payload [2, 32, 16]:

ValueError: Weight shape mismatch for 'model.layers.0.mlp.fc2_experts_weights':
model expects [2, 32, 8], got [2, 32, 16]

#733 improves generic validation and loading, but it intentionally preserves canonical QMoE groups under one decoder layout. Its QMoE coverage verifies that an ordinary projection override such as self_attn.q_proj does not disturb uniform expert groups; it does not construct QMoE nodes from expert-layer overrides.

Olive ORT GenAI ModelBuilder also currently rejects Olive checkpoints with quantization_config.moe == true, so there is no alternative validated export path for this checkpoint.

Scope after #733

Build on #733 rather than duplicating its config parser, module matcher, codec registry, or generic weight binder.

1. Resolve the effective Olive layouts for each layer's fused experts.gate_up_proj and experts.down_proj source names.
2. Require both projections in one layer to have compatible bits, group_size, symmetry, and supported storage format because one QMoE node represents both.
3. Reject incompatible FC1/FC2 settings before graph construction, naming the layer, source paths, and both effective layouts.
4. Thread the compatible layer-local descriptor into that layer's MoELayer.
5. Use the descriptor consistently for:
   • QMoE initializer shapes
   • zero-point presence and shape
   • expert_weight_bits
   • block_size
   • expert sidecar normalization and binding
6. Preserve the existing uniform QMoE path unchanged. Do not silently dequantize or use the dense all-expert fallback for an incompatible packed checkpoint.

Deployment-unit contract

For the current QMoE ABI:

• canonical Olive identities remain separate: gate_up_proj and down_proj
• the deployable selection unit is the pair of routed-expert projections in one decoder layer
• different layers may use different compatible QMoE layouts because the operator attributes are node-local
• FC1 and FC2 within one layer cannot use different bits or block sizes

The #2645 rule-based modes currently promote down_proj independently. Those checkpoints remain valid PyTorch/Hugging Face artifacts, but they are not QMoE-deployable unless Olive co-promotes the corresponding gate_up_proj or the runtime ABI gains separate FC1/FC2 attributes.

Suggested stages

Stage 1: QMoE layer-plan validation

• Add a helper that resolves paired expert layouts from Load quantized checkpoint weights per component #733's module policy.
• Distinguish uniform, paired-compatible, excluded, and incompatible expert layers.
• Fail closed for partial quantization, mixed FC1/FC2 layouts, unsupported formats, and unsupported QMoE modes.
• Add config/graph tests without making a runtime support claim.

Stage 2: layer-local graph and weight binding

• Construct each MoELayer with its resolved expert layout.
• Emit layer-local QMoE attributes and matching FC1/FC2 parameter shapes.
• Extend the Load quantized checkpoint weights per component #733 typed loading path or a QMoE-specific adapter to validate and bind each layer using the same descriptor.
• Add a two-layer checkpoint test with different qualified layouts and verify exact initializer binding.

Stage 3: runtime qualification

• Verify which per-layer INT4/INT8 combinations are accepted by the targeted ORT QMoE builds and execution providers.
• Run numerical parity against the reloaded Hugging Face/Olive checkpoint.
• Document the exact ORT/CUDA/EP support matrix before declaring each mode supported.

Acceptance criteria

• Load quantized checkpoint weights per component #733's ordinary per-projection and per-component behavior remains unchanged.
• Existing uniform Olive INT4 QMoE checkpoints retain the same graph and weight layout.
• Incompatible same-layer FC1/FC2 overrides fail before producing or binding an invalid model.
• For every supported layer-local configuration, packed checkpoint shapes exactly match the generated QMoE initializers.
• A test covers at least two QMoE layers with different qualified layouts and inspects node attributes plus final weight bindings.
• Runtime-qualified tests demonstrate inference parity for every advertised per-layer mode.

Olive dependency

Until this contract is implemented and qualified, score-based fused-MoE SelectiveMixedPrecision can either:

• co-promote gate_up_proj and down_proj as one per-layer deployment unit and target this future path; or
• allow independent projection selection but document the result as PyTorch-checkpoint-only.

September 21, 2026 update: mixed-width QMoE target

microsoft/onnxruntime#32657 now prioritizes a mixed-width fused-QMoE path for
Qwen-class models:

• FC1 routed expert gate/up weights: INT2
• FC2 routed expert down weights: INT4

This changes the forward-looking contract assumed above. The existing
single-width QMoE ABI still requires compatible FC1/FC2 layouts, but the proposed
ORT contract will represent independent effective widths for FC1 and FC2 while
preserving legacy expert_weight_bits behavior.

Once the ORT contract is pinned, this issue should track the Mobius exporter
workstream from the upstream delivery plan:

1. Resolve each layer's experts.gate_up_proj and experts.down_proj layouts
   independently instead of collapsing them into one decoder-level quantization
   config.
2. Carry separate FC1 and FC2 descriptors through parameter allocation, packed
   sidecar normalization, initializer binding, and QMoE emission.
3. Construct FC1 packed shapes from the INT2 layout and FC2 packed shapes from
   the INT4 layout, including independent scale and zero-point validation.
4. Emit the finalized ORT mixed-width attributes or descriptor only after its
   schema and compatibility semantics are agreed.
5. Preserve existing uniform INT4 QMoE graphs and checkpoints unchanged.
6. Qualify a tiny Olive-produced mixed-width checkpoint against the ORT CPU
   mixed-QMoE reference, including exact packed bytes, gate/up interleaving,
   down-projection orientation, save/reload, external-data binding, and
   numerical parity.

Contract-dependent graph emission is blocked on the corresponding ORT QMoE
schema change. Schema-neutral checkpoint inspection, packing oracles, and test
fixture preparation can proceed beforehand.

Related tracking:

• microsoft/onnxruntime#32657: QMoE contract, CPU reference, and CUDA runtime
• microsoft/Olive#2638: fused-expert SMP recipe and checkpoint qualification