From 4ed080c14871cdfee41080a061a9857d48ef9d6b Mon Sep 17 00:00:00 2001 From: "Oke Iyanuoluwa E." Date: Sat, 17 Oct 2026 10:25:18 +0100 Subject: [PATCH] Pin the eager experts implementation in moe=True quantized checkpoints transformers picks the grouped_mm fused-experts forward by default, even on CPU. It transposes the expert weight, which a storage-only QuantTensor can't do, so running a reloaded Rtn/Gptq/KQuant(moe=True) checkpoint crashed with "Shape-movement / view ops ... are not supported on an Olive QuantTensor". save_pretrained doesn't persist the experts implementation, so finalize() now writes "experts_implementation": "eager" to the saved config.json (after save_metadata, which rewrites that file). from_pretrained then loads the checkpoint with the eager forward, which uses the weight without moving it. Fixes #2619 --- docs/source/features/quantization.md | 10 +++++----- olive/passes/pytorch/quant_utils.py | 30 ++++++++++++++++++++++++++++ test/passes/pytorch/test_rtn.py | 29 +++++++++++++++++++++++++++ 3 files changed, 64 insertions(+), 5 deletions(-) diff --git a/docs/source/features/quantization.md b/docs/source/features/quantization.md index 88ef37f991..3d3fa39376 100644 --- a/docs/source/features/quantization.md +++ b/docs/source/features/quantization.md @@ -117,16 +117,16 @@ before dispatching to their matmul kernel. Because Olive's 3D fused-expert `Quan `RuntimeError` at inference time -- even for architectures whose checkpoint layout (`is_transposed=False`) is fully supported for quantization. -**Workaround**: after loading a `moe=True`-quantized checkpoint for native PyTorch inference (as opposed to -consuming it via Mobius / ORT GenAI `ModelBuilder`), force the eager path once, before running any forward -pass: +To avoid this, `Rtn`, `Gptq` and `KQuant` with `moe=True` write `"experts_implementation": "eager"` to the +saved `config.json`, so `from_pretrained` loads the checkpoint with the eager forward by default. + +If you switch the implementation yourself after loading (or load with an explicit +`experts_implementation=...` argument), set it back to eager before running any forward pass: ```python model.set_experts_implementation("eager") ``` -This is tracked as a follow-up in [#2619](https://github.com/microsoft/Olive/issues/2619). - ### `modules_to_not_convert` and `overrides` `modules_to_not_convert` lists module-name patterns to exclude entirely, and `overrides` maps module-name diff --git a/olive/passes/pytorch/quant_utils.py b/olive/passes/pytorch/quant_utils.py index 28c37dd77c..4adedca9a9 100644 --- a/olive/passes/pytorch/quant_utils.py +++ b/olive/passes/pytorch/quant_utils.py @@ -7,11 +7,13 @@ import contextlib import inspect +import json import logging import re from collections.abc import Mapping from copy import deepcopy from dataclasses import dataclass +from pathlib import Path from typing import TYPE_CHECKING, Callable import torch @@ -1598,6 +1600,31 @@ def _retie_meta_parameters_for_save( ) +def _pin_eager_experts_implementation(output_model_path: str) -> None: + """Make the saved MoE checkpoint load with the ``eager`` fused-experts forward. + + transformers (>= 5) picks ``grouped_mm`` for fused experts by default, even on CPU. Its + ``_grouped_linear`` transposes the expert weight, which a storage-only ``QuantTensor`` cannot + do, so running a reloaded ``moe=True`` checkpoint crashes. The ``eager`` implementation uses + the weight without moving it. ``save_pretrained`` does not persist the implementation choice, + so write it to ``config.json`` where ``from_pretrained`` reads it back. + """ + config_path = Path(output_model_path) / "config.json" + if not config_path.exists(): + return + with config_path.open() as f: + config = json.load(f) + if config.get("experts_implementation") == "eager": + return + config["experts_implementation"] = "eager" + with config_path.open("w") as f: + json.dump(config, f, indent=2) + logger.info( + "Set experts_implementation='eager' in %s so the quantized fused experts can run after reloading.", + config_path, + ) + + def finalize( model: HfModelHandler, output_model_path: str, @@ -1694,5 +1721,8 @@ def finalize( ) save_model.save_pretrained(output_model_path, **save_kwargs) model.save_metadata(output_model_path) + if getattr(quant_config, "moe", False): + # after save_metadata, which rewrites config.json + _pin_eager_experts_implementation(output_model_path) return inherit_hf_from_hf(model, output_model_path, adapter_path=model.adapter_path) diff --git a/test/passes/pytorch/test_rtn.py b/test/passes/pytorch/test_rtn.py index 2791044a11..12452b2198 100644 --- a/test/passes/pytorch/test_rtn.py +++ b/test/passes/pytorch/test_rtn.py @@ -388,6 +388,35 @@ def test_rtn_moe_real_forward_after_reload(tmp_path: Path): assert not torch.isinf(out.logits).any() +def test_rtn_moe_checkpoint_runs_with_default_experts_implementation(tmp_path): + """A reloaded ``moe=True`` checkpoint must run without forcing ``eager`` by hand. + + transformers picks ``grouped_mm`` for fused experts by default, and that path transposes the + expert weight, which a ``QuantTensor`` cannot do. ``finalize()`` pins ``eager`` in the saved + config so the default load path works. + """ + import json + + input_model = _make_local_tiny_qwen3_moe(tmp_path / "input_model") + # drop the test helper's eager pin so the input model uses transformers' default + config_path = Path(input_model.model_path) / "config.json" + config = json.loads(config_path.read_text()) + config.pop("experts_implementation", None) + config_path.write_text(json.dumps(config, indent=2)) + + quantizer = create_pass_from_dict(Rtn, {"bits": 4, "group_size": -1, "moe": True}, disable_search=True) + output_path = tmp_path / "rtn" + out = quantizer.run(input_model, str(output_path)) + + assert json.loads((output_path / "config.json").read_text())["experts_implementation"] == "eager" + reloaded = out.load_model().eval() + assert reloaded.config._experts_implementation == "eager" + assert isinstance(reloaded.model.layers[0].mlp.experts.gate_up_proj.data, QuantTensor) + with torch.no_grad(): + logits = reloaded(torch.randint(0, 32, (1, 8))).logits + assert torch.isfinite(logits).all() + + def _make_local_tiny_tied_llama(save_path) -> HfModelHandler: """Save a tiny ``LlamaForCausalLM`` with ``tie_word_embeddings=True`` (built locally, no hub access).""" from transformers import LlamaConfig, LlamaForCausalLM