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.