Guard RTN MoE quantization by expert layout

docs/source/features/quantization.md

@@ -87,15 +87,14 @@ Mixture-of-Experts (MoE) experts at full precision. Three independent category f
 | --- | --- | --- |
 | `lm_head` | `false` | Also quantize the language-model head. |
 | `embeds` | `false` | Also quantize the input embeddings. |
-| `moe` | `false` | Also quantize MoE expert weights (both 3D fused-parameter experts such as gpt-oss / newer Qwen3-MoE, and `nn.ModuleList`-style experts such as Mixtral). |
+| `moe` | `false` | Also quantize MoE expert weights: classic per-expert `nn.ModuleList` experts (e.g. Mixtral, PhiMoE on older `transformers`) are always supported, and fused-expert modules are supported whenever the resolved experts class reports `is_transposed=False` (the K-last `(E, OUT, K)` layout used by most fused-experts architectures). Architectures that store a transposed `(E, K, OUT)` layout, such as gpt-oss, are rejected. Architectures whose experts class does not report `is_transposed` at all -- either because it predates the `transformers` fused-experts refactor, or because it has not adopted it (e.g. llama4, aria) -- are also rejected, since Olive cannot independently verify their layout. |
 
 The `moe` flag is **fail-closed**: when `moe` is `false`, every module under an experts subtree is skipped even
 if it looks like a plain `nn.Linear`. If the model config indicates an MoE architecture but Olive cannot resolve
 the experts subtree for that (unrecognized) architecture, the pass raises a clear error *before* modifying any
 parameter, rather than silently quantizing the experts.
 
-Only weight parameters are quantized. On fused-expert modules that also expose 2D bias parameters (e.g.
-gpt-oss's `gate_up_proj_bias` / `down_proj_bias`), the biases are left in full precision.
+Only weight parameters are quantized. Fused expert 2D bias parameters, when present, remain in full precision.
 
 ### `moe` and ONNX export
 

olive/passes/pytorch/moe_support.py

@@ -0,0 +1,95 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Shared layout-safety checks for fused MoE expert quantization."""
+
+from __future__ import annotations
+
+import torch
+
+
+class MoeSupportError(ValueError):
+    """Raised when an MoE quantization layout cannot be proven safe to quantize."""
+
+
+def check_moe_layout_support(
+    experts_modules: list[torch.nn.Module],
+    *,
+    model_type: str,
+    operation: str,
+) -> None:
+    """Fail closed unless every fused-experts module reports a K-last layout.
+
+    The quantizers using this check group along each weight tensor's last dimension.
+    Consequently, only fused weights stored ``(num_experts, out_features, in_features)``
+    are safe: the input/contraction dimension K must be last.
+
+    ``is_transposed`` is a boolean attribute exposed by ``transformers``'s
+    ``use_experts_implementation`` decorator on fused-experts modules (``False`` for K-last
+    architectures such as Mixtral/Qwen3-MoE/DeepSeek-V3/etc., ``True`` for architectures such
+    as gpt-oss). It is assigned once, from the decorator's own argument, inside the wrapped
+    ``__init__`` -- not derived from ``config`` -- so no checkpoint or config field can
+    influence it for a standard (non-``trust_remote_code``) experts implementation. Its
+    absence indicates either an undecorated experts implementation (e.g. an older
+    ``transformers`` release, or an architecture such as llama4/aria that has not adopted the
+    fused-experts decorator) or an unrecognized implementation, so it is treated as unsafe.
+
+    Classic per-expert ``nn.ModuleList`` experts (e.g. PhiMoE, DeepSeek-V3, or Mixtral on
+    older ``transformers`` releases without the fused-experts refactor) are exempt only when
+    the ``ModuleList`` itself owns no direct 3D parameter: Olive's quantizer selection only
+    ever groups a fused-experts module's *own* 3D parameters, and a plain ``ModuleList``
+    container cannot have one (each per-expert child is a plain 2D ``nn.Linear``, whose
+    weight is unconditionally ``(out, in)`` with K last). A ``ModuleList`` that does carry a
+    direct 3D parameter falls through to the normal ``is_transposed`` check below, which
+    rejects it (such a parameter has no ``is_transposed`` metadata to trust).
+
+    Note: this trusts ``is_transposed`` as reported by the resolved experts class. A custom
+    experts implementation loaded via ``trust_remote_code`` that misreports its own layout
+    (e.g. sets ``is_transposed=False`` while actually storing transposed weights) is not
+    caught here; that is treated as user-introduced misuse of an explicitly opted-in trust
+    boundary, not a layout Olive can independently verify.
+
+    Args:
+        experts_modules: The fused-experts (or classic ``nn.ModuleList``) modules discovered
+            for this model, one per MoE decoder layer.
+        model_type: The resolved HF ``config.model_type`` of the model being quantized. Not
+            used to decide the layout (that is entirely driven by ``is_transposed``); included
+            only for diagnostic context in error messages.
+        operation: Human-readable label for the calling pass/operation, prefixed onto error
+            messages. Purely cosmetic: it does not affect validation logic.
+
+    Raises:
+        MoeSupportError: If a fused-experts module's ``is_transposed`` attribute is missing,
+            not a ``bool``, or ``True``.
+
+    """
+    for experts in experts_modules:
+        if isinstance(experts, torch.nn.ModuleList) and not any(
+            param.dim() == 3 for param in experts.parameters(recurse=False)
+        ):
+            # No direct 3D parameter: Olive's quantizer selection never groups anything on
+            # this module itself, only on its per-expert nn.Linear children (K is
+            # unconditionally last for those), so there is no transposed-layout risk here.
+            continue
+
+        actual_class = type(experts).__name__
+        is_transposed = getattr(experts, "is_transposed", None)
+        if not isinstance(is_transposed, bool):
+            raise MoeSupportError(
+                f"{operation} cannot verify the layout of experts module '{actual_class}' "
+                f"(model_type='{model_type}') because attribute 'is_transposed' is missing "
+                "or not a boolean. This typically means an experts implementation that has "
+                "not adopted the fused-experts decorator (e.g. an older transformers release, "
+                "or an architecture such as llama4/aria), or an unrecognized implementation. "
+                "Quantization is refused rather than risking grouping along the wrong "
+                "dimension. Re-run with moe=False so the experts stay in full precision."
+            )
+        if is_transposed:
+            raise MoeSupportError(
+                f"{operation} refuses experts module '{actual_class}' (model_type='{model_type}') "
+                "because it reports a transposed fused-weight layout (E, K, OUT), while "
+                "last-dimension grouping requires (E, OUT, K) with K last. Architectures such "
+                "as gpt-oss store this transposed layout. Re-run with moe=False so the "
+                "experts stay in full precision."
+            )

olive/passes/pytorch/rtn.py

@@ -10,6 +10,7 @@
 import torch
 
 from olive.passes import Pass
+from olive.passes.pytorch.moe_support import check_moe_layout_support
 from olive.passes.pytorch.quant_utils import finalize, get_quantizer_config, prepare_model
 
 if TYPE_CHECKING:
@@ -44,5 +45,20 @@ def _run_for_config(
 
         """
         wrapper, qcfg, retie_word_embeddings = prepare_model(model, config, allow_quantized=True)
+        # Gate on this invocation's own ``moe`` request, not ``qcfg.moe`` -- ``prepare_model``
+        # ORs in any pre-existing checkpoint's ``moe`` flag (see ``quant_utils.prepare_model``),
+        # so ``qcfg.moe`` can be True even when this run itself passed ``moe=False``.
+        if getattr(config, "moe", False):
+            experts_modules = [
+                experts
+                for layer in wrapper.get_layer_wrappers()
+                if (experts := layer.get_experts(return_name=False)) is not None
+            ]
+            if experts_modules:
+                check_moe_layout_support(
+                    experts_modules,
+                    model_type=wrapper.model_type,
+                    operation="RTN MoE quantization",
+                )
         device = "cuda" if torch.cuda.is_available() else "cpu"
         return finalize(model, output_model_path, wrapper, qcfg, device, retie_word_embeddings=retie_word_embeddings)

test/passes/pytorch/test_moe_support.py

@@ -0,0 +1,92 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Tests for shared fused-MoE layout safety checks."""
+
+import pytest
+import torch
+
+from olive.passes.pytorch.moe_support import MoeSupportError, check_moe_layout_support
+
+
+class _FakeExperts(torch.nn.Module):
+    def __init__(self, **flags):
+        super().__init__()
+        for key, value in flags.items():
+            setattr(self, key, value)
+
+
+def make_fake_experts(class_name: str = "FakeExperts", **flags) -> torch.nn.Module:
+    """Build a stand-in experts module whose class name is ``class_name``."""
+    return type(class_name, (_FakeExperts,), {})(**flags)
+
+
+@pytest.mark.parametrize("model_type", ["mixtral", "qwen3_moe", "some_future_moe_architecture"])
+def test_check_support_accepts_any_model_type_when_not_transposed(model_type: str):
+    """``is_transposed=False`` is accepted for any model_type, not just a fixed allow-list."""
+    check_moe_layout_support(
+        [make_fake_experts(is_transposed=False)],
+        model_type=model_type,
+        operation="Test MoE quantization",
+    )
+
+
+def test_check_support_rejects_missing_capability():
+    with pytest.raises(MoeSupportError, match="is_transposed") as exc_info:
+        check_moe_layout_support([make_fake_experts()], model_type="qwen3_moe", operation="Test MoE quantization")
+    assert "moe=False" in str(exc_info.value)
+
+
+@pytest.mark.parametrize("bad_value", [None, 0, 1, "False", "", [], torch.tensor(False)])
+def test_check_support_rejects_non_bool_is_transposed(bad_value):
+    """Non-boolean values (e.g. a stray ``None``/tensor/string) must not be treated as falsy-safe."""
+    with pytest.raises(MoeSupportError, match="is_transposed") as exc_info:
+        check_moe_layout_support(
+            [make_fake_experts(is_transposed=bad_value)], model_type="qwen3_moe", operation="Test MoE quantization"
+        )
+    assert "moe=False" in str(exc_info.value)
+
+
+def test_check_support_rejects_transposed_layout():
+    with pytest.raises(MoeSupportError, match=r"\(E, K, OUT\)") as exc_info:
+        check_moe_layout_support(
+            [make_fake_experts(is_transposed=True)],
+            model_type="gpt_oss",
+            operation="Test MoE quantization",
+        )
+    message = str(exc_info.value)
+    assert "moe=False" in message
+    assert "gpt_oss" in message
+
+
+def test_check_support_exempts_module_list_experts_with_no_direct_3d_param():
+    """Classic per-expert ``nn.ModuleList`` experts are exempt from the layout check.
+
+    E.g. PhiMoE/Mixtral on transformers releases without the fused-experts refactor are
+    unconditionally K-last (each child is a plain 2D ``nn.Linear``), so they must not be
+    rejected for a missing ``is_transposed`` attribute.
+    """
+    module_list = torch.nn.ModuleList([torch.nn.Linear(4, 4) for _ in range(2)])
+    check_moe_layout_support([module_list], model_type="any_model_type", operation="Test MoE quantization")
+
+
+def test_check_support_rejects_module_list_with_direct_3d_param():
+    """A ``ModuleList`` that also owns a direct 3D parameter is not exempt.
+
+    Olive's quantizer selection would group that parameter's last dimension directly, so it
+    needs the same ``is_transposed`` guarantee as any other fused-experts module -- and a
+    bare ``ModuleList`` has no such attribute to trust.
+    """
+    module_list = torch.nn.ModuleList([torch.nn.Linear(4, 4)])
+    module_list.fused_weight = torch.nn.Parameter(torch.zeros(2, 4, 4))
+    with pytest.raises(MoeSupportError, match="is_transposed"):
+        check_moe_layout_support([module_list], model_type="any_model_type", operation="Test MoE quantization")
+
+
+def test_check_support_checks_every_module_in_a_mixed_list():
+    """A single transposed module among otherwise-safe ones still triggers rejection."""
+    safe = make_fake_experts("SafeExperts", is_transposed=False)
+    unsafe = make_fake_experts("UnsafeExperts", is_transposed=True)
+    with pytest.raises(MoeSupportError, match="UnsafeExperts"):
+        check_moe_layout_support([safe, unsafe], model_type="mixtral", operation="Test MoE quantization")

test/passes/pytorch/test_rtn.py

@@ -14,10 +14,22 @@
 from olive.model import HfModelHandler
 from olive.passes.olive_pass import create_pass_from_dict
 from olive.passes.pytorch.gptq import Gptq
+from olive.passes.pytorch.moe_support import MoeSupportError
+from olive.passes.pytorch.quant_utils import prepare_model
 from olive.passes.pytorch.rtn import Rtn
 from test.utils import get_tiny_phi3
 
 
+def _save_trivial_tokenizer(save_path: Path, vocab_size: int) -> None:
+    """Save a local tokenizer so pass metadata serialization never needs the hub."""
+    from tokenizers import Tokenizer, models, pre_tokenizers
+    from transformers import PreTrainedTokenizerFast
+
+    tokenizer = Tokenizer(models.WordLevel({f"t{i}": i for i in range(vocab_size)}, unk_token="t0"))
+    tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
+    PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token="t0", pad_token="t0").save_pretrained(save_path)
+
+
 def _is_quant(module: torch.nn.Module) -> bool:
     if not isinstance(module, (torch.nn.Linear, torch.nn.Embedding)):
         return False
@@ -55,6 +67,112 @@ def _make_local_tiny_mixtral(save_path):
     return HfModelHandler(model_path=str(save_path))
 
 
+def _make_local_tiny_qwen3_moe(save_path) -> HfModelHandler:
+    """Save a tiny K-last fused-experts model without downloading a checkpoint."""
+    from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM
+
+    torch.manual_seed(0)
+    save_path = Path(save_path)
+    save_path.mkdir(parents=True, exist_ok=True)
+    config = Qwen3MoeConfig(  # pylint: disable=unexpected-keyword-arg
+        vocab_size=32,
+        hidden_size=16,
+        intermediate_size=16,
+        moe_intermediate_size=8,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        num_experts=2,
+        num_experts_per_tok=1,
+        decoder_sparse_step=1,
+        head_dim=8,
+        experts_implementation="eager",
+    )
+    Qwen3MoeForCausalLM(config).save_pretrained(save_path)
+    _save_trivial_tokenizer(save_path, config.vocab_size)
+    return HfModelHandler(model_path=str(save_path))
+
+
+def test_rtn_moe_refuses_transposed_layout_before_finalize(tmp_path: Path, monkeypatch):
+    """A transposed-layout (``is_transposed=True``) experts module fails before finalize.
+
+    The rejection must happen before quantization or output serialization.
+    """
+    input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model")
+
+    def patched_prepare_model(*args, **kwargs):
+        wrapper, qcfg, retie = prepare_model(*args, **kwargs)
+        for layer in wrapper.get_layer_wrappers():
+            experts = layer.get_experts(return_name=False)
+            if experts is not None:
+                experts.is_transposed = True
+        return wrapper, qcfg, retie
+
+    def unexpected_finalize(*args, **kwargs):
+        pytest.fail("finalize must not run after MoE layout support is rejected")
+
+    monkeypatch.setattr("olive.passes.pytorch.rtn.prepare_model", patched_prepare_model)
+    monkeypatch.setattr("olive.passes.pytorch.rtn.finalize", unexpected_finalize)
+
+    quantizer = create_pass_from_dict(Rtn, {"moe": True}, disable_search=True)
+    output_path = tmp_path / "rtn"
+    with pytest.raises(MoeSupportError, match=r"\(E, K, OUT\).*moe=False"):
+        quantizer.run(input_model, str(output_path))
+    assert not output_path.exists()
+
+
+def test_rtn_moe_false_does_not_run_layout_gate(tmp_path: Path, monkeypatch):
+    """The default dense path must not invoke fused-experts layout validation."""
+    input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model")
+
+    def unexpected_gate(*args, **kwargs):
+        pytest.fail("MoE layout support check must not run when moe=False")
+
+    monkeypatch.setattr("olive.passes.pytorch.rtn.check_moe_layout_support", unexpected_gate)
+    quantizer = create_pass_from_dict(Rtn, {"moe": False, "group_size": -1}, disable_search=True)
+    out = quantizer.run(input_model, str(tmp_path / "rtn"))
+
+    loaded = out.load_model()
+    experts = loaded.model.layers[0].mlp.experts
+    assert not any(isinstance(param.data, QuantTensor) for param in experts.parameters())
+    assert isinstance(loaded.model.layers[0].self_attn.q_proj.weight.data, QuantTensor)
+
+
+def test_rtn_moe_gate_ignores_prior_checkpoint_moe_flag(tmp_path: Path, monkeypatch):
+    """Regression test: a second RTN pass with moe=False must not re-run the layout gate.
+
+    ``prepare_model`` ORs a pre-existing checkpoint's ``moe`` flag into the merged
+    ``qcfg.moe`` (see ``quant_utils.prepare_model``), so gating on ``qcfg.moe`` would make
+    this second, moe=False invocation incorrectly re-run fused-experts layout validation
+    -- something this run never asked for. The gate must key off this invocation's own
+    ``config.moe`` request instead.
+    """
+    input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model")
+    first_pass = create_pass_from_dict(Rtn, {"moe": True, "group_size": -1}, disable_search=True)
+    quantized = first_pass.run(input_model, str(tmp_path / "rtn_first"))
+
+    def unexpected_gate(*args, **kwargs):
+        pytest.fail("MoE layout support check must not re-run when this invocation requests moe=False")
+
+    monkeypatch.setattr("olive.passes.pytorch.rtn.check_moe_layout_support", unexpected_gate)
+    second_pass = create_pass_from_dict(Rtn, {"moe": False, "lm_head": True, "group_size": -1}, disable_search=True)
+    out = second_pass.run(quantized, str(tmp_path / "rtn_second"))
+
+    loaded = out.load_model()
+    assert isinstance(loaded.lm_head.weight.data, QuantTensor)
+
+
+def test_rtn_moe_k_last_layout_quantizes_experts(tmp_path: Path):
+    """A K-last (``is_transposed=False``) fused-experts model quantizes successfully end to end."""
+    input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model")
+    quantizer = create_pass_from_dict(Rtn, {"moe": True, "group_size": -1}, disable_search=True)
+    out = quantizer.run(input_model, str(tmp_path / "rtn"))
+
+    loaded = out.load_model()
+    experts = loaded.model.layers[0].mlp.experts
+    assert any(isinstance(param.data, QuantTensor) for param in experts.parameters())
+
+
 @pytest.mark.parametrize("group_size", [-1, 16])
 @pytest.mark.parametrize("sym", [True, False])
 @pytest.mark.parametrize("lm_head", [True, False])