diff --git a/CHANGELOG.md b/CHANGELOG.md index f89ee93ff..19d4de273 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/src/mobius/_weight_utils.py b/src/mobius/_weight_utils.py index 1520ddc1d..b01148480 100644 --- a/src/mobius/_weight_utils.py +++ b/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 (``_qweight``, see Olive's +# ``olive/common/quant/state_dict.py``), while GPTQ/AWQ store dotted sibling +# buffers on the owning 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,10 +1149,15 @@ 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 = ( @@ -1132,7 +1165,8 @@ def _is_packed_key(key: str, float_key: str) -> bool: ( 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, ) diff --git a/src/mobius/_weight_utils_test.py b/src/mobius/_weight_utils_test.py index 916e1489d..ca81cc927 100644 --- a/src/mobius/_weight_utils_test.py +++ b/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) diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 22bea0e0a..926229fb8 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/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", diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 5bdf1ecd7..0e3263663 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/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 diff --git a/src/mobius/models/moe.py b/src/mobius/models/moe.py index 870d00175..b05f4a590 100644 --- a/src/mobius/models/moe.py +++ b/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,7 +674,7 @@ 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): @@ -662,7 +682,7 @@ def _rename_moe_expert_weights( 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] diff --git a/src/mobius/models/moe_test.py b/src/mobius/models/moe_test.py new file mode 100644 index 000000000..3674c7452 --- /dev/null +++ b/src/mobius/models/moe_test.py @@ -0,0 +1,359 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Weight-preprocessing tests for the generic MoE family (``qwen3_moe`` & co.). + +Focus: packed *fused* expert tensors must survive +:func:`~mobius.models.moe._rename_moe_expert_weights` untouched so +``pack_qmoe_expert_weights`` can map them onto ``com.microsoft::QMoE`` +``fc1``/``fc2`` parameters. + +An Olive-quantized Qwen3-MoE checkpoint stores each layer's routed experts as +expert-major *packed* sidecars — ``experts.gate_up_proj_qweight`` +``[E, 2*moe_inter, hidden*bits/8]`` uint8 next to ``experts.gate_up_proj_scales`` +``[E, 2*moe_inter, hidden/group]`` bf16 (and ``_qzeros`` when asymmetric). +Every one of those keys contains ``.experts.gate_up_proj`` and is 3-D, so the +fused-expert *splitter* used to claim them: it reinterpreted packed bytes as +float rows **and** wrote every sidecar of a projection to the same +``experts.{i}.gate_proj.weight`` key, so only the last one survived and the QMoE +parameters were filled with garbage (or nothing at all). + +All configs are tiny and synthetic — no checkpoint download. +""" + +from __future__ import annotations + +import math + +import onnx_ir as ir +import pytest +import torch + +from mobius._configs import QuantizationConfig +from mobius._testing import make_config +from mobius.models.moe import MoECausalLMModel, _rename_moe_expert_weights + +# Tiny mirror of the real Olive Qwen3-MoE ABI (hidden=2048, moe_inter=768, +# bits=4, group=128): every shape below keeps the same relationships. +_E, _H, _INT, _BLK, _BITS = 4, 64, 32, 16, 4 +_FC1_OUT = 2 * _INT # gate rows then up rows +# Packed/blocked column counts are derived from the *input* dim of a projection: +# gate_up_proj and the attention projections consume ``hidden``, down_proj +# consumes ``moe_intermediate``. +_HIDDEN_PACKED = _H * _BITS // 8 +_INTER_PACKED = _INT * _BITS // 8 +_HIDDEN_BLOCKS = _H // _BLK +_INTER_BLOCKS = _INT // _BLK +_LAYER = "model.layers.0." + + +def _quantization(*, sym: bool = True, **overrides) -> QuantizationConfig: + return QuantizationConfig( + bits=_BITS, group_size=_BLK, quant_method="olive", sym=sym, **overrides + ) + + +def _moe_config(quantization: QuantizationConfig | None) -> object: + return make_config( + num_hidden_layers=1, + hidden_size=_H, + intermediate_size=64, + moe_intermediate_size=_INT, + num_local_experts=_E, + num_experts_per_tok=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + vocab_size=32, + quantization=quantization, + ) + + +def _packed_expert_state_dict(*, sym: bool = True) -> dict[str, torch.Tensor]: + """Fused expert-major Olive sidecars, as found in the Qwen3-MoE checkpoint. + + Olive suffixes the *parameter* name with an underscore + (``gate_up_proj_qweight``), not a dotted sibling buffer — see Olive's + ``olive/common/quant/state_dict.py``. Scales are bf16 in the real + checkpoint. + """ + p = f"{_LAYER}mlp." + state_dict = { + p + "experts.gate_up_proj_qweight": torch.randint( + 0, 256, (_E, _FC1_OUT, _HIDDEN_PACKED), dtype=torch.uint8 + ), + p + "experts.gate_up_proj_scales": torch.rand( + _E, _FC1_OUT, _HIDDEN_BLOCKS, dtype=torch.bfloat16 + ), + p + "experts.down_proj_qweight": torch.randint( + 0, 256, (_E, _H, _INTER_PACKED), dtype=torch.uint8 + ), + p + "experts.down_proj_scales": torch.rand( + _E, _H, _INTER_BLOCKS, dtype=torch.bfloat16 + ), + p + "gate.weight": torch.rand(_E, _H), + } + if not sym: + state_dict[p + "experts.gate_up_proj_qzeros"] = torch.randint( + 0, 256, (_E, _FC1_OUT, math.ceil(_HIDDEN_BLOCKS * _BITS / 8)), dtype=torch.uint8 + ) + state_dict[p + "experts.down_proj_qzeros"] = torch.randint( + 0, 256, (_E, _H, math.ceil(_INTER_BLOCKS * _BITS / 8)), dtype=torch.uint8 + ) + return state_dict + + +def _packed_decoder_state_dict(*, sym: bool = True) -> dict[str, torch.Tensor]: + """Full single-layer checkpoint: packed attention + packed routed experts. + + Attention projections are packed per-linear (``.weight_qweight`` + ``[N, hidden*bits/8]``); norms, router, embedding and LM head stay float, + matching what Olive emits for Qwen3-MoE. + """ + state_dict = _packed_expert_state_dict(sym=sym) + kv_out = 2 * 16 # num_key_value_heads * head_dim + for proj, out_features in ( + ("q_proj", 4 * 16), + ("k_proj", kv_out), + ("v_proj", kv_out), + ("o_proj", _H), + ): + base = f"{_LAYER}self_attn.{proj}." + state_dict[base + "weight_qweight"] = torch.randint( + 0, 256, (out_features, _HIDDEN_PACKED), dtype=torch.uint8 + ) + state_dict[base + "weight_scales"] = torch.rand( + out_features, _HIDDEN_BLOCKS, dtype=torch.bfloat16 + ) + state_dict[f"{_LAYER}input_layernorm.weight"] = torch.rand(_H) + state_dict[f"{_LAYER}post_attention_layernorm.weight"] = torch.rand(_H) + state_dict["model.norm.weight"] = torch.rand(_H) + state_dict["model.embed_tokens.weight"] = torch.rand(32, _H) + state_dict["lm_head.weight"] = torch.rand(32, _H) + return state_dict + + +class TestRenameMoEExpertWeightsPacked: + """Packed fused sidecars must pass through the HF→ONNX renamer untouched.""" + + def test_packed_olive_sidecars_stay_distinct_and_unsplit(self): + """``_qweight`` and ``_scales`` keep their own keys, shapes and dtypes. + + Regression guard: the splitter matched on the ``.experts.gate_up_proj`` + substring, so both sidecars were split into per-expert + ``experts.{i}.gate_proj.weight`` keys — the same key for both — and the + packed payload was lost. + """ + state_dict = _packed_expert_state_dict() + original = {k: (tuple(v.shape), v.dtype) for k, v in state_dict.items()} + + out = _rename_moe_expert_weights(dict(state_dict)) + + assert set(out) == set(state_dict), "packed keys must pass through unchanged" + for key, tensor in out.items(): + assert (tuple(tensor.shape), tensor.dtype) == original[key] + assert tensor is state_dict[key] + # No per-expert float-style keys were fabricated from packed tensors. + assert not any(f".experts.{i}." in k for k in out for i in range(_E)) + + @pytest.mark.parametrize("suffix", ["_qweight", "_scales", "_qzeros"]) + def test_packed_granite_fused_linears_pass_through(self, suffix): + """GraniteMoE's fused ``input_linear``/``output_linear`` are guarded too. + + Their branches keyed off the ``.input_linear.weight`` / + ``.output_linear.weight`` substring, which a packed + ``...input_linear.weight_qweight`` key also contains. + + Scope: this only pins the *generic renamer's* preservation of packed + keys. ``GraniteMoECausalLMModel`` preprocesses through the base + ``CausalLMModel`` (``qmoe_target_path=None``), so it never reaches + ``pack_qmoe_expert_weights`` — quantized GraniteMoE export stays + unwired, and this test does not claim otherwise. + """ + state_dict = { + f"{_LAYER}block_sparse_moe.input_linear.weight{suffix}": torch.zeros( + _E, _FC1_OUT, _HIDDEN_PACKED, dtype=torch.uint8 + ), + f"{_LAYER}block_sparse_moe.output_linear.weight{suffix}": torch.zeros( + _E, _H, _INTER_PACKED, dtype=torch.uint8 + ), + } + + out = _rename_moe_expert_weights(state_dict) + + # ``block_sparse_moe`` → ``mlp`` still applies (path rename only). + assert set(out) == { + f"{_LAYER}mlp.input_linear.weight{suffix}", + f"{_LAYER}mlp.output_linear.weight{suffix}", + } + + @pytest.mark.parametrize("suffix", [".qweight", ".scales", ".qzeros"]) + def test_packed_dotted_fused_experts_pass_through(self, suffix): + """GPTQ/AWQ dotted sidecars on a fused expert parameter are guarded too.""" + state_dict = { + f"{_LAYER}mlp.experts.gate_up_proj{suffix}": torch.zeros( + _E, _FC1_OUT, _HIDDEN_PACKED, dtype=torch.uint8 + ), + f"{_LAYER}mlp.experts.down_proj{suffix}": torch.zeros( + _E, _H, _INTER_PACKED, dtype=torch.uint8 + ), + } + + out = _rename_moe_expert_weights(dict(state_dict)) + + assert set(out) == set(state_dict) + + def test_unquantized_fused_experts_still_split(self): + """Float fused experts keep the dense per-expert unfusing behaviour.""" + p = f"{_LAYER}mlp." + state_dict = { + p + "experts.gate_up_proj": torch.rand(_E, _FC1_OUT, _H), + p + "experts.down_proj": torch.rand(_E, _H, _INT), + p + "gate.weight": torch.rand(_E, _H), + } + + out = _rename_moe_expert_weights(state_dict) + + assert p + "experts.gate_up_proj" not in out + assert p + "experts.down_proj" not in out + for i in range(_E): + assert out[f"{p}experts.{i}.gate_proj.weight"].shape == (_INT, _H) + assert out[f"{p}experts.{i}.up_proj.weight"].shape == (_INT, _H) + assert out[f"{p}experts.{i}.down_proj.weight"].shape == (_H, _INT) + assert out[p + "gate.weight"].shape == (_E, _H) + + def test_unquantized_granite_fused_linears_still_split(self): + """Float GraniteMoE fused linears keep splitting (and the router rename).""" + p = f"{_LAYER}block_sparse_moe." + state_dict = { + p + "input_linear.weight": torch.rand(_E, _FC1_OUT, _H), + p + "output_linear.weight": torch.rand(_E, _H, _INT), + p + "router.layer.weight": torch.rand(_E, _H), + } + + out = _rename_moe_expert_weights(state_dict) + + assert f"{_LAYER}mlp.gate.weight" in out + for i in range(_E): + assert out[f"{_LAYER}mlp.experts.{i}.gate_proj.weight"].shape == (_INT, _H) + assert out[f"{_LAYER}mlp.experts.{i}.down_proj.weight"].shape == (_H, _INT) + + def test_router_rename_still_applies_to_packed_keys(self): + """Module-path renames are suffix-preserving, so packed routers rename too.""" + out = _rename_moe_expert_weights( + { + f"{_LAYER}mlp.router.weight_qweight": torch.zeros( + _E, _HIDDEN_PACKED, dtype=torch.uint8 + ) + } + ) + + assert set(out) == {f"{_LAYER}mlp.gate.weight_qweight"} + + +class TestMoECausalLMPackedQMoEExport: + """End-to-end ``preprocess_weights`` for an Olive-packed Qwen3-MoE layer.""" + + def test_packed_experts_become_qmoe_parameters(self): + model = MoECausalLMModel(_moe_config(_quantization())) + + out = model.preprocess_weights(_packed_expert_state_dict()) + + p = f"{_LAYER}mlp." + fc1, fc2 = out[p + "fc1_experts_weights"], out[p + "fc2_experts_weights"] + assert fc1.shape == (_E, _FC1_OUT, _HIDDEN_PACKED) + assert fc1.dtype == torch.uint8 + assert fc2.shape == (_E, _H, _INTER_PACKED) + assert fc2.dtype == torch.uint8 + # Scales stay a separate tensor (they used to collide with qweight). + assert out[p + "fc1_scales"].shape == (_E, _FC1_OUT, _HIDDEN_BLOCKS) + assert out[p + "fc2_scales"].shape == (_E, _H, _INTER_BLOCKS) + assert out[p + "fc1_scales"].dtype == torch.bfloat16 + # No dense per-expert fallback keys leaked through. + assert not any(".mlp.experts." in k for k in out) + + def test_symmetric_quantization_emits_no_zero_points(self): + model = MoECausalLMModel(_moe_config(_quantization(sym=True))) + + out = model.preprocess_weights(_packed_expert_state_dict()) + + assert model.model.layers[0].mlp.fc1_experts_zero_points is None + assert model.model.layers[0].mlp.fc2_experts_zero_points is None + assert not any("zero_point" in k for k in out) + + def test_asymmetric_quantization_packs_zero_points(self): + """The ``_qzeros`` sidecar survives the renamer and lands on QMoE too.""" + model = MoECausalLMModel(_moe_config(_quantization(sym=False))) + + out = model.preprocess_weights(_packed_expert_state_dict(sym=False)) + + p = f"{_LAYER}mlp." + params = dict(model.named_parameters()) + for key in (p + "fc1_experts_zero_points", p + "fc2_experts_zero_points"): + assert out[key].dtype == torch.uint8 + assert tuple(params[key].shape) == tuple(out[key].shape) + + def test_every_produced_key_binds_to_a_named_parameter(self): + """Whole-layer round trip: names *and* shapes match the built module.""" + model = MoECausalLMModel(_moe_config(_quantization())) + params = dict(model.named_parameters()) + + out = model.preprocess_weights(_packed_decoder_state_dict()) + + assert out, "preprocess_weights returned an empty state dict" + for name, tensor in out.items(): + assert name in params, f"{name} does not bind to any model parameter" + assert tuple(params[name].shape) == tuple(tensor.shape), ( + f"shape mismatch for {name}: model expects " + f"{tuple(params[name].shape)}, got {tuple(tensor.shape)}" + ) + # The fused QMoE parameters are actually populated, not skipped. + assert f"{_LAYER}mlp.fc1_experts_weights" in out + assert f"{_LAYER}mlp.fc2_experts_weights" in out + + def test_attention_is_quantized_while_router_and_tables_stay_float(self): + """Only the checkpoint's packed modules become ``QuantizedLinear``.""" + model = MoECausalLMModel(_moe_config(_quantization())) + layer = model.model.layers[0] + + for proj in ("q_proj", "k_proj", "v_proj", "o_proj"): + assert type(getattr(layer.self_attn, proj)).__name__ == "QuantizedLinear" + # Routed experts go through fused QMoE, not per-expert dense MLPs. + assert layer.mlp.experts is None + assert type(layer.mlp.gate).__name__ == "SoftmaxTopKGate" + assert type(layer.input_layernorm).__name__ == "RMSNorm" + assert type(model.lm_head).__name__ == "Linear" + assert type(model.model.embed_tokens).__name__ == "Embedding" + # Router, norms and the embedding/LM-head tables keep float parameters + # (Olive leaves them unquantized), unlike the uint8 QMoE payloads. + params = dict(model.named_parameters()) + for name in ( + f"{_LAYER}mlp.gate.weight", + f"{_LAYER}input_layernorm.weight", + f"{_LAYER}post_attention_layernorm.weight", + "model.norm.weight", + "model.embed_tokens.weight", + "lm_head.weight", + ): + assert params[name].dtype.is_floating_point(), f"{name} must stay float" + assert params[f"{_LAYER}mlp.fc1_experts_weights"].dtype == ir.DataType.UINT8 + assert params[f"{_LAYER}self_attn.q_proj.weight"].dtype == ir.DataType.UINT8 + + def test_unquantized_model_uses_dense_expert_fallback(self): + """Without quantization the fused float experts still un-fuse and bind.""" + model = MoECausalLMModel(_moe_config(None)) + params = dict(model.named_parameters()) + p = f"{_LAYER}mlp." + + out = model.preprocess_weights( + { + p + "experts.gate_up_proj": torch.rand(_E, _FC1_OUT, _H), + p + "experts.down_proj": torch.rand(_E, _H, _INT), + p + "gate.weight": torch.rand(_E, _H), + } + ) + + assert model.model.layers[0].mlp.experts is not None + assert not any("fc1_experts_weights" in k for k in out) + for name, tensor in out.items(): + assert tuple(params[name].shape) == tuple(tensor.shape) diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index d17225e0a..e86e53c04 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -916,6 +916,50 @@ def test_glm4_moe_int4_quantizes_shared_expert(self): layer = module.model.layers[0] assert type(layer.mlp.shared_expert.down_proj).__name__ == "QuantizedLinear" + def test_qwen3_moe_olive_int4_emits_qmoe_and_matmulnbits(self): + """Olive-int4 Qwen3-MoE fuses routed experts into QMoE, attention into MatMulNBits. + + Graph/parameter-ABI check only — no weights are loaded and + ``preprocess_weights`` is not exercised here (see + ``src/mobius/models/moe_test.py`` for the state-dict side). It pins the + shapes the fused expert path *expects*: one ``com.microsoft::QMoE`` per + layer with expert-major ``[experts, 2*moe_inter, hidden*bits/8]`` + weights, plus MatMulNBits for the four quantized attention projections. + """ + from mobius._configs import QuantizationConfig + + qc = QuantizationConfig(bits=4, group_size=32, quant_method="olive", sym=True) + config = self._shared_moe_config("qwen3_moe", qc) + module = registry.get("qwen3_moe")(config) + pkg = CausalLMTask().build(module, config) + model = pkg["model"] + + qmoe = [n for n in model.graph if n.op_type == "QMoE"] + nbits = [n for n in model.graph if n.op_type == "MatMulNBits"] + assert len(qmoe) == 1, f"routed experts must fuse to one QMoE, got {len(qmoe)}" + assert len(nbits) == 4, f"attention q/k/v/o must emit 4 MatMulNBits, got {len(nbits)}" + + initializers = model.graph.initializers + experts = config.num_local_experts + assert tuple(initializers["model.layers.0.mlp.fc1_experts_weights"].shape) == ( + experts, + 2 * config.moe_intermediate_size, + config.hidden_size * qc.bits // 8, + ) + assert tuple(initializers["model.layers.0.mlp.fc2_experts_weights"].shape) == ( + experts, + config.hidden_size, + config.moe_intermediate_size * qc.bits // 8, + ) + assert tuple(initializers["model.layers.0.mlp.fc1_scales"].shape) == ( + experts, + 2 * config.moe_intermediate_size, + config.hidden_size // qc.group_size, + ) + # Symmetric quantization carries no zero points for the routed experts. + moe_initializers = [n for n in initializers if n.startswith("model.layers.0.mlp.")] + assert not any("zero_point" in name for name in moe_initializers) + class TestBuildGraphVisionLanguage: """Verify multimodal models build correctly."""