diff --git a/docs/source/features/quantization.md b/docs/source/features/quantization.md index dd606bf98a..476f92ec74 100644 --- a/docs/source/features/quantization.md +++ b/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 diff --git a/olive/passes/pytorch/moe_support.py b/olive/passes/pytorch/moe_support.py new file mode 100644 index 0000000000..a71b9d9780 --- /dev/null +++ b/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." + ) diff --git a/olive/passes/pytorch/rtn.py b/olive/passes/pytorch/rtn.py index 1180e0a1dd..fbbb826555 100644 --- a/olive/passes/pytorch/rtn.py +++ b/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) diff --git a/test/passes/pytorch/test_moe_support.py b/test/passes/pytorch/test_moe_support.py new file mode 100644 index 0000000000..2252a05132 --- /dev/null +++ b/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") diff --git a/test/passes/pytorch/test_rtn.py b/test/passes/pytorch/test_rtn.py index 51609405dd..1beb4dc0bb 100644 --- a/test/passes/pytorch/test_rtn.py +++ b/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])