Pin the eager experts implementation in moe=True quantized checkpoints

docs/source/features/quantization.md

@@ -123,16 +123,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

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)

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