Fix Olive-quantized Qwen3 MoE export

CHANGELOG.md

@@ -7,6 +7,59 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Packed fused MoE experts (Olive/GPTQ/AWQ) survive HF weight renaming
+
+#### Fixed
+
+- MoE exports whose routed experts go through the fused `com.microsoft::QMoE`
+  packer (`qwen3_moe`, Mixtral, OLMoE, Qwen2-MoE, Ernie4.5-MoE, GLM4-MoE) no
+  longer fail on **packed quantized** fused expert tensors.
+  `_rename_moe_expert_weights` matched packed sidecars by substring
+  (`.experts.gate_up_proj` also matches `experts.gate_up_proj_qweight`) and
+  split them as if they were float weights, writing the `qweight`, `scales` and
+  `qzeros` of one projection to the *same* per-expert `.weight` key — so only
+  the last one survived, and the restacked tensor no longer matched the QMoE
+  parameter, aborting the export at weight binding:
+
+  ```
+  ValueError: Weight shape mismatch for 'model.layers.0.mlp.fc1_experts_weights':
+  model expects [4, 64, 32], got [4, 256]
+  ```
+
+  Packed tensors now pass through untouched and reach
+  `pack_qmoe_expert_weights` in the expert-major layout it expects (Qwen3-MoE
+  Olive int4: `[128, 1536, 1024]` uint8 weights + `[128, 1536, 16]` bf16
+  scales). Unquantized fused experts still un-fuse into the dense per-expert
+  fallback.
+
+#### Added
+
+- `mobius._weight_utils.is_packed_quant_key` plus the shared
+  `OLIVE_PACKED_QUANT_SUFFIXES` / `DOTTED_PACKED_QUANT_SUFFIXES` /
+  `PACKED_QUANT_SUFFIXES` constants: one predicate for
+  `_qweight`/`_scales`/`_qzeros` (Olive) and `.qweight`/`.scales`/`.qzeros`
+  (GPTQ/AWQ) sidecar keys, reused by `preprocess_quantized_weights`.
+
+---
+
+### Qwen3-MoE packages load in ONNX Runtime GenAI
+
+#### Fixed
+
+- Exported `qwen3_moe` packages no longer fail to load with
+  `RuntimeError: Unsupported model_type in config.json: qwen3_moe`.
+  `_ORT_GENAI_MODEL_TYPE` had no `qwen3_moe` entry, so `--config` mode wrote
+  the HuggingFace type straight into `genai_config.json` and ORT GenAI
+  rejected it (its LLM type registry has no `qwen3_moe`). Qwen3-MoE now
+  resolves to the accepted `qwen3` type: both `qwen2` and `qwen3` dispatch to
+  ORT GenAI's `DecoderOnly_Model`, but its tokenizer tag fallback
+  (`tokenizer_tag_utils.cpp`) only supplies the Qwen3 reasoning-token IDs
+  (`bor` 151667 / `eor` 151668) for `qwen3` — under `qwen2` they are absent and
+  `tokenizer.bor_token_id` / `eor_token_id` throws. The pre-existing dense
+  `qwen3 -> qwen2` alias is unchanged.
+
+---
+
 ### Fixed
 
 - An exported graph is no longer transcribed into its own workflow component.

src/mobius/_weight_utils.py

@@ -24,6 +24,42 @@
 
 logger = logging.getLogger(__name__)
 
+# Key suffixes marking a *packed quantized* sidecar tensor rather than a float
+# weight. A "sidecar" here is an auxiliary tensor that rides alongside a
+# quantized parameter's canonical ``.weight`` key instead of replacing it —
+# the packed integer payload (``qweight``) plus its per-group ``scales`` and
+# ``qzeros`` are each stored under their own key, so one logical weight is
+# split across several sidecar keys until they are unpacked back into a
+# single dequantized tensor. Olive appends them with an underscore directly
+# to the parameter name (``<pname>_qweight``, see Olive's
+# ``olive/common/quant/state_dict.py``), while GPTQ/AWQ store dotted sibling
+# buffers on the owning module (``<module>.qweight``). Both conventions occur
+# in raw HF checkpoints.
+OLIVE_PACKED_QUANT_SUFFIXES: frozenset[str] = frozenset({"_qweight", "_scales", "_qzeros"})
+DOTTED_PACKED_QUANT_SUFFIXES: frozenset[str] = frozenset({".qweight", ".scales", ".qzeros"})
+# str.endswith() requires a tuple (not a set), so the combined predicate below
+# needs a tuple even though the two suffix sets above are otherwise unordered.
+PACKED_QUANT_SUFFIXES: tuple[str, ...] = tuple(
+    OLIVE_PACKED_QUANT_SUFFIXES | DOTTED_PACKED_QUANT_SUFFIXES
+)
+
+
+def is_packed_quant_key(name: str) -> bool:
+    """Whether ``name`` is a packed-quantization sidecar key.
+
+    Packed keys carry the quantized payload (``qweight``), the per-group
+    ``scales`` or the ``qzeros`` of a quantized parameter, in either the Olive
+    underscore convention (``…experts.gate_up_proj_qweight``) or the GPTQ/AWQ
+    dotted convention (``…gate_proj.qweight``).
+
+    HF→ONNX key rewriting must leave these tensors intact until
+    :func:`preprocess_quantized_weights` unpacks them: they are packed bytes,
+    not float weights, so reshaping or splitting them corrupts the payload —
+    and several sidecars of one projection would collapse onto a single
+    renamed ``.weight`` key.
+    """
+    return name.endswith(PACKED_QUANT_SUFFIXES)
+
 
 def supported_qmoe_quantization(
     quantization: QuantizationConfig | None,
@@ -1077,14 +1113,6 @@ def preprocess_quantized_weights(
         if qmoe_target_path is not None
         else False
     )
-    packed_suffixes = (
-        "_qweight",
-        "_scales",
-        "_qzeros",
-        ".qweight",
-        ".scales",
-        ".qzeros",
-    )
     if (
         use_qmoe
         and quantization is not None
@@ -1107,7 +1135,7 @@ def preprocess_quantized_weights(
         packed_expert_keys = [
             key
             for key in state_dict
-            if qmoe_target_path in key and ".experts." in key and key.endswith(packed_suffixes)
+            if qmoe_target_path in key and ".experts." in key and is_packed_quant_key(key)
         ]
         if packed_expert_keys:
             raise ValueError(
@@ -1121,18 +1149,24 @@ def preprocess_quantized_weights(
                 "supported_qmoe_quantization) for MoE models instead."
             )
 
-    def _is_packed_key(key: str, float_key: str) -> bool:
+    def _is_packed_sidecar_of(key: str, float_key: str) -> bool:
+        """Whether ``key`` is a packed sidecar of the specific ``float_key``.
+
+        Unlike the module-level :func:`is_packed_quant_key` suffix predicate,
+        this matches the *exact* sidecar keys of one named float parameter.
+        """
         owner = float_key.removesuffix(".weight")
-        return any(key == float_key + suffix for suffix in packed_suffixes[:3]) or any(
-            key == owner + suffix for suffix in packed_suffixes[3:]
+        return any(key == float_key + suffix for suffix in OLIVE_PACKED_QUANT_SUFFIXES) or any(
+            key == owner + suffix for suffix in DOTTED_PACKED_QUANT_SUFFIXES
         )
 
     packed_embedding_or_head = (
         next(
             (
                 key
                 for key in state_dict
-                if _is_packed_key(key, embed_key) or _is_packed_key(key, head_key)
+                if _is_packed_sidecar_of(key, embed_key)
+                or _is_packed_sidecar_of(key, head_key)
             ),
             None,
         )

src/mobius/_weight_utils_test.py

@@ -10,6 +10,7 @@
 
 from mobius._configs import QuantizationConfig
 from mobius._weight_utils import (
+    is_packed_quant_key,
     merge_lora_weights,
     preprocess_awq_weights,
     preprocess_gptq_weights,
@@ -1393,3 +1394,43 @@ def test_non_expert_keys_pass_through(self):
         assert torch.equal(
             result["model.embed_tokens.weight"], sd["model.embed_tokens.weight"]
         )
+
+
+class TestIsPackedQuantKey:
+    """Shared predicate for packed-quantization sidecar keys."""
+
+    @pytest.mark.parametrize(
+        "key",
+        [
+            # Olive underscore convention (suffix on the parameter name).
+            "model.layers.0.mlp.experts.gate_up_proj_qweight",
+            "model.layers.0.mlp.experts.gate_up_proj_scales",
+            "model.layers.0.mlp.experts.down_proj_qzeros",
+            "model.layers.0.self_attn.q_proj.weight_qweight",
+            # GPTQ/AWQ dotted convention (sibling buffers of the module).
+            "model.layers.0.mlp.experts.0.gate_proj.qweight",
+            "model.layers.0.mlp.experts.0.gate_proj.scales",
+            "model.layers.0.mlp.experts.0.gate_proj.qzeros",
+        ],
+    )
+    def test_packed_keys_detected(self, key):
+        assert is_packed_quant_key(key)
+
+    @pytest.mark.parametrize(
+        "key",
+        [
+            # Float weights, including the fused MoE tensors that *are* split.
+            "model.layers.0.mlp.experts.gate_up_proj",
+            "model.layers.0.mlp.experts.down_proj",
+            "model.layers.0.mlp.gate.weight",
+            "model.layers.0.self_attn.q_proj.bias",
+            "model.layers.0.block_sparse_moe.input_linear.weight",
+            # Already-unpacked names produced downstream by the preprocessors.
+            "model.layers.0.mlp.fc1_experts_weights",
+            "model.layers.0.self_attn.q_proj.zero_points",
+            # Suffix must be terminal, not merely present.
+            "model.layers.0.mlp.experts.gate_up_proj_qweight.extra",
+        ],
+    )
+    def test_unpacked_keys_rejected(self, key):
+        assert not is_packed_quant_key(key)

src/mobius/integrations/ort_genai/auto_export.py

@@ -75,6 +75,17 @@ def _revision_kwargs(revision: str | None) -> dict[str, str]:
     "llama": "llama",
     "qwen2": "qwen2",
     "qwen3": "qwen2",
+    # Qwen3-MoE shares the dense Qwen3 decoder contract (same inputs, position
+    # IDs and KV cache layout); only the MLP differs, and that is fused into
+    # the exported graph. ORT GenAI has no "qwen3_moe" entry in its LLM type
+    # registry (see onnxruntime-genai/src/models/model_type.h), so passing the
+    # HF type through fails to load with "Unsupported model_type in
+    # config.json: qwen3_moe". It maps to "qwen3" rather than reusing the dense
+    # "qwen3" -> "qwen2" alias: both types dispatch to DecoderOnly_Model, but
+    # ORT GenAI's tokenizer tag fallback (tokenizer_tag_utils.cpp) only supplies
+    # Qwen3 reasoning-token IDs (bor 151667 / eor 151668) for "qwen3"; under
+    # "qwen2" those are absent and tokenizer.bor_token_id/eor_token_id throws.
+    "qwen3_moe": "qwen3",
     "phi3": "phi3",
     "phi": "phi",
     "phi4mm": "phi4mm",

src/mobius/integrations/ort_genai/auto_export_test.py

@@ -123,6 +123,14 @@ def test_phi4mm_model_types(self):
         assert _resolve_ort_genai_model_type("phi4_multimodal") == "phi4mm"
         assert _resolve_ort_genai_model_type("phi") == "phi"
 
+    def test_qwen3_moe_maps_to_supported_decoder_type(self):
+        # ORT GenAI's LLM type registry has no "qwen3_moe" entry, so passing
+        # the HF type through fails to load with "Unsupported model_type in
+        # config.json: qwen3_moe". Qwen3-MoE maps to the accepted "qwen3" type
+        # (not the dense "qwen3" -> "qwen2" alias) so ORT GenAI's tokenizer tag
+        # fallback still supplies the Qwen3 reasoning-token IDs.
+        assert _resolve_ort_genai_model_type("qwen3_moe") == "qwen3"
+
     def test_gemma4_unified_model_types(self):
         # The gemma-4-12B unified checkpoint (model_type "gemma4_unified")
         # reuses the multimodal "gemma4" ORT GenAI pipeline; its standalone
@@ -161,6 +169,14 @@ def test_multimodal_keeps_hf_type(self):
     def test_decoder_only_falls_back_to_hf_when_config_missing(self):
         assert _select_ort_model_type(None, "qwen3", is_decoder_only=True) == "qwen2"
 
+    def test_decoder_only_qwen3_moe_resolves_to_supported_type(self):
+        # hf_model_id mode: both the package config and the HF config report
+        # "qwen3_moe"; the decoder-only preference must still resolve through
+        # the alias rather than emitting the unsupported HF type.
+        assert (
+            _select_ort_model_type("qwen3_moe", "qwen3_moe", is_decoder_only=True) == "qwen3"
+        )
+
     def test_decoder_only_unknown_config_falls_back_to_hf(self):
         # An unrecognised config.model_type (not in _ORT_GENAI_MODEL_TYPE) must
         # not pass straight through as an invalid ORT type; fall back to the
@@ -2085,6 +2101,25 @@ class FakeConfig:
         # "gemma2" maps to "gemma" in _ORT_GENAI_MODEL_TYPE
         assert data["model"]["type"] == "gemma"
 
+    def test_config_mode_qwen3_moe_emits_supported_decoder_type(self, tmp_path):
+        """Qwen3-MoE --config exports must not emit the unsupported HF type.
+
+        Regression: a decoder-only Qwen3-MoE package wrote
+        ``"type": "qwen3_moe"`` into genai_config.json, and loading it raised
+        ``RuntimeError: Unsupported model_type in config.json: qwen3_moe``
+        because ORT GenAI's LLM registry has no such type. The emitted type is
+        literally ``"qwen3"``: ORT GenAI's tokenizer tag fallback keys the
+        Qwen3 reasoning-token IDs off that name.
+        """
+        from mobius.integrations.ort_genai.auto_export import write_ort_genai_config
+
+        pkg = _make_fake_llm_pkg("qwen3_moe")
+        result = write_ort_genai_config(pkg, str(tmp_path), hf_model_id=None)
+
+        with open(result["genai_config"]) as f:
+            data = json.load(f)
+        assert data["model"]["type"] == "qwen3"
+
     def test_config_mode_gemma3_text_vlm_uses_multimodal_model_type(self, tmp_path):
         """Gemma3 VLM --config exports use ORT's multimodal gemma3 type."""
         import dataclasses

src/mobius/models/moe.py

@@ -12,7 +12,7 @@
 from onnxscript import OpBuilder, nn
 
 from mobius._configs import ArchitectureConfig
-from mobius._weight_utils import preprocess_quantized_weights
+from mobius._weight_utils import is_packed_quant_key, preprocess_quantized_weights
 from mobius.components import (
     Attention,
     Embedding,
@@ -611,6 +611,22 @@ def _rename_moe_expert_weights(
        ``input_linear.weight [N, 2*inter, hidden]`` → per-expert gate_proj + up_proj
        ``output_linear.weight [N, hidden, inter]`` → per-expert down_proj
        ``router.layer.weight`` → ``gate.weight``
+
+    Packed quantized expert tensors (Olive ``…_qweight``/``_scales``/``_qzeros``,
+    GPTQ/AWQ ``….qweight``/``.scales``/``.qzeros`` — see
+    :func:`~mobius._weight_utils.is_packed_quant_key`) are **never** split.
+    Splitting them would reinterpret packed bytes as float rows *and* collapse
+    every sidecar of one projection onto the same per-expert ``.weight`` key,
+    keeping only the last one. Call sites that continue into
+    :func:`_preprocess_moe_weights` (the :class:`MoECausalLMModel` family) need
+    the fused expert-major layout intact for ``pack_qmoe_expert_weights``; for
+    every other call site this function simply leaves the packed keys alone.
+
+    Module-path renames still apply to packed keys under the Olive convention,
+    where the suffix hangs off a retained ``.weight`` component
+    (``…mlp.router.weight_qweight`` → ``…mlp.gate.weight_qweight``). Dotted
+    GPTQ/AWQ sidecars have no ``.weight`` component (``…mlp.router.qweight``),
+    so no rename pattern matches them and they pass through as-is.
     """
     # Step 0: GraniteMoE uses block_sparse_moe; rename to mlp to match our attribute.
     state_dict = {
@@ -632,6 +648,10 @@ def _rename_moe_expert_weights(
     # Second pass: split fused 3D expert weights and rename routers
     fused: dict[str, torch.Tensor] = {}
     for name, tensor in list(renamed.items()):
+        # Packed quantized sidecars are never split: they must stay fused for the
+        # downstream QMoE packer, and splitting also collides their keys. Only
+        # float fused tensors are split into per-expert ``.weight`` keys.
+        is_packed = is_packed_quant_key(name)
         # GraniteMoE: router.layer.weight → gate.weight
         if ".router.layer.weight" in name:
             new_name = name.replace(".router.layer.weight", ".gate.weight")
@@ -643,7 +663,7 @@ def _rename_moe_expert_weights(
             fused[new_name] = tensor
             del renamed[name]
         # GraniteMoE: input_linear.weight [N, 2*inter, hidden] → gate_proj + up_proj
-        elif ".input_linear.weight" in name and tensor.dim() == 3:
+        elif not is_packed and ".input_linear.weight" in name and tensor.dim() == 3:
             prefix = name.replace(".input_linear.weight", "")
             num_experts = tensor.shape[0]
             half = tensor.shape[1] // 2
@@ -654,15 +674,15 @@ def _rename_moe_expert_weights(
                 fused[f"{prefix}.experts.{i}.up_proj.weight"] = up_w
             del renamed[name]
         # GraniteMoE: output_linear.weight [N, hidden, inter] → down_proj
-        elif ".output_linear.weight" in name and tensor.dim() == 3:
+        elif not is_packed and ".output_linear.weight" in name and tensor.dim() == 3:
             prefix = name.replace(".output_linear.weight", "")
             num_experts = tensor.shape[0]
             for i in range(num_experts):
                 fused[f"{prefix}.experts.{i}.down_proj.weight"] = tensor[i]
             del renamed[name]
         # Fused gate_up_proj [N, 2*inter, hidden] → per-expert gate_proj + up_proj
         # (Mixtral, OLMoE, Qwen2-MoE, PhiMoE)
-        elif ".experts.gate_up_proj" in name and tensor.dim() == 3:
+        elif not is_packed and ".experts.gate_up_proj" in name and tensor.dim() == 3:
             prefix = name.split(".experts.gate_up_proj")[0]
             num_experts = tensor.shape[0]
             half = tensor.shape[1] // 2
@@ -674,7 +694,12 @@ def _rename_moe_expert_weights(
             del renamed[name]
         # Fused experts.down_proj [N, hidden, inter] → per-expert down_proj
         # Only match the fused format (3D tensor), not per-expert experts.{i}.down_proj
-        elif ".experts.down_proj" in name and tensor.dim() == 3 and "experts." in name:
+        elif (
+            not is_packed
+            and ".experts.down_proj" in name
+            and tensor.dim() == 3
+            and "experts." in name
+        ):
             parts = name.split(".experts.down_proj")
             if len(parts) == 2 and not parts[0].endswith(tuple("0123456789")):
                 prefix = parts[0]