From 5a329b840ee530604d1ae6e21d5fe37265a79a87 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 23:12:45 +0000
Subject: [PATCH 01/26] Quantize fused experts under grouped_mm by intercepting
 _grouped_linear

transformers' grouped_mm experts forward calls the module-level
_grouped_linear once per projection with the full 3-D weight, in the same
(num_experts, out, in) layout the eager forward slices per expert. The
_QuantFusedExperts wrapper now replaces it too: each expert's slice passes its
own weight quantizer and the batch the shared input quantizer, so the grouped
forward calibrates and fake-quantizes the same quantizers as eager.

A wrapper that defines its weights in get_expert_weight, which the grouped
forward never reaches, must say it quantizes on the eager forward only
(_eager_only), and its grouped forward raises; the FP8 and non-gated wrappers
do. Calibration still runs every model on eager: the pipelines switch to the
grouped forward in the next PR.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/quantization/plugins/huggingface.py | 126 ++++++-
 .../plugins/test_fused_experts.py             | 309 +++++++++++++++++-
 2 files changed, 419 insertions(+), 16 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index a0b896b7903..ed11c8d5bd5 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -53,6 +53,11 @@
 
 logger = logging.getLogger(__name__)
 
+try:
+    import transformers.integrations.moe as _transformers_moe
+except ImportError:
+    _transformers_moe = None
+
 try:
     import kitchen
     from kitchen.fa import KitchenFlashAttentionModule
@@ -793,9 +798,21 @@ class _QuantFusedExperts(_QuantFunctionalMixin):
     Verified compatible models: Mixtral, Qwen2-MoE, Qwen3-MoE, Qwen3.5-MoE,
     DeepSeek-V3, Jamba, OLMoE.
 
-    Limitation: only works when ``experts_implementation="eager"`` (default).
-    ``batched_mm`` / ``grouped_mm`` backends use ``torch.bmm`` /
-    ``torch._grouped_mm`` instead of ``F.linear`` and are not intercepted.
+    Backends: both the ``eager`` (per-expert ``F.linear``) and ``grouped_mm``
+    (``transformers.integrations.moe._grouped_linear`` →
+    ``torch._grouped_mm`` / ``torch.nn.functional.grouped_mm`` / fallback)
+    forwards are intercepted. Both calibrate the same per-expert weight quantizers
+    and shared input quantizers, and export the same keys; the activation scales
+    they calibrate can differ. Only one backend's hook fires per forward
+    (selected by ``config._experts_implementation``), so installing both is safe.
+    Calibration still runs every model on ``eager``
+    (``force_eager_experts_impl_on_the_fly``).
+
+    Limitation: the ``grouped_mm`` path assumes the non-transposed weight layout
+    ``(num_experts, out_dim, in_dim)`` used by all standard fused experts. The
+    transposed layout (e.g. GptOss, which has its own quantized wrapper) is not
+    handled here. ``batched_mm`` (``torch.bmm``) and ``sonicmoe`` (opaque CUDA
+    kernel) backends are not intercepted.
 
     The non-gated variant (``up_proj`` instead of ``gate_up_proj``, used by
     NemotronH) is handled by :class:`_QuantNonGatedFusedExperts` via the
@@ -812,6 +829,22 @@ class _QuantFusedExperts(_QuantFunctionalMixin):
     # Under expert parallelism each rank's experts see only the tokens routed to them, so calibration
     # shares the layer's input amax over the expert-parallel group (see max_calibrate).
     _share_input_amax_across_ep = True
+    # Whether the wrapper quantizes only on the eager per-expert forward, so a model holding it must
+    # calibrate there.
+    _eager_only = False
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        # The grouped_mm forward quantizes the stored weight and never reaches get_expert_weight, so
+        # a wrapper that defines its weight domain there would calibrate the wrong weights on it.
+        if (
+            cls.get_expert_weight is not _QuantFusedExperts.get_expert_weight
+            and not cls._eager_only
+        ):
+            raise TypeError(
+                f"{cls.__name__} overrides get_expert_weight, which the grouped_mm forward bypasses; "
+                "set _eager_only = True so that its models calibrate on the eager forward."
+            )
 
     @property
     def _first_proj_input_quantizer_attr(self) -> str:
@@ -822,7 +855,7 @@ def _first_proj_weight_quantizers_attr(self) -> str:
         return f"{self._first_proj_attr}_weight_quantizers"
 
     def _get_expert_idx_from_first_proj(self, weight: torch.Tensor) -> int:
-        """Recover expert index from a first-projection weight slice's storage offset.
+        """Recover expert index from the first-projection weight slice's storage offset.
 
         When HF indexes ``<first_proj>[idx]``, the result is a view sharing the
         same underlying storage.  The offset delta divided by the stride along
@@ -859,12 +892,36 @@ def _setup(self):
         # Per forward: each projection's experts, fake-quantized together (see _quantized_weight).
         self._register_temp_attribute("_stacked_weights", {})
 
+    def _quantize_grouped_weight(self, weight: torch.Tensor, quantizers: nn.ModuleList):
+        """Fake-quantize a fused 3-D weight per expert.
+
+        Each expert slice passes through its own weight quantizer (collecting
+        ``amax`` during calibration, fake-quantizing afterwards), matching the
+        eager per-expert ``F.linear`` semantics. Returns the weight untouched when
+        every quantizer is disabled so passthrough stays exact (and cheap).
+        """
+        if all(not q.is_enabled for q in quantizers):
+            return weight
+        if torch.is_grad_enabled():
+            # Slices written into one buffer would chain a full-size gradient copy per expert into
+            # the backward (QAT), so stack them there.
+            return torch.stack([quantizers[idx](weight[idx]) for idx in range(self.num_experts)])
+        # Each slice goes into the output as it is made: a stack would hold all of them beside it.
+        out = torch.empty_like(weight)
+        for idx in range(self.num_experts):
+            out[idx] = quantizers[idx](weight[idx])
+        return out
+
     @property
     def functionals_to_replace(self):
         _orig_linear = torch.nn.functional.linear
+        # First-projection quantizers, by layout: gate_up_proj_* (gated) or up_proj_* (non-gated).
+        # Bound once; the objects are stable across the forward.
+        first_input_q = getattr(self, self._first_proj_input_quantizer_attr)
+        first_weight_qs = getattr(self, self._first_proj_weight_quantizers_attr)
 
         # The HF fused expert forward calls F.linear exactly twice per expert
-        # in strict alternation: first for gate_up_proj, then for down_proj.
+        # in strict alternation: first the up/gate_up projection, then down_proj.
         # forward() resets the toggle before each call to super().forward().
         def _quantized_linear(input, weight, bias=None):
             if self._down_proj_linear:
@@ -876,19 +933,51 @@ def _quantized_linear(input, weight, bias=None):
             else:
                 idx = self._get_expert_idx_from_first_proj(weight)
                 self._current_expert_idx = idx
-                input = getattr(self, self._first_proj_input_quantizer_attr)(input)
+                input = first_input_q(input)
                 weight = self._quantized_weight(
                     self._first_proj_attr,
-                    getattr(self, self._first_proj_weight_quantizers_attr),
+                    first_weight_qs,
                     idx,
                     input.dtype,
                 )
             self._down_proj_linear = not self._down_proj_linear
             return _orig_linear(input, weight, bias)
 
-        return [
-            (torch.nn.functional, "linear", _quantized_linear),
-        ]
+        replacements: list[tuple] = [(torch.nn.functional, "linear", _quantized_linear)]
+
+        # grouped_mm backend: HF's grouped_mm_experts_forward calls the module-level
+        # _grouped_linear exactly twice (gate_up then down) with the FULL 3-D weight
+        # in the same (num_experts, out, in) layout as eager F.linear, so the same
+        # per-expert weight quantizers / shared input quantizers apply. Only one
+        # backend's hook fires per forward, so patching both F.linear and
+        # _grouped_linear is safe.
+        if _transformers_moe is not None and hasattr(_transformers_moe, "_grouped_linear"):
+            _orig_grouped_linear = _transformers_moe._grouped_linear
+
+            def _quantized_grouped_linear(input, weight, offs, bias=None, is_transposed=False):
+                if self._eager_only:
+                    # Quantizing the stored weight here would bypass get_expert_weight, which defines
+                    # this wrapper's weights.
+                    raise RuntimeError(
+                        f"{type(self).__name__} quantizes on the eager experts forward only, but its "
+                        'grouped_mm forward ran. Load the model with experts_implementation="eager".'
+                    )
+                quantizers = (
+                    self.down_proj_weight_quantizers if self._down_proj_linear else first_weight_qs
+                )
+                if self._down_proj_linear:
+                    input = self.down_proj_input_quantizer(input)
+                else:
+                    input = first_input_q(input)
+                weight = self._quantize_grouped_weight(weight, quantizers)
+                self._down_proj_linear = not self._down_proj_linear
+                return _orig_grouped_linear(
+                    input, weight, offs, bias=bias, is_transposed=is_transposed
+                )
+
+            replacements.append((_transformers_moe, "_grouped_linear", _quantized_grouped_linear))
+
+        return replacements
 
     def _quantized_weight(self, name, quantizers, idx, dtype):
         """Expert ``idx``'s ``name`` projection, fake-quantized by its own weight quantizer.
@@ -987,11 +1076,18 @@ class _QuantNonGatedFusedExperts(_QuantFusedExperts):
 
     _first_proj_attr = "up_proj"
     _is_gated = False
+    # The grouped_mm interception is validated for the gated layout only, so these keep
+    # calibrating on the eager forward.
+    _eager_only = True
 
 
 class _QuantFP8Experts(_QuantFusedExperts):
     """Scale-aware PTQ for native Transformers FP8 experts, dequantized per expert."""
 
+    # The block scales apply in get_expert_weight, which only the eager forward reaches: transformers'
+    # FP8 grouped_mm forward runs its own kernels, which nothing here intercepts.
+    _eager_only = True
+
     @property
     def _first_proj_attr(self):
         return "gate_up_proj" if self.has_gate else "up_proj"
@@ -1041,6 +1137,11 @@ def unpack_weight(self):
                 setattr(self, name, nn.Parameter(unpacked, requires_grad=False))
 
 
+def _is_quant_fused_experts_module(module):
+    """Return True for a converted HF fused-MoE-experts quantization wrapper."""
+    return isinstance(module, _QuantFusedExperts)
+
+
 def _get_fused_experts_quantizer_attr_names(module):
     """Return quantizer attribute names for a converted fused-experts module."""
     first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
@@ -1052,11 +1153,6 @@ def _get_fused_experts_quantizer_attr_names(module):
     )
 
 
-def _is_quant_fused_experts_module(module):
-    """Return True for a converted HF fused-MoE-experts quantization wrapper."""
-    return isinstance(module, _QuantFusedExperts)
-
-
 @contextmanager
 def patch_compressed_linear_loading():
     """Context manager that patches CompressedLinear to survive custom ``_init_weights`` calls.
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 0847227c412..7f0b13c397d 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -17,6 +17,7 @@
 
 import copy
 import warnings
+import weakref
 from contextlib import nullcontext
 
 import pytest
@@ -55,6 +56,7 @@
     register_sparse_moe_on_the_fly,
 )
 from modelopt.torch.quantization.utils.core_utils import weight_attr_names
+from modelopt.torch.utils.moe import torch_a2a_experts_forward
 
 
 def test_native_fp8_experts_dequantized_before_nvfp4_export():
@@ -285,6 +287,121 @@ def _route_once_to_each_expert(model):
     model.moe.experts(hidden_states, top_k_index, top_k_weights)
 
 
+# ---------------------------------------------------------------------------
+# grouped_mm backend fixtures: a synthetic fused-experts module that can run
+# either the eager (F.linear) forward or the real transformers grouped_mm
+# forward, selected by ``config._experts_implementation``. Used to verify
+# grouped_mm calibration matches eager and that local-Hessian capture works in
+# grouped mode. The class avoids the ``@use_experts_implementation`` decorator so
+# it imports on transformers builds that predate the grouped_mm backend.
+# ---------------------------------------------------------------------------
+try:
+    import transformers.integrations.moe as _moe_integration
+
+    HAS_GROUPED_MM = hasattr(_moe_integration, "_grouped_linear")
+except ImportError:
+    HAS_GROUPED_MM = False
+
+
+class _ExpertsConfig:
+    """Minimal config carrying the experts-implementation selector."""
+
+    def __init__(self, impl="grouped_mm"):
+        self._experts_implementation = impl
+        self.hidden_act = "silu"
+        self.ep_group = None  # the expert-parallel group a torch_a2a forward dispatches over
+
+
+class _SyntheticGroupedExperts(nn.Module):
+    """Fused experts dispatching to eager or the real grouped_mm forward.
+
+    Mirrors the HF transformers 5.x fused-experts contract
+    (``has_gate``/``has_bias``/``is_transposed``/``_apply_gate`` + 3-D
+    ``gate_up_proj``/``down_proj``) that ``grouped_mm_experts_forward`` expects.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.num_experts = NUM_EXPERTS
+        self.hidden_dim = HIDDEN_DIM
+        self.intermediate_dim = INTERMEDIATE_DIM
+        self.has_gate = True
+        self.has_bias = False
+        self.is_transposed = False
+        self.is_concatenated = True
+        # What transformers' @use_experts_implementation sets (5.18+ reads it in grouped_mm).
+        self._is_expert_parallel = False
+        self.gate_up_proj = nn.Parameter(
+            torch.randn(NUM_EXPERTS, 2 * INTERMEDIATE_DIM, HIDDEN_DIM) * 0.02
+        )
+        self.down_proj = nn.Parameter(torch.randn(NUM_EXPERTS, HIDDEN_DIM, INTERMEDIATE_DIM) * 0.02)
+        self.act_fn = nn.SiLU()
+
+    def _apply_gate(self, gate_up_out):
+        gate, up = gate_up_out.chunk(2, dim=-1)
+        return self.act_fn(gate) * up
+
+    def _eager_forward(self, hidden_states, top_k_index, top_k_weights):
+        final_hidden_states = torch.zeros_like(hidden_states)
+        with torch.no_grad():
+            expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+        for expert_idx in expert_hit:
+            expert_idx = expert_idx[0]
+            if expert_idx == self.num_experts:
+                continue
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            current_state = hidden_states[token_idx]
+            gate, up = F.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
+            current_hidden_states = self.act_fn(gate) * up
+            current_hidden_states = F.linear(current_hidden_states, self.down_proj[expert_idx])
+            current_hidden_states = (
+                current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
+            )
+            final_hidden_states.index_add_(
+                0, token_idx, current_hidden_states.to(final_hidden_states.dtype)
+            )
+        return final_hidden_states
+
+    def forward(self, hidden_states, top_k_index, top_k_weights):
+        if self.config._experts_implementation == "grouped_mm":
+            return _moe_integration.grouped_mm_experts_forward(
+                self, hidden_states, top_k_index, top_k_weights
+            )
+        if self.config._experts_implementation == "torch_a2a":
+            return torch_a2a_experts_forward(
+                self, hidden_states, top_k_index, top_k_weights, self.config.ep_group
+            )
+        return self._eager_forward(hidden_states, top_k_index, top_k_weights)
+
+
+class _GroupedSparseMoeBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.gate = _SyntheticTopKRouter()
+        self.experts = _SyntheticGroupedExperts(config)
+
+    def forward(self, hidden_states):
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        _, top_k_weights, top_k_index = self.gate(hidden_states)
+        hidden_states = self.experts(hidden_states, top_k_index, top_k_weights)
+        return hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+
+
+class _GroupedMoEModel(nn.Module):
+    """Minimal MoE model whose experts run the selected backend."""
+
+    def __init__(self, impl):
+        super().__init__()
+        self.config = _ExpertsConfig(impl=impl)
+        self.moe = _GroupedSparseMoeBlock(self.config)
+
+    def forward(self, x):
+        return self.moe(x)
+
+
 # ---------------------------------------------------------------------------
 # Tests for _is_fused_experts_module
 # ---------------------------------------------------------------------------
@@ -355,6 +472,28 @@ def test_fused_experts_only_registration(self):
         assert QuantModuleRegistry.get(block_type) is None
         self._cleanup_registry(expert_type)
 
+    def test_a_wrapper_that_defines_its_weights_in_get_expert_weight_is_eager_only(self):
+        """The grouped_mm forward quantizes the stored weight without get_expert_weight, so a wrapper
+        that defines its weight domain there must calibrate on the eager forward. One that does not
+        say so is refused when it is defined, rather than calibrating the stored weights."""
+        with pytest.raises(TypeError, match="_eager_only"):
+
+            class _Scaled(_QuantFusedExperts):
+                def get_expert_weight(self, name, idx, dtype=None):
+                    return super().get_expert_weight(name, idx, dtype)
+
+        class _ScaledOnEager(_QuantFusedExperts):
+            _eager_only = True
+
+            def get_expert_weight(self, name, idx, dtype=None):
+                return super().get_expert_weight(name, idx, dtype)
+
+        # Inheriting the override does not take it off the eager forward either.
+        with pytest.raises(TypeError, match="_eager_only"):
+
+            class _ScaledOnGroupedMM(_ScaledOnEager):
+                _eager_only = False
+
     def test_convert_creates_quantizers(self):
         """After conversion, fused experts should have shared input and per-expert weight quantizers."""
         model = _TinyMoEModel()
@@ -1595,6 +1734,174 @@ def test_unrelated_dotted_number_unchanged(self):
         )
 
 
+# ---------------------------------------------------------------------------
+# Tests for the grouped_mm backend: _QuantFusedExperts intercepts the real
+# transformers grouped_mm forward (via _grouped_linear), and grouped calibration
+# must produce the same per-expert quantizer state as the eager F.linear path.
+# ---------------------------------------------------------------------------
+@pytest.mark.skipif(not HAS_GROUPED_MM, reason="transformers grouped_mm backend unavailable")
+class TestGroupedMMExperts:
+    @staticmethod
+    def _cleanup_registry(mod_type):
+        if QuantModuleRegistry.get(mod_type) is not None:
+            QuantModuleRegistry.unregister(mod_type)
+
+    @staticmethod
+    def _quant_cfg():
+        return {
+            "quant_cfg": [
+                {"quantizer_name": "*", "enable": False},
+                {
+                    "quantizer_name": "*gate_up_proj_input_quantizer",
+                    "cfg": {"num_bits": 8, "axis": None},
+                },
+                {
+                    "quantizer_name": "*down_proj_input_quantizer",
+                    "cfg": {"num_bits": 8, "axis": None},
+                },
+                {
+                    "quantizer_name": "*gate_up_proj_weight_quantizer",
+                    "cfg": {"num_bits": 8, "axis": 0},
+                },
+                {
+                    "quantizer_name": "*down_proj_weight_quantizer",
+                    "cfg": {"num_bits": 8, "axis": 0},
+                },
+            ],
+            "algorithm": "max",
+        }
+
+    @staticmethod
+    def _forward_loop_all_experts(m):
+        # Route every token to all experts (round-robin) so eager and grouped both
+        # calibrate every expert's weight quantizer — eager only sees routed experts.
+        torch.manual_seed(0)
+        seq = 16
+        hidden = torch.randn(seq, HIDDEN_DIM)
+        top_k_index = torch.stack(
+            [torch.arange(seq) % NUM_EXPERTS, (torch.arange(seq) + 1) % NUM_EXPERTS], dim=1
+        )
+        top_k_weights = torch.softmax(torch.randn(seq, TOP_K), dim=1)
+        with torch.no_grad():
+            m.moe.experts(hidden, top_k_index, top_k_weights)
+
+    def test_forward_passthrough_matches(self):
+        """Converted grouped experts with their quantizers disabled match the unconverted forward."""
+        self._cleanup_registry(_SyntheticGroupedExperts)
+        torch.manual_seed(0)
+        model = _GroupedMoEModel("grouped_mm")
+        ref = _SyntheticGroupedExperts(_ExpertsConfig("grouped_mm"))
+        ref.load_state_dict(model.moe.experts.state_dict())
+
+        register_fused_experts_on_the_fly(model)
+        converted = QuantModuleRegistry.convert(model.moe.experts)
+        for q in converted.modules():
+            if isinstance(q, TensorQuantizer):
+                q.disable()
+
+        seq = 8
+        hidden = torch.randn(seq, HIDDEN_DIM)
+        top_k_index = torch.randint(0, NUM_EXPERTS, (seq, TOP_K))
+        top_k_weights = torch.softmax(torch.randn(seq, TOP_K), dim=1)
+        with torch.no_grad():
+            out_ref = ref(hidden, top_k_index, top_k_weights)
+            out_test = converted(hidden, top_k_index, top_k_weights)
+        assert torch.allclose(out_ref, out_test, atol=1e-6), (
+            f"Max diff: {(out_ref - out_test).abs().max().item()}"
+        )
+        self._cleanup_registry(_SyntheticGroupedExperts)
+
+    def test_calibration_matches_eager(self):
+        """grouped_mm PTQ calibration yields the same quantizer state as eager PTQ.
+
+        Identical weights + identical routing ⇒ the shared input quantizers see the
+        same activation set and each per-expert weight quantizer sees the same slice,
+        so amax must match (max abs is order-free; only the down-proj input flows
+        through a GEMM, hence the small tolerance there).
+        """
+        self._cleanup_registry(_SyntheticGroupedExperts)
+        torch.manual_seed(0)
+        eager_model = _GroupedMoEModel("eager")
+        grouped_model = _GroupedMoEModel("grouped_mm")
+        grouped_model.load_state_dict(eager_model.state_dict())
+
+        for model in (eager_model, grouped_model):
+            impl = model.config._experts_implementation
+            # quantize() calibrates every model on eager; calibrate on the backend under test.
+            mtq.quantize(model, self._quant_cfg())
+            model.config._experts_implementation = impl
+            mtq.calibrate(model, "max", self._forward_loop_all_experts)
+
+        e = eager_model.moe.experts
+        g = grouped_model.moe.experts
+
+        # Shared input quantizers (down flows through a GEMM -> small tolerance).
+        for q_name in ("gate_up_proj_input_quantizer", "down_proj_input_quantizer"):
+            ea = getattr(e, q_name).amax
+            ga = getattr(g, q_name).amax
+            assert ea is not None and ga is not None, f"{q_name} not calibrated"
+            assert torch.allclose(ea, ga, atol=1e-3, rtol=1e-2), (
+                f"{q_name} amax mismatch: eager={ea} grouped={ga}"
+            )
+
+        # Per-expert weight amax (weights identical -> tight match), all experts calibrated.
+        for idx in range(NUM_EXPERTS):
+            for q_list in ("gate_up_proj_weight_quantizers", "down_proj_weight_quantizers"):
+                eq = getattr(e, q_list)[idx].amax
+                gq = getattr(g, q_list)[idx].amax
+                assert eq is not None and gq is not None, f"{q_list}[{idx}] not calibrated"
+                assert torch.allclose(eq, gq, atol=1e-5), f"{q_list}[{idx}] weight amax mismatch"
+        self._cleanup_registry(_SyntheticGroupedExperts)
+
+    def test_fake_quantizing_a_projection_keeps_one_experts_slice_alive(self):
+        """Without autograd each expert's fake-quantized slice goes into the output as it is made:
+        a projection's slices all alive beside the output would double its memory."""
+        model = _GroupedMoEModel("grouped_mm")
+        try:
+            mtq.quantize(model, self._quant_cfg(), self._forward_loop_all_experts)
+        finally:
+            self._cleanup_registry(_SyntheticGroupedExperts)
+        experts = model.moe.experts
+        weight, quantizers = experts.gate_up_proj, experts.gate_up_proj_weight_quantizers
+        made = []
+
+        def one_alive(_quantizer, _args, output):
+            assert all(ref() is None for ref in made), "an earlier expert's slice is still alive"
+            made.append(weakref.ref(output))
+
+        handles = [q.register_forward_hook(one_alive) for q in quantizers]
+        try:
+            with torch.no_grad():
+                out = experts._quantize_grouped_weight(weight, quantizers)
+        finally:
+            for handle in handles:
+                handle.remove()
+
+        assert len(made) == NUM_EXPERTS
+        with torch.no_grad():
+            expected = torch.stack([q(weight[idx]) for idx, q in enumerate(quantizers)])
+        torch.testing.assert_close(out, expected, rtol=0, atol=0)
+
+    def test_an_eager_only_wrapper_refuses_the_grouped_forward(self):
+        """A wrapper whose weights are defined on the eager forward raises if the grouped one runs,
+        rather than quantize the stored weight."""
+
+        class _EagerOnlyExperts(_QuantFusedExperts):
+            _eager_only = True
+
+        self._cleanup_registry(_SyntheticGroupedExperts)
+        QuantModuleRegistry.register({_SyntheticGroupedExperts: "test.SyntheticGroupedExperts"})(
+            _EagerOnlyExperts
+        )
+        try:
+            model = _GroupedMoEModel("grouped_mm")
+            QuantModuleRegistry.convert(model.moe.experts)
+            with pytest.raises(RuntimeError, match="eager experts forward only"):
+                self._forward_loop_all_experts(model)
+        finally:
+            self._cleanup_registry(_SyntheticGroupedExperts)
+
+
 # ---------------------------------------------------------------------------
 # Tests for the non-gated fused-experts path (NemotronH NemotronHExperts):
 # single up_proj (no gate half) + down_proj. Quantizers are named after the
@@ -1881,7 +2188,7 @@ def _experts_type():
 
         return Qwen3VLMoeTextExperts
 
-    def test_left_to_generic_fused_experts_wrapper(self):
+    def test_left_to_generic_fused_experts_wrapper_class(self):
         """The experts must be left to the generic fused-experts wrapper."""
         # A static registration would shadow on-the-fly detection.
         assert QuantModuleRegistry.get(self._experts_type()) is None

From 54b50f58a786ff674dc3a4b5fa605d8d66e17c93 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 23:13:52 +0000
Subject: [PATCH 02/26] Leave fused experts stored transposed unquantized

_fused_experts_wrapper (renamed from _fused_experts_wrapper_class) matches
fused experts on their parameters' shapes alone, so it also claimed a
container stored transposed, (num_experts, in_dim, out_dim). The fused-experts
export splits gate/up along dim 1, the input dim of that layout, and wrote
silently wrong weights. register_fused_experts_on_the_fly now leaves such a
container unquantized with a warning; GptOss and Llama4 have their own
wrappers. The Qwen3 recipe test's synthetic experts take the standard layout.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/models/step3p5/modeling_ptq.py |  2 +-
 .../torch/quantization/plugins/huggingface.py | 51 ++++++++++++++-----
 tests/unit/recipe/test_qwen3_8_recipe.py      |  2 +-
 .../plugins/test_fused_experts.py             | 36 ++++++++++---
 4 files changed, 70 insertions(+), 21 deletions(-)

diff --git a/modelopt/torch/models/step3p5/modeling_ptq.py b/modelopt/torch/models/step3p5/modeling_ptq.py
index 9e3d40d49b8..85860affe17 100644
--- a/modelopt/torch/models/step3p5/modeling_ptq.py
+++ b/modelopt/torch/models/step3p5/modeling_ptq.py
@@ -137,7 +137,7 @@ def _is_expert_indexed_moe_linear(module: nn.Module) -> bool:
     ``weight`` of shape ``[num_experts, out_features, in_features]``, whose
     ``forward(x, expert_id)`` runs ``F.linear`` against the selected expert's slice. The
     weights therefore live on the projection submodule rather than on the expert container,
-    which is what :func:`_fused_experts_wrapper_class` looks for, and the module is not an
+    which is what :func:`_fused_experts_wrapper` looks for, and the module is not an
     ``nn.Linear``, so neither the fused-experts path nor the plain linear path claims it.
 
     Detection is structural rather than keyed on class names so new Step revisions are picked
diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index ed11c8d5bd5..6cad7693bc5 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -1461,26 +1461,33 @@ def register_sparse_moe_on_the_fly(model):
             )
 
 
-def _fused_experts_wrapper_class(module):
-    """Return the _QuantFusedExperts subclass for a fused MoE expert container, or None.
+def _fused_experts_wrapper(module):
+    """Return the ``_QuantFusedExperts`` subclass for a fused-MoE-experts module, else None.
 
-    Two 3-D fused layouts are recognized, both requiring ``num_experts`` and a
-    3-D ``down_proj`` parameter:
+    Detects the standardized HuggingFace transformers 5.0+ fused expert pattern: ``num_experts``,
+    a 3-D ``down_proj`` and a 3-D first projection, either
 
-    * gated (``_QuantFusedExperts``): a 3-D ``gate_up_proj`` fusing gate+up. Matches
+    * gated (:class:`_QuantFusedExperts`): a 3-D ``gate_up_proj`` fusing gate+up. Matches
       ``MixtralExperts``, ``Qwen2MoeExperts``, ``Qwen3MoeExperts``,
       ``Qwen3_5MoeExperts``, ``DeepseekV3NaiveMoe``, ``JambaExperts``,
       ``OlmoeExperts``, ``MiniMaxM2Experts``, ``MiniMaxM3VLExperts``, etc.
     * non-gated (``_QuantNonGatedFusedExperts``): a 3-D ``up_proj`` with no
       ``gate_proj`` and no ``gate_up_proj``. Matches NemotronH ``NemotronHExperts``.
 
-    Returns ``None`` for non-standard layouts (GptOss, GraniteMoE,
-    Llama4TextExperts) which have their own explicit registrations.
+    Matching is by shape alone, so GptOss and Llama4TextExperts -- transposed layouts with their own
+    explicit registrations -- match too; :func:`register_fused_experts_on_the_fly` never reaches
+    them while those registrations stand.
 
     ``act_fn`` is not required: these wrappers only intercept the two ``F.linear``
     calls, so modules with a custom gated activation (e.g. ``MiniMaxM3VLExperts``)
     are still supported.
     """
+    # ``act_fn`` is a common marker of the standard HF fused-experts pattern, but it is NOT required by
+    # _QuantFusedExperts: the wrapper only quantizes the two projections; the activation is applied by the
+    # module's OWN forward. Models with a custom activation applied inline -- e.g. MiniMax-M3's SwiGLU via
+    # ``_apply_gate`` (``swiglu_alpha``/``swiglu_limit``, no ``act_fn``) -- are otherwise standard fused
+    # experts, so gate on the STRUCTURAL signature (num_experts + two 3-D fused projections) and treat
+    # ``act_fn`` as optional.
     if not hasattr(module, "num_experts"):
         return None
     down = getattr(module, "down_proj", None)
@@ -1500,10 +1507,10 @@ def _fused_experts_wrapper_class(module):
 def _is_fused_experts_module(module):
     """Check if a module is a fused MoE expert container compatible with _QuantFusedExperts.
 
-    See :func:`_fused_experts_wrapper_class` for the recognized layouts (gated
+    See :func:`_fused_experts_wrapper` for the recognized layouts (gated
     ``gate_up_proj`` and non-gated ``up_proj``).
     """
-    return _fused_experts_wrapper_class(module) is not None
+    return _fused_experts_wrapper(module) is not None
 
 
 def register_fused_experts_on_the_fly(model):
@@ -1525,13 +1532,31 @@ def register_fused_experts_on_the_fly(model):
 
         visited_types.add(mod_type)
 
-        wrapper_cls = _fused_experts_wrapper_class(module)
-        if wrapper_cls is not None:
+        wrapper = _fused_experts_wrapper(module)
+        if wrapper is not None and (
+            getattr(module, "is_transposed", False)
+            # Stored `(num_experts, out_dim, in_dim)`, the first projection takes the dim down_proj
+            # puts out; a gated container stored transposed never does, whether it says so or not.
+            or getattr(module, wrapper._first_proj_attr).shape[-1] != module.down_proj.shape[1]
+        ):
+            # `_fused_experts_wrapper` matches on parameter SHAPE only, so it also claims a
+            # transposed `(num_experts, in_dim, out_dim)` container. The fused-experts export
+            # (export/moe_utils.py) then splits gate/up by slicing dim 1 at `intermediate_dim`, the
+            # wrong axis for that layout, and writes silently wrong weights. Leave it unquantized.
+            warnings.warn(
+                f"Fused MoE experts '{name}' ({mod_type.__name__}) store their weights transposed "
+                "(num_experts, in_dim, out_dim), which the generic _QuantFusedExperts export path "
+                "does not support: it splits gate/up along dim 1, which is the input dim for this "
+                "layout. Leaving them unquantized; such a class needs its own quantized wrapper, as "
+                "GptOssExperts and Llama4TextExperts have."
+            )
+            continue
+        if wrapper is not None:
             print(
                 f"\033[1mDetected fused MoE experts '{name}' of type {mod_type.__name__}, "
-                f"registering with {wrapper_cls.__name__}.\033[0m"
+                f"registering with {wrapper.__name__}.\033[0m"
             )
-            QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(wrapper_cls)
+            QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(wrapper)
 
 
 def force_eager_experts_impl_on_the_fly(model):
diff --git a/tests/unit/recipe/test_qwen3_8_recipe.py b/tests/unit/recipe/test_qwen3_8_recipe.py
index 36abe30e14d..576406c508b 100644
--- a/tests/unit/recipe/test_qwen3_8_recipe.py
+++ b/tests/unit/recipe/test_qwen3_8_recipe.py
@@ -45,7 +45,7 @@ def __init__(self):
         super().__init__()
         self.num_experts = 2
         self.intermediate_dim = _H
-        self.gate_up_proj = nn.Parameter(torch.randn(2, _H, 2 * _H))
+        self.gate_up_proj = nn.Parameter(torch.randn(2, 2 * _H, _H))
         self.down_proj = nn.Parameter(torch.randn(2, _H, _H))
 
     def forward(self, hidden_states):
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 7f0b13c397d..fd4102e490e 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -46,7 +46,7 @@
 from modelopt.torch.quantization.model_calib import local_hessian_calibrate
 from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer
 from modelopt.torch.quantization.plugins.huggingface import (
-    _fused_experts_wrapper_class,
+    _fused_experts_wrapper,
     _is_fused_experts_module,
     _is_sparse_sequaential_moe_block,
     _QuantFusedExperts,
@@ -472,6 +472,30 @@ def test_fused_experts_only_registration(self):
         assert QuantModuleRegistry.get(block_type) is None
         self._cleanup_registry(expert_type)
 
+    @pytest.mark.parametrize("declared", [True, False], ids=["declared", "by-shape"])
+    def test_transposed_fused_experts_are_left_unquantized(self, declared):
+        """Matched by shape, a transposed container would export with gate/up split along its input
+        dim; registration leaves it unquantized, with a warning, rather than let export write wrong
+        weights or abort the rest of the model's quantization. A gated container's shapes give it
+        away when it does not say so."""
+
+        class _TransposedExperts(_SyntheticFusedExperts):
+            pass
+
+        experts = _TransposedExperts()
+        if declared:
+            experts.is_transposed = True
+        else:
+            for name in ("gate_up_proj", "down_proj"):
+                weight = getattr(experts, name).detach().transpose(1, 2).contiguous()
+                setattr(experts, name, nn.Parameter(weight))
+        model = nn.Module()
+        model.experts = experts
+
+        with pytest.warns(UserWarning, match="store their weights transposed"):
+            register_fused_experts_on_the_fly(model)
+        assert QuantModuleRegistry.get(_TransposedExperts) is None
+
     def test_a_wrapper_that_defines_its_weights_in_get_expert_weight_is_eager_only(self):
         """The grouped_mm forward quantizes the stored weight without get_expert_weight, so a wrapper
         that defines its weight domain there must calibrate on the eager forward. One that does not
@@ -1916,10 +1940,10 @@ def _cleanup_registry(mod_type):
     def test_detected_and_picks_nongated_wrapper(self):
         module = _SyntheticNonGatedFusedExperts()
         assert _is_fused_experts_module(module) is True
-        assert _fused_experts_wrapper_class(module) is _QuantNonGatedFusedExperts
+        assert _fused_experts_wrapper(module) is _QuantNonGatedFusedExperts
 
     def test_gated_still_picks_base_wrapper(self):
-        assert _fused_experts_wrapper_class(_SyntheticFusedExperts()) is _QuantFusedExperts
+        assert _fused_experts_wrapper(_SyntheticFusedExperts()) is _QuantFusedExperts
 
     def test_register_uses_nongated_wrapper(self):
         model = _TinyNonGatedMoEModel()
@@ -2107,7 +2131,7 @@ def __init__(self):
                 self.act_fn = nn.SiLU()
 
         module = _SplitGatedExperts()
-        assert _fused_experts_wrapper_class(module) is None
+        assert _fused_experts_wrapper(module) is None
         assert _is_fused_experts_module(module) is False
 
     def test_get_quant_config_resolves_nongated_experts(self):
@@ -2188,11 +2212,11 @@ def _experts_type():
 
         return Qwen3VLMoeTextExperts
 
-    def test_left_to_generic_fused_experts_wrapper_class(self):
+    def test_left_to_generic_fused_experts_wrapper(self):
         """The experts must be left to the generic fused-experts wrapper."""
         # A static registration would shadow on-the-fly detection.
         assert QuantModuleRegistry.get(self._experts_type()) is None
-        assert _fused_experts_wrapper_class(_make_qwen3_vl_moe_experts()) is _QuantFusedExperts
+        assert _fused_experts_wrapper(_make_qwen3_vl_moe_experts()) is _QuantFusedExperts
 
     def test_convert_and_forward_matches_reference(self):
         """Conversion must succeed and stay numerically transparent before calibration."""

From 2d287cf1715378a60484e510b7ee604175c7ce83 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 00:07:31 +0000
Subject: [PATCH 03/26] Keep a disabled weight quantizer's pre-quant scale and
 rotation on the grouped forward

TensorQuantizer applies its pre-quant scale and rotation before it checks
whether it is disabled, so the eager forward still transforms a disabled
expert's weight. The grouped forward's passthrough shortcut returned the stored
weight whenever every quantizer was disabled; it now does only when every one
would pass its slice through unchanged.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/quantization/plugins/huggingface.py | 13 ++++++---
 .../plugins/test_fused_experts.py             | 27 +++++++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 6cad7693bc5..d56c80fde3f 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -897,10 +897,17 @@ def _quantize_grouped_weight(self, weight: torch.Tensor, quantizers: nn.ModuleLi
 
         Each expert slice passes through its own weight quantizer (collecting
         ``amax`` during calibration, fake-quantizing afterwards), matching the
-        eager per-expert ``F.linear`` semantics. Returns the weight untouched when
-        every quantizer is disabled so passthrough stays exact (and cheap).
+        eager per-expert ``F.linear`` semantics. Returns the weight untouched when every
+        quantizer would pass its slice through unchanged, so passthrough stays exact (and cheap).
         """
-        if all(not q.is_enabled for q in quantizers):
+        # A disabled quantizer still applies its pre-quant scale and rotation, as on the eager path.
+        if all(
+            isinstance(q, TensorQuantizer)
+            and not q.is_enabled
+            and q.pre_quant_scale is None
+            and not q.rotate_is_enabled
+            for q in quantizers
+        ):
             return weight
         if torch.is_grad_enabled():
             # Slices written into one buffer would chain a full-size gradient copy per expert into
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index fd4102e490e..5f0dab39346 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -1906,6 +1906,33 @@ def one_alive(_quantizer, _args, output):
             expected = torch.stack([q(weight[idx]) for idx, q in enumerate(quantizers)])
         torch.testing.assert_close(out, expected, rtol=0, atol=0)
 
+    def test_disabled_quantizers_keep_their_pre_quant_scale(self):
+        """A disabled weight quantizer still applies its pre-quant scale, as on the eager path, so the
+        grouped forward does not pass the stored weight through unscaled."""
+        self._cleanup_registry(_SyntheticGroupedExperts)
+        torch.manual_seed(0)
+        models = [_GroupedMoEModel("eager"), _GroupedMoEModel("grouped_mm")]
+        models[1].load_state_dict(models[0].state_dict())
+        scale = torch.rand(HIDDEN_DIM) + 0.5
+        hidden = torch.randn(8, HIDDEN_DIM)
+        top_k_index = torch.randint(0, NUM_EXPERTS, (8, TOP_K))
+        top_k_weights = torch.softmax(torch.randn(8, TOP_K), dim=1)
+        outputs = []
+        try:
+            register_fused_experts_on_the_fly(models[0])
+            for model in models:
+                experts = QuantModuleRegistry.convert(model.moe.experts)
+                for q in experts.modules():
+                    if isinstance(q, TensorQuantizer):
+                        q.disable()
+                for q in experts.gate_up_proj_weight_quantizers:
+                    q.pre_quant_scale = scale
+                with torch.no_grad():
+                    outputs.append(experts(hidden, top_k_index, top_k_weights))
+        finally:
+            self._cleanup_registry(_SyntheticGroupedExperts)
+        torch.testing.assert_close(outputs[1], outputs[0], rtol=1e-4, atol=1e-6)
+
     def test_an_eager_only_wrapper_refuses_the_grouped_forward(self):
         """A wrapper whose weights are defined on the eager forward raises if the grouped one runs,
         rather than quantize the stored weight."""

From 144db5a4d7369e8cc8a0d12dec4680b2f9f72c73 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 00:07:52 +0000
Subject: [PATCH 04/26] Pin the grouped forward's fake-quantized output and
 gradients against eager

The autograd branch of _quantize_grouped_weight, which stacks every expert's
fake-quantized slice, ran in no grouped forward or backward. On the same
calibrated quantizers, the grouped forward's output and its gradients to the
input and to both fused weights now match the eager forward's.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../plugins/test_fused_experts.py             | 29 +++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 5f0dab39346..b1b67fd0d67 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -1933,6 +1933,35 @@ def test_disabled_quantizers_keep_their_pre_quant_scale(self):
             self._cleanup_registry(_SyntheticGroupedExperts)
         torch.testing.assert_close(outputs[1], outputs[0], rtol=1e-4, atol=1e-6)
 
+    def test_fake_quantized_grouped_forward_and_backward_match_eager(self):
+        """Under autograd (QAT) the grouped forward fake-quantizes every expert's slice at once; its
+        output and gradients match the eager forward's on the same calibrated quantizers."""
+        self._cleanup_registry(_SyntheticGroupedExperts)
+        torch.manual_seed(0)
+        models = [_GroupedMoEModel("eager"), _GroupedMoEModel("grouped_mm")]
+        models[1].load_state_dict(models[0].state_dict())
+        try:
+            # quantize() calibrates every model on eager, so both get the same quantizers.
+            for model in models:
+                mtq.quantize(model, self._quant_cfg(), self._forward_loop_all_experts)
+            models[1].config._experts_implementation = "grouped_mm"
+        finally:
+            self._cleanup_registry(_SyntheticGroupedExperts)
+        hidden = torch.randn(16, HIDDEN_DIM)
+        top_k_index = torch.stack(
+            [torch.arange(16) % NUM_EXPERTS, (torch.arange(16) + 1) % NUM_EXPERTS], 1
+        )
+        top_k_weights = torch.softmax(torch.randn(16, TOP_K), dim=1)
+        results = []
+        for model in models:
+            x = hidden.clone().requires_grad_()
+            out = model.moe.experts(x, top_k_index, top_k_weights)
+            out.square().sum().backward()
+            experts = model.moe.experts
+            results.append((out, x.grad, experts.gate_up_proj.grad, experts.down_proj.grad))
+        for grouped, eager in zip(results[1], results[0]):
+            torch.testing.assert_close(grouped, eager, rtol=1e-4, atol=1e-6)
+
     def test_an_eager_only_wrapper_refuses_the_grouped_forward(self):
         """A wrapper whose weights are defined on the eager forward raises if the grouped one runs,
         rather than quantize the stored weight."""

From 9122ae1b31ee280ac60999fa414b752e93b1bd62 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 00:08:45 +0000
Subject: [PATCH 05/26] Unregister the grouped test experts in a finally, as
 the class's other tests do

QuantModuleRegistry is process-global: a failing assertion in either test left
the synthetic experts registered for the rest of the session.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../plugins/test_fused_experts.py             | 90 ++++++++++---------
 1 file changed, 48 insertions(+), 42 deletions(-)

diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index b1b67fd0d67..874d59e1bf3 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -1817,23 +1817,25 @@ def test_forward_passthrough_matches(self):
         ref = _SyntheticGroupedExperts(_ExpertsConfig("grouped_mm"))
         ref.load_state_dict(model.moe.experts.state_dict())
 
-        register_fused_experts_on_the_fly(model)
-        converted = QuantModuleRegistry.convert(model.moe.experts)
-        for q in converted.modules():
-            if isinstance(q, TensorQuantizer):
-                q.disable()
+        try:
+            register_fused_experts_on_the_fly(model)
+            converted = QuantModuleRegistry.convert(model.moe.experts)
+            for q in converted.modules():
+                if isinstance(q, TensorQuantizer):
+                    q.disable()
 
-        seq = 8
-        hidden = torch.randn(seq, HIDDEN_DIM)
-        top_k_index = torch.randint(0, NUM_EXPERTS, (seq, TOP_K))
-        top_k_weights = torch.softmax(torch.randn(seq, TOP_K), dim=1)
-        with torch.no_grad():
-            out_ref = ref(hidden, top_k_index, top_k_weights)
-            out_test = converted(hidden, top_k_index, top_k_weights)
-        assert torch.allclose(out_ref, out_test, atol=1e-6), (
-            f"Max diff: {(out_ref - out_test).abs().max().item()}"
-        )
-        self._cleanup_registry(_SyntheticGroupedExperts)
+            seq = 8
+            hidden = torch.randn(seq, HIDDEN_DIM)
+            top_k_index = torch.randint(0, NUM_EXPERTS, (seq, TOP_K))
+            top_k_weights = torch.softmax(torch.randn(seq, TOP_K), dim=1)
+            with torch.no_grad():
+                out_ref = ref(hidden, top_k_index, top_k_weights)
+                out_test = converted(hidden, top_k_index, top_k_weights)
+            assert torch.allclose(out_ref, out_test, atol=1e-6), (
+                f"Max diff: {(out_ref - out_test).abs().max().item()}"
+            )
+        finally:
+            self._cleanup_registry(_SyntheticGroupedExperts)
 
     def test_calibration_matches_eager(self):
         """grouped_mm PTQ calibration yields the same quantizer state as eager PTQ.
@@ -1849,33 +1851,37 @@ def test_calibration_matches_eager(self):
         grouped_model = _GroupedMoEModel("grouped_mm")
         grouped_model.load_state_dict(eager_model.state_dict())
 
-        for model in (eager_model, grouped_model):
-            impl = model.config._experts_implementation
-            # quantize() calibrates every model on eager; calibrate on the backend under test.
-            mtq.quantize(model, self._quant_cfg())
-            model.config._experts_implementation = impl
-            mtq.calibrate(model, "max", self._forward_loop_all_experts)
-
-        e = eager_model.moe.experts
-        g = grouped_model.moe.experts
-
-        # Shared input quantizers (down flows through a GEMM -> small tolerance).
-        for q_name in ("gate_up_proj_input_quantizer", "down_proj_input_quantizer"):
-            ea = getattr(e, q_name).amax
-            ga = getattr(g, q_name).amax
-            assert ea is not None and ga is not None, f"{q_name} not calibrated"
-            assert torch.allclose(ea, ga, atol=1e-3, rtol=1e-2), (
-                f"{q_name} amax mismatch: eager={ea} grouped={ga}"
-            )
+        try:
+            for model in (eager_model, grouped_model):
+                impl = model.config._experts_implementation
+                # quantize() calibrates every model on eager; calibrate on the backend under test.
+                mtq.quantize(model, self._quant_cfg())
+                model.config._experts_implementation = impl
+                mtq.calibrate(model, "max", self._forward_loop_all_experts)
+
+            e = eager_model.moe.experts
+            g = grouped_model.moe.experts
+
+            # Shared input quantizers (down flows through a GEMM -> small tolerance).
+            for q_name in ("gate_up_proj_input_quantizer", "down_proj_input_quantizer"):
+                ea = getattr(e, q_name).amax
+                ga = getattr(g, q_name).amax
+                assert ea is not None and ga is not None, f"{q_name} not calibrated"
+                assert torch.allclose(ea, ga, atol=1e-3, rtol=1e-2), (
+                    f"{q_name} amax mismatch: eager={ea} grouped={ga}"
+                )
 
-        # Per-expert weight amax (weights identical -> tight match), all experts calibrated.
-        for idx in range(NUM_EXPERTS):
-            for q_list in ("gate_up_proj_weight_quantizers", "down_proj_weight_quantizers"):
-                eq = getattr(e, q_list)[idx].amax
-                gq = getattr(g, q_list)[idx].amax
-                assert eq is not None and gq is not None, f"{q_list}[{idx}] not calibrated"
-                assert torch.allclose(eq, gq, atol=1e-5), f"{q_list}[{idx}] weight amax mismatch"
-        self._cleanup_registry(_SyntheticGroupedExperts)
+            # Per-expert weight amax (weights identical -> tight match), all experts calibrated.
+            for idx in range(NUM_EXPERTS):
+                for q_list in ("gate_up_proj_weight_quantizers", "down_proj_weight_quantizers"):
+                    eq = getattr(e, q_list)[idx].amax
+                    gq = getattr(g, q_list)[idx].amax
+                    assert eq is not None and gq is not None, f"{q_list}[{idx}] not calibrated"
+                    assert torch.allclose(eq, gq, atol=1e-5), (
+                        f"{q_list}[{idx}] weight amax mismatch"
+                    )
+        finally:
+            self._cleanup_registry(_SyntheticGroupedExperts)
 
     def test_fake_quantizing_a_projection_keeps_one_experts_slice_alive(self):
         """Without autograd each expert's fake-quantized slice goes into the output as it is made:

From 6372ed9b2f824d8162a55ab5c683284c89716694 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 00:09:36 +0000
Subject: [PATCH 06/26] Leave the torch_a2a test fixtures to the PR whose tests
 drive them

No test in this PR runs the synthetic experts on torch_a2a; the expert-parallel
tests, and the dispatch branch, ep_group field and import they need, come with
the calibration change stacked on it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/unit/torch/quantization/plugins/test_fused_experts.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 874d59e1bf3..3f781a2fbe4 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -56,7 +56,6 @@
     register_sparse_moe_on_the_fly,
 )
 from modelopt.torch.quantization.utils.core_utils import weight_attr_names
-from modelopt.torch.utils.moe import torch_a2a_experts_forward
 
 
 def test_native_fp8_experts_dequantized_before_nvfp4_export():
@@ -309,7 +308,6 @@ class _ExpertsConfig:
     def __init__(self, impl="grouped_mm"):
         self._experts_implementation = impl
         self.hidden_act = "silu"
-        self.ep_group = None  # the expert-parallel group a torch_a2a forward dispatches over
 
 
 class _SyntheticGroupedExperts(nn.Module):
@@ -369,10 +367,6 @@ def forward(self, hidden_states, top_k_index, top_k_weights):
             return _moe_integration.grouped_mm_experts_forward(
                 self, hidden_states, top_k_index, top_k_weights
             )
-        if self.config._experts_implementation == "torch_a2a":
-            return torch_a2a_experts_forward(
-                self, hidden_states, top_k_index, top_k_weights, self.config.ep_group
-            )
         return self._eager_forward(hidden_states, top_k_index, top_k_weights)
 
 

From 7c8390953d2156c9757d2bb82a67e5513f3aec7a Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 00:47:07 +0000
Subject: [PATCH 07/26] Fake-quantize the grouped forward's experts in one
 launch where the format allows

Under autograd the eager forward fake-quantizes a projection's experts with
_try_stacked_fake_quant, one kernel launch for dynamic block-16 NVFP4; the
grouped forward quantized them one expert at a time. It now takes the same
stacked kernel first, so QAT keeps that speed whichever backend runs. The GPU
test that pins the stacked kernel's equivalence now runs on both backends.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/quantization/plugins/huggingface.py           | 4 ++++
 .../torch/quantization/plugins/test_fused_experts_cuda.py    | 5 ++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index d56c80fde3f..70ec694eb16 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -910,6 +910,10 @@ def _quantize_grouped_weight(self, weight: torch.Tensor, quantizers: nn.ModuleLi
         ):
             return weight
         if torch.is_grad_enabled():
+            # One kernel launch for every expert where the format has one, as on the eager path.
+            stacked = _try_stacked_fake_quant(weight, quantizers)
+            if stacked is not None:
+                return stacked
             # Slices written into one buffer would chain a full-size gradient copy per expert into
             # the backward (QAT), so stack them there.
             return torch.stack([quantizers[idx](weight[idx]) for idx in range(self.num_experts)])
diff --git a/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py b/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py
index a0669da1c1c..da7c96c7b97 100644
--- a/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py
+++ b/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py
@@ -29,11 +29,14 @@
 )
 
 
+@pytest.mark.parametrize("impl", ["eager", "grouped_mm"])
 @pytest.mark.parametrize("pass_through_bwd", [True, False])
-def test_stacked_nvfp4_experts_match_per_expert_quantizers(monkeypatch, pass_through_bwd):
+def test_stacked_nvfp4_experts_match_per_expert_quantizers(monkeypatch, pass_through_bwd, impl):
     model = get_tiny_qwen3_moe().cuda().train()
     input_ids = torch.randint(0, model.config.vocab_size, (2, 16), device="cuda")
     mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(input_ids))
+    # Train on the backend under test; both fake-quantize each projection's experts in one launch.
+    model.config._experts_implementation = impl
     experts = [m for m in model.modules() if isinstance(m, hf_plugin._QuantFusedExperts)]
     assert experts
     for module in experts:

From 8c84ef910c9f04243a3d44b2b1c8ade413673656 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 01:17:52 +0000
Subject: [PATCH 08/26] Let the eager-keyed weight-input captures skip the
 grouped forward's batch

local Hessian, imatrix and GPTQ capture an expert's inputs on its shared input
quantizer, keyed by _current_expert_idx, which only the eager loop sets. The
grouped forward runs the same quantizer on every expert's rows at once, so a
capture credited the whole batch to whichever expert the index last named. The
grouped interceptor now clears the index: the captures skip that batch, as they
did before it was intercepted. Splitting it per expert comes with the grouped
calibration in the next PR.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/quantization/plugins/huggingface.py |  4 +++
 .../plugins/test_fused_experts.py             | 30 +++++++++++++++++--
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 70ec694eb16..d893022b314 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -973,6 +973,10 @@ def _quantized_grouped_linear(input, weight, offs, bias=None, is_transposed=Fals
                         f"{type(self).__name__} quantizes on the eager experts forward only, but its "
                         'grouped_mm forward ran. Load the model with experts_implementation="eager".'
                     )
+                # The batch holds every expert's rows: no single expert for the eager loop's index to
+                # name, so a weight-input capture keyed on it skips the batch rather than credit
+                # one expert with all of it.
+                self._current_expert_idx = None
                 quantizers = (
                     self.down_proj_weight_quantizers if self._down_proj_linear else first_weight_qs
                 )
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 3f781a2fbe4..d863978a41d 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -43,7 +43,10 @@
 )
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name
-from modelopt.torch.quantization.model_calib import local_hessian_calibrate
+from modelopt.torch.quantization.model_calib import (
+    _register_weight_input_hooks,
+    local_hessian_calibrate,
+)
 from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer
 from modelopt.torch.quantization.plugins.huggingface import (
     _fused_experts_wrapper,
@@ -55,7 +58,7 @@
     register_fused_experts_on_the_fly,
     register_sparse_moe_on_the_fly,
 )
-from modelopt.torch.quantization.utils.core_utils import weight_attr_names
+from modelopt.torch.quantization.utils.core_utils import module_name_maps, weight_attr_names
 
 
 def test_native_fp8_experts_dequantized_before_nvfp4_export():
@@ -1962,6 +1965,29 @@ def test_fake_quantized_grouped_forward_and_backward_match_eager(self):
         for grouped, eager in zip(results[1], results[0]):
             torch.testing.assert_close(grouped, eager, rtol=1e-4, atol=1e-6)
 
+    def test_weight_input_captures_skip_the_grouped_forward(self):
+        """The weight-input captures key on the eager loop's expert index, which names no expert in the
+        grouped forward's batch of every expert's rows: they skip it rather than credit one expert
+        with all of it."""
+        model = _GroupedMoEModel("eager")
+        rows = {}
+
+        def capture(quantizer, _weight, inputs):
+            rows[id(quantizer)] = rows.get(id(quantizer), 0) + inputs.shape[0]
+
+        try:
+            mtq.quantize(model, self._quant_cfg(), self._forward_loop_all_experts)
+            handles = _register_weight_input_hooks(model, module_name_maps(model), capture)
+            model.config._experts_implementation = "grouped_mm"
+            try:
+                self._forward_loop_all_experts(model)
+            finally:
+                for handle in handles:
+                    handle.remove()
+        finally:
+            self._cleanup_registry(_SyntheticGroupedExperts)
+        assert rows == {}
+
     def test_an_eager_only_wrapper_refuses_the_grouped_forward(self):
         """A wrapper whose weights are defined on the eager forward raises if the grouped one runs,
         rather than quantize the stored weight."""

From c7fb068e5fb6bcfe44fb9a5eda7fa79efe6e1cdd Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 01:19:49 +0000
Subject: [PATCH 09/26] Take the stacked fake-quant kernel on the grouped
 forward without autograd too

The eager forward skips it without autograd because it then quantizes only the
routed experts; the grouped forward quantizes every expert either way, and the
kernel writes one output of the weight's shape, the buffer the slice-by-slice
path fills. The GPU test now also checks that the backend under test ran, and
that the no-autograd forward takes the kernel and matches the per-expert path.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/quantization/plugins/huggingface.py    | 10 ++++++----
 .../plugins/test_fused_experts_cuda.py           | 16 ++++++++++++++++
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index d893022b314..6df940f0b19 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -909,11 +909,13 @@ def _quantize_grouped_weight(self, weight: torch.Tensor, quantizers: nn.ModuleLi
             for q in quantizers
         ):
             return weight
+        # One kernel launch for every expert where the format has one. Unlike the eager path, which
+        # without autograd quantizes only the routed experts, this one quantizes every expert either
+        # way, and the kernel's one output is the buffer the slice-by-slice path below fills.
+        stacked = _try_stacked_fake_quant(weight, quantizers)
+        if stacked is not None:
+            return stacked
         if torch.is_grad_enabled():
-            # One kernel launch for every expert where the format has one, as on the eager path.
-            stacked = _try_stacked_fake_quant(weight, quantizers)
-            if stacked is not None:
-                return stacked
             # Slices written into one buffer would chain a full-size gradient copy per expert into
             # the backward (QAT), so stack them there.
             return torch.stack([quantizers[idx](weight[idx]) for idx in range(self.num_experts)])
diff --git a/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py b/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py
index da7c96c7b97..2d447a163d9 100644
--- a/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py
+++ b/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py
@@ -18,6 +18,7 @@
 import pytest
 import torch
 from _test_utils.torch.transformers_models import get_tiny_qwen3_moe
+from transformers.integrations import moe as moe_integration
 
 import modelopt.torch.quantization as mtq
 import modelopt.torch.quantization.plugins.huggingface as hf_plugin
@@ -65,14 +66,29 @@ def loss_and_grads(checkpointing):
         "stacked_nvfp4_fake_quant",
         lambda *args: calls.append(args) or stacked_fake_quant(*args),
     )
+    grouped_calls = []
+    grouped_linear = moe_integration._grouped_linear
+    monkeypatch.setattr(
+        moe_integration,
+        "_grouped_linear",
+        lambda *args, **kwargs: grouped_calls.append(1) or grouped_linear(*args, **kwargs),
+    )
     results = []
     for checkpointing in (False, True):
         calls.clear()
         results.append(loss_and_grads(checkpointing))
         # Checkpointing reruns the forward in backward.
         assert len(calls) == 2 * len(experts) * (1 + checkpointing)
+    assert bool(grouped_calls) == (impl == "grouped_mm"), "the backend under test did not run"
+    # Without autograd only the grouped forward, which quantizes every expert, takes the kernel.
+    calls.clear()
+    with torch.no_grad():
+        logits = model(input_ids).logits
+    assert len(calls) == (2 * len(experts) if impl == "grouped_mm" else 0)
 
     monkeypatch.setattr(hf_plugin, "_try_stacked_fake_quant", lambda *args: None)
+    with torch.no_grad():
+        assert torch.equal(model(input_ids).logits, logits)
     ref_loss, ref_grads = loss_and_grads(checkpointing=False)
     for loss, grads in results:
         assert torch.equal(loss, ref_loss)

From 123e5f70d623b8f7831110fe1939ecfb5f5f7817 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 01:20:26 +0000
Subject: [PATCH 10/26] Name expert parallelism in the eager-only wrapper's
 grouped-forward refusal

The torch_a2a expert-parallel forward runs the grouped linear too, so loading
with experts_implementation="eager" alone does not avoid it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 CHANGELOG.rst                                      | 1 +
 modelopt/torch/quantization/plugins/huggingface.py | 3 ++-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index f299e566015..2cf6430585c 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -61,6 +61,7 @@ Changelog
 **Bug Fixes**
 
 - ``hf_quant_config.json`` now declares ``pre_quant_scale`` only for INT4 / W4A8 layers that have the tensor; loading INT4 / W4A8 fused-MoE checkpoints in TensorRT-LLM also needs a release whose MoE loader treats it as optional. AWQ, SmoothQuant and SVDQuant now warn when they skip quantized fused-MoE experts.
+- Hugging Face fused MoE experts stored transposed, ``(num_experts, in_dim, out_dim)``, without a quantized wrapper of their own (for example Aria's) are now left unquantized with a warning. They were quantized with gate and up split along the wrong axis, so their exported weights were wrong; the Hugging Face export now refuses such a model's MoE block instead.
 - Calibrating a Megatron-Core MLA model now gives ``linear_q_down_proj`` and ``linear_kv_down_proj`` one per-tensor input and weight range, as vLLM and TensorRT-LLM fuse them into one GEMM with one scale per tensor. Re-calibrate existing MLA checkpoints to get the shared ranges.
 - ``export_mcore_gpt_to_hf`` and ``capture_kv_cache_export_state`` now write NVFP4 KV-cache ``k_scale`` / ``v_scale`` as amax / 448, like the HF exporter and as vLLM and TensorRT-LLM expect. They wrote amax / (6 * 448), which overflows in the attention kernels. Re-export Megatron checkpoints with an NVFP4 KV cache.
 - Preserve accumulated histogram counts when computing an entropy calibration threshold with NumPy histograms.
diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 6df940f0b19..c44bc1e05fd 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -973,7 +973,8 @@ def _quantized_grouped_linear(input, weight, offs, bias=None, is_transposed=Fals
                     # this wrapper's weights.
                     raise RuntimeError(
                         f"{type(self).__name__} quantizes on the eager experts forward only, but its "
-                        'grouped_mm forward ran. Load the model with experts_implementation="eager".'
+                        'grouped_mm forward ran. Load the model with experts_implementation="eager", '
+                        "and without expert parallelism, whose forward runs it too."
                     )
                 # The batch holds every expert's rows: no single expert for the eager loop's index to
                 # name, so a weight-input capture keyed on it skips the batch rather than credit

From fb41c9e8caaa6db338a246b93b87f04acabc58d7 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 07:17:18 +0000
Subject: [PATCH 11/26] Say the grouped forward's weight-input captures skip
 its batch

The grouped interceptor clears the expert index, so on grouped_mm local
Hessian, imatrix and GPTQ calibrate the experts' weights without their
inputs, while max calibration is unaffected. The class docstring said both
backends calibrate the same quantizers without that exception.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/quantization/plugins/huggingface.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index c44bc1e05fd..4346661644d 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -814,6 +814,11 @@ class _QuantFusedExperts(_QuantFunctionalMixin):
     handled here. ``batched_mm`` (``torch.bmm``) and ``sonicmoe`` (opaque CUDA
     kernel) backends are not intercepted.
 
+    On ``grouped_mm`` the weight-input captures keyed on one expert (``local_hessian``,
+    ``imatrix``, ``gptq``) skip the batch, which holds every expert's rows, so those
+    algorithms calibrate the experts' weights without their inputs; max calibration
+    is unaffected.
+
     The non-gated variant (``up_proj`` instead of ``gate_up_proj``, used by
     NemotronH) is handled by :class:`_QuantNonGatedFusedExperts` via the
     ``_first_proj_attr`` / ``_is_gated`` hooks below; each layout names the

From 4b5bc94f26d8f7ed3a042a356074d8597de2cb27 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 07:17:28 +0000
Subject: [PATCH 12/26] Pin that a disabled weight quantizer's rotation applies
 on the grouped forward

The passthrough shortcut skips only quantizers that would leave their slice
unchanged; the rotation case fails with its rotate_is_enabled term dropped.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../plugins/test_fused_experts.py             | 20 +++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index d863978a41d..163ce030b85 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -1909,9 +1909,18 @@ def one_alive(_quantizer, _args, output):
             expected = torch.stack([q(weight[idx]) for idx, q in enumerate(quantizers)])
         torch.testing.assert_close(out, expected, rtol=0, atol=0)
 
-    def test_disabled_quantizers_keep_their_pre_quant_scale(self):
-        """A disabled weight quantizer still applies its pre-quant scale, as on the eager path, so the
-        grouped forward does not pass the stored weight through unscaled."""
+    @pytest.mark.parametrize("transform", ["pre_quant_scale", "rotate"])
+    def test_disabled_quantizers_keep_their_pre_quant_scale_and_rotation(
+        self, monkeypatch, transform
+    ):
+        """A disabled weight quantizer still applies its pre-quant scale and rotation, as on the eager
+        path, so the grouped forward does not pass the stored weight through unchanged."""
+        # A stand-in rotation that changes any weight it is applied to.
+        monkeypatch.setattr(
+            tensor_quantizer_module,
+            "normalized_hadamard_transform",
+            lambda inputs, rotate_fp32=False, block_size=None: inputs.flip(-1),
+        )
         self._cleanup_registry(_SyntheticGroupedExperts)
         torch.manual_seed(0)
         models = [_GroupedMoEModel("eager"), _GroupedMoEModel("grouped_mm")]
@@ -1929,7 +1938,10 @@ def test_disabled_quantizers_keep_their_pre_quant_scale(self):
                     if isinstance(q, TensorQuantizer):
                         q.disable()
                 for q in experts.gate_up_proj_weight_quantizers:
-                    q.pre_quant_scale = scale
+                    if transform == "rotate":
+                        q.set_from_attribute_config({"rotate": True})
+                    else:
+                        q.pre_quant_scale = scale
                 with torch.no_grad():
                     outputs.append(experts(hidden, top_k_index, top_k_weights))
         finally:

From 791ce88e8012b00f9fdc6eec358f53ce6cc74957 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 07:45:35 +0000
Subject: [PATCH 13/26] Drop a comment that repeats the docstring, and a move
 of a helper

_fused_experts_wrapper's body comment said again, at greater length, what
its docstring says about act_fn, and _is_quant_fused_experts_module had
moved above _get_fused_experts_quantizer_attr_names with nothing calling
for it. Both now read as on main.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/quantization/plugins/huggingface.py    | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 4346661644d..b01c6aced3a 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -1160,11 +1160,6 @@ def unpack_weight(self):
                 setattr(self, name, nn.Parameter(unpacked, requires_grad=False))
 
 
-def _is_quant_fused_experts_module(module):
-    """Return True for a converted HF fused-MoE-experts quantization wrapper."""
-    return isinstance(module, _QuantFusedExperts)
-
-
 def _get_fused_experts_quantizer_attr_names(module):
     """Return quantizer attribute names for a converted fused-experts module."""
     first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
@@ -1176,6 +1171,11 @@ def _get_fused_experts_quantizer_attr_names(module):
     )
 
 
+def _is_quant_fused_experts_module(module):
+    """Return True for a converted HF fused-MoE-experts quantization wrapper."""
+    return isinstance(module, _QuantFusedExperts)
+
+
 @contextmanager
 def patch_compressed_linear_loading():
     """Context manager that patches CompressedLinear to survive custom ``_init_weights`` calls.
@@ -1505,12 +1505,6 @@ def _fused_experts_wrapper(module):
     calls, so modules with a custom gated activation (e.g. ``MiniMaxM3VLExperts``)
     are still supported.
     """
-    # ``act_fn`` is a common marker of the standard HF fused-experts pattern, but it is NOT required by
-    # _QuantFusedExperts: the wrapper only quantizes the two projections; the activation is applied by the
-    # module's OWN forward. Models with a custom activation applied inline -- e.g. MiniMax-M3's SwiGLU via
-    # ``_apply_gate`` (``swiglu_alpha``/``swiglu_limit``, no ``act_fn``) -- are otherwise standard fused
-    # experts, so gate on the STRUCTURAL signature (num_experts + two 3-D fused projections) and treat
-    # ``act_fn`` as optional.
     if not hasattr(module, "num_experts"):
         return None
     down = getattr(module, "down_proj", None)

From b68e15dd92023c5d1a703eeee537e3439c20e91b Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 08:08:38 +0000
Subject: [PATCH 14/26] Say in the changelog that an old modelopt_state for
 such experts no longer restores

Restore converts the model as quantize does, which no longer registers the
transposed experts type, so a saved state's quantizer keys for those experts
match no module (ApplyModeError). The entry now tells users so.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 CHANGELOG.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 2cf6430585c..a5c54cad3fe 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -61,7 +61,7 @@ Changelog
 **Bug Fixes**
 
 - ``hf_quant_config.json`` now declares ``pre_quant_scale`` only for INT4 / W4A8 layers that have the tensor; loading INT4 / W4A8 fused-MoE checkpoints in TensorRT-LLM also needs a release whose MoE loader treats it as optional. AWQ, SmoothQuant and SVDQuant now warn when they skip quantized fused-MoE experts.
-- Hugging Face fused MoE experts stored transposed, ``(num_experts, in_dim, out_dim)``, without a quantized wrapper of their own (for example Aria's) are now left unquantized with a warning. They were quantized with gate and up split along the wrong axis, so their exported weights were wrong; the Hugging Face export now refuses such a model's MoE block instead.
+- Hugging Face fused MoE experts stored transposed, ``(num_experts, in_dim, out_dim)``, without a quantized wrapper of their own (for example Aria's) are now left unquantized with a warning. They were quantized with gate and up split along the wrong axis, so their exported weights were wrong; the Hugging Face export now refuses such a model's MoE block instead, and a ``modelopt_state`` saved for such a model by an earlier release no longer restores.
 - Calibrating a Megatron-Core MLA model now gives ``linear_q_down_proj`` and ``linear_kv_down_proj`` one per-tensor input and weight range, as vLLM and TensorRT-LLM fuse them into one GEMM with one scale per tensor. Re-calibrate existing MLA checkpoints to get the shared ranges.
 - ``export_mcore_gpt_to_hf`` and ``capture_kv_cache_export_state`` now write NVFP4 KV-cache ``k_scale`` / ``v_scale`` as amax / 448, like the HF exporter and as vLLM and TensorRT-LLM expect. They wrote amax / (6 * 448), which overflows in the attention kernels. Re-export Megatron checkpoints with an NVFP4 KV cache.
 - Preserve accumulated histogram counts when computing an entropy calibration threshold with NumPy histograms.

From d0c695d7321cd6d878a9c90eb7dd1425893c8c54 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 08:08:39 +0000
Subject: [PATCH 15/26] Import transformers' moe module in the GPU test as the
 plugin does

The test imported transformers.integrations.moe and read its private
_grouped_linear unconditionally, while the plugin and the unit tests treat
both as optional. It now skips only the grouped_mm case without them, and
spies on _grouped_linear only where it exists.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../plugins/test_fused_experts_cuda.py        | 35 ++++++++++++++-----
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py b/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py
index 2d447a163d9..156532af3cc 100644
--- a/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py
+++ b/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py
@@ -18,19 +18,37 @@
 import pytest
 import torch
 from _test_utils.torch.transformers_models import get_tiny_qwen3_moe
-from transformers.integrations import moe as moe_integration
 
 import modelopt.torch.quantization as mtq
 import modelopt.torch.quantization.plugins.huggingface as hf_plugin
 from modelopt.torch.quantization import tensor_quant
 
+# As the plugin and the unit tests treat it: grouped_mm is intercepted where transformers has it.
+try:
+    from transformers.integrations import moe as moe_integration
+
+    HAS_GROUPED_MM = hasattr(moe_integration, "_grouped_linear")
+except ImportError:
+    HAS_GROUPED_MM = False
+
 pytestmark = pytest.mark.skipif(
     not hasattr(tensor_quant.triton_kernel, "fp4_fake_quant_block"),
     reason="fp4_fake_quant_block requires compute >= 8.9",
 )
 
 
-@pytest.mark.parametrize("impl", ["eager", "grouped_mm"])
+@pytest.mark.parametrize(
+    "impl",
+    [
+        "eager",
+        pytest.param(
+            "grouped_mm",
+            marks=pytest.mark.skipif(
+                not HAS_GROUPED_MM, reason="transformers grouped_mm backend unavailable"
+            ),
+        ),
+    ],
+)
 @pytest.mark.parametrize("pass_through_bwd", [True, False])
 def test_stacked_nvfp4_experts_match_per_expert_quantizers(monkeypatch, pass_through_bwd, impl):
     model = get_tiny_qwen3_moe().cuda().train()
@@ -67,12 +85,13 @@ def loss_and_grads(checkpointing):
         lambda *args: calls.append(args) or stacked_fake_quant(*args),
     )
     grouped_calls = []
-    grouped_linear = moe_integration._grouped_linear
-    monkeypatch.setattr(
-        moe_integration,
-        "_grouped_linear",
-        lambda *args, **kwargs: grouped_calls.append(1) or grouped_linear(*args, **kwargs),
-    )
+    if HAS_GROUPED_MM:
+        grouped_linear = moe_integration._grouped_linear
+        monkeypatch.setattr(
+            moe_integration,
+            "_grouped_linear",
+            lambda *args, **kwargs: grouped_calls.append(1) or grouped_linear(*args, **kwargs),
+        )
     results = []
     for checkpointing in (False, True):
         calls.clear()

From 814f60c67802ed191964f5a12c3d180b5d225011 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 08:08:44 +0000
Subject: [PATCH 16/26] Fake-quantize the grouped weight's slices in strict
 step with its quantizers

The per-slice loops ran over num_experts while the stacked kernel goes by the
weight's rows: a weight with more rows than quantizers would have sent
unwritten rows into the grouped matmul. A strict zip makes any mismatch fail.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/quantization/plugins/huggingface.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index b01c6aced3a..3fe20972135 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -923,11 +923,12 @@ def _quantize_grouped_weight(self, weight: torch.Tensor, quantizers: nn.ModuleLi
         if torch.is_grad_enabled():
             # Slices written into one buffer would chain a full-size gradient copy per expert into
             # the backward (QAT), so stack them there.
-            return torch.stack([quantizers[idx](weight[idx]) for idx in range(self.num_experts)])
+            return torch.stack([q(w) for q, w in zip(quantizers, weight, strict=True)])
         # Each slice goes into the output as it is made: a stack would hold all of them beside it.
+        # Strict, so every row of the output is written or the call fails.
         out = torch.empty_like(weight)
-        for idx in range(self.num_experts):
-            out[idx] = quantizers[idx](weight[idx])
+        for idx, (q, w) in enumerate(zip(quantizers, weight, strict=True)):
+            out[idx] = q(w)
         return out
 
     @property

From 57e12d72916ed8faa29ce57aa696b47cf9e63cd5 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 17:25:57 +0000
Subject: [PATCH 17/26] Keep the quantizers' dtype in the grouped forward
 without autograd

Without autograd the per-slice path wrote each fake-quantized slice into a
buffer of the weight's dtype, so a pre-quant scale wider than the weights
(FP32 scales on a BF16 model) was rounded back to BF16, while the autograd
path's stack kept FP32: the same forward computed from different weights in
the two grad modes. The buffer now takes the quantizers' output dtype. The
new test runs the grouped forward both ways on BF16 weights with FP32 input
and weight scales and fails without the change.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/quantization/plugins/huggingface.py | 12 +++++--
 .../plugins/test_fused_experts.py             | 31 +++++++++++++++++++
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 3fe20972135..6bbf01ba7aa 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -925,10 +925,16 @@ def _quantize_grouped_weight(self, weight: torch.Tensor, quantizers: nn.ModuleLi
             # the backward (QAT), so stack them there.
             return torch.stack([q(w) for q, w in zip(quantizers, weight, strict=True)])
         # Each slice goes into the output as it is made: a stack would hold all of them beside it.
-        # Strict, so every row of the output is written or the call fails.
-        out = torch.empty_like(weight)
+        # Strict, so every row of the output is written or the call fails. The output keeps the
+        # quantizers' dtype, as the stack does: a pre-quant scale can widen it, and narrowing it
+        # here would compute from other weights than with autograd.
+        out: torch.Tensor | None = None
         for idx, (q, w) in enumerate(zip(quantizers, weight, strict=True)):
-            out[idx] = q(w)
+            quantized = q(w)
+            if out is None:
+                out = torch.empty_like(weight, dtype=quantized.dtype)
+            out[idx] = quantized
+            del quantized  # not alive beside the next slice
         return out
 
     @property
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 163ce030b85..1b820e9cbb4 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -1948,6 +1948,37 @@ def test_disabled_quantizers_keep_their_pre_quant_scale_and_rotation(
             self._cleanup_registry(_SyntheticGroupedExperts)
         torch.testing.assert_close(outputs[1], outputs[0], rtol=1e-4, atol=1e-6)
 
+    def test_a_widening_pre_quant_scale_gives_one_grouped_forward_in_both_grad_modes(self):
+        """FP32 pre-quant scales on a BF16 model widen what the quantizers return. The grouped forward
+        computes from those FP32 slices with or without autograd, rather than rounding the weights to
+        BF16 only without it."""
+        self._cleanup_registry(_SyntheticGroupedExperts)
+        torch.manual_seed(0)
+        model = _GroupedMoEModel("grouped_mm").to(torch.bfloat16)
+        hidden = torch.randn(8, HIDDEN_DIM, dtype=torch.bfloat16)
+        top_k_index = torch.randint(0, NUM_EXPERTS, (8, TOP_K))
+        top_k_weights = torch.softmax(torch.randn(8, TOP_K), dim=1).to(torch.bfloat16)
+        outputs = []
+        try:
+            register_fused_experts_on_the_fly(model)
+            experts = QuantModuleRegistry.convert(model.moe.experts)
+            for q in experts.modules():
+                if isinstance(q, TensorQuantizer):
+                    q.disable()
+            for name in ("gate_up_proj", "down_proj"):
+                in_dim = getattr(experts, name).shape[-1]
+                getattr(experts, f"{name}_input_quantizer").pre_quant_scale = (
+                    torch.rand(in_dim) + 0.5
+                )
+                for q in getattr(experts, f"{name}_weight_quantizers"):
+                    q.pre_quant_scale = torch.rand(in_dim) + 0.5
+            for grad in (True, False):
+                with torch.set_grad_enabled(grad):
+                    outputs.append(experts(hidden, top_k_index, top_k_weights).detach())
+        finally:
+            self._cleanup_registry(_SyntheticGroupedExperts)
+        torch.testing.assert_close(outputs[1], outputs[0], rtol=0, atol=0)
+
     def test_fake_quantized_grouped_forward_and_backward_match_eager(self):
         """Under autograd (QAT) the grouped forward fake-quantizes every expert's slice at once; its
         output and gradients match the eager forward's on the same calibrated quantizers."""

From 4f1b64a31d5742fd4d4f1e1ee63d12b4bb0f20d3 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 17:26:08 +0000
Subject: [PATCH 18/26] Pin the export refusal and the restore break for
 transposed experts

The changelog says the Hugging Face export refuses an MoE block whose
transposed experts were left unquantized, and that a state saved when an
earlier release quantized them no longer restores. One test runs the
export's MoE preparation on a block is_moe recognizes and expects
NotImplementedError; the other quantizes such a container the way an
earlier release did, then restores its state into a fresh model and expects
ApplyModeError.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../plugins/test_fused_experts.py             | 52 +++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 1b820e9cbb4..2f5110325cb 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -31,6 +31,7 @@
 from transformers import PretrainedConfig
 from transformers.integrations.finegrained_fp8 import FP8Experts
 
+import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
 import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module
 from modelopt.torch.export.model_utils import _release_exported_tensors
@@ -41,6 +42,9 @@
     get_quantization_format,
     process_layer_quant_config,
 )
+from modelopt.torch.export.unified_export_hf import _prepare_moe_inputs
+from modelopt.torch.models import is_moe
+from modelopt.torch.opt.conversion import ApplyModeError
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name
 from modelopt.torch.quantization.model_calib import (
@@ -493,6 +497,54 @@ class _TransposedExperts(_SyntheticFusedExperts):
             register_fused_experts_on_the_fly(model)
         assert QuantModuleRegistry.get(_TransposedExperts) is None
 
+    def test_a_transposed_containers_moe_block_is_refused_at_export(self):
+        """Left unquantized, a transposed container reaches the export inside an MoE block it
+        recognizes, with no handler for its experts, so the export refuses the block rather than
+        write it."""
+
+        class _TransposedExperts(_SyntheticFusedExperts):
+            pass
+
+        model = _TinyMoEModel()
+        model.moe.experts = _TransposedExperts()
+        model.moe.experts.is_transposed = True
+        with pytest.warns(UserWarning, match="store their weights transposed"):
+            register_fused_experts_on_the_fly(model)
+        assert is_moe(model.moe)
+
+        with pytest.raises(NotImplementedError):
+            _prepare_moe_inputs(model, torch.float16, is_modelopt_qlora=False)
+
+    def test_a_state_saved_with_transposed_experts_quantized_no_longer_restores(self):
+        """An earlier release quantized such a container, wrongly. Restoring that state now finds its
+        experts left unquantized and refuses the quantizers saved for them."""
+
+        class _TransposedExperts(_SyntheticFusedExperts):
+            pass
+
+        def build():
+            model = _TinyMoEModel()
+            model.moe.experts = _TransposedExperts()
+            model.moe.experts.is_transposed = True
+            return model
+
+        model = build()
+        # As an earlier release did: matched by shape, with the gated wrapper.
+        QuantModuleRegistry.register({_TransposedExperts: "test.TransposedExperts"})(
+            _QuantFusedExperts
+        )
+        try:
+            mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(torch.randn(2, 4, HIDDEN_DIM)))
+            state = mto.modelopt_state(model)
+        finally:
+            self._cleanup_registry(_TransposedExperts)
+
+        with (
+            pytest.warns(UserWarning, match="store their weights transposed"),
+            pytest.raises(ApplyModeError),
+        ):
+            mto.restore_from_modelopt_state(build(), state)
+
     def test_a_wrapper_that_defines_its_weights_in_get_expert_weight_is_eager_only(self):
         """The grouped_mm forward quantizes the stored weight without get_expert_weight, so a wrapper
         that defines its weight domain there must calibrate on the eager forward. One that does not

From f85148d4447f0cda3a68c4f5f677374054c5307f Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 17:26:12 +0000
Subject: [PATCH 19/26] File the restore break under Backward Breaking Changes

A modelopt_state saved for a model with transposed generic fused experts no
longer restores; that is a break for whoever holds one, not a bug fix, so
it gets its own entry there.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 CHANGELOG.rst | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index a5c54cad3fe..90aa9a89b4c 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -50,6 +50,7 @@ Changelog
 - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata.
 - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm.
 - The Windows FVD example rejects I3D checkpoints requiring unrestricted pickle loading. Custom TorchScript archives now require ``--weights-sha256`` with a trusted checksum; tensor state dictionaries continue to load directly.
+- A ``modelopt_state`` saved by an earlier release for a Hugging Face model whose fused MoE experts are stored transposed without a quantized wrapper of their own (for example Aria's) no longer restores, since those experts are now left unquantized. The checkpoints such a state exported had wrong expert weights, so there is nothing to migrate.
 
 - Layerwise-calibration resume (``checkpoint_dir``) now loads checkpoint files with ``weights_only=True``. A checkpoint directory written by an earlier release whose layer inputs hold objects other than tensors, plain containers or Hugging Face dynamic caches can no longer be resumed; recalibrate in a fresh directory, or disable checkpointing.
 
@@ -61,7 +62,7 @@ Changelog
 **Bug Fixes**
 
 - ``hf_quant_config.json`` now declares ``pre_quant_scale`` only for INT4 / W4A8 layers that have the tensor; loading INT4 / W4A8 fused-MoE checkpoints in TensorRT-LLM also needs a release whose MoE loader treats it as optional. AWQ, SmoothQuant and SVDQuant now warn when they skip quantized fused-MoE experts.
-- Hugging Face fused MoE experts stored transposed, ``(num_experts, in_dim, out_dim)``, without a quantized wrapper of their own (for example Aria's) are now left unquantized with a warning. They were quantized with gate and up split along the wrong axis, so their exported weights were wrong; the Hugging Face export now refuses such a model's MoE block instead, and a ``modelopt_state`` saved for such a model by an earlier release no longer restores.
+- Hugging Face fused MoE experts stored transposed, ``(num_experts, in_dim, out_dim)``, without a quantized wrapper of their own (for example Aria's) are now left unquantized with a warning. They were quantized with gate and up split along the wrong axis, so their exported weights were wrong; the Hugging Face export now refuses such a model's MoE block instead.
 - Calibrating a Megatron-Core MLA model now gives ``linear_q_down_proj`` and ``linear_kv_down_proj`` one per-tensor input and weight range, as vLLM and TensorRT-LLM fuse them into one GEMM with one scale per tensor. Re-calibrate existing MLA checkpoints to get the shared ranges.
 - ``export_mcore_gpt_to_hf`` and ``capture_kv_cache_export_state`` now write NVFP4 KV-cache ``k_scale`` / ``v_scale`` as amax / 448, like the HF exporter and as vLLM and TensorRT-LLM expect. They wrote amax / (6 * 448), which overflows in the attention kernels. Re-export Megatron checkpoints with an NVFP4 KV cache.
 - Preserve accumulated histogram counts when computing an entropy calibration threshold with NumPy histograms.

From 77c6f7cfb19dab647b5e62b7a806823869e65044 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 17:41:58 +0000
Subject: [PATCH 20/26] Promote the grouped weight's buffer as the stack does
 when experts disagree

Without autograd the per-slice buffer took the first slice's dtype, so an
expert whose pre-quant scale widens it past the first one's had its slice
narrowed, while the autograd path's torch.stack promotes every slice to the
widest dtype. The buffer now moves to the promoted dtype when a later slice
is wider. The widening test gains a case where only a later expert is
widened; it fails without the change.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/quantization/plugins/huggingface.py |  8 +++++---
 .../quantization/plugins/test_fused_experts.py     | 14 ++++++++------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 6bbf01ba7aa..798ee49ac4a 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -925,14 +925,16 @@ def _quantize_grouped_weight(self, weight: torch.Tensor, quantizers: nn.ModuleLi
             # the backward (QAT), so stack them there.
             return torch.stack([q(w) for q, w in zip(quantizers, weight, strict=True)])
         # Each slice goes into the output as it is made: a stack would hold all of them beside it.
-        # Strict, so every row of the output is written or the call fails. The output keeps the
-        # quantizers' dtype, as the stack does: a pre-quant scale can widen it, and narrowing it
-        # here would compute from other weights than with autograd.
+        # Strict, so every row of the output is written or the call fails. The output takes the
+        # dtype the stack promotes the slices to: a pre-quant scale can widen one, and narrowing
+        # it here would compute from other weights than with autograd.
         out: torch.Tensor | None = None
         for idx, (q, w) in enumerate(zip(quantizers, weight, strict=True)):
             quantized = q(w)
             if out is None:
                 out = torch.empty_like(weight, dtype=quantized.dtype)
+            elif quantized.dtype != out.dtype:
+                out = out.to(torch.promote_types(out.dtype, quantized.dtype))
             out[idx] = quantized
             del quantized  # not alive beside the next slice
         return out
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 2f5110325cb..7c8dac3d500 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -2000,10 +2000,11 @@ def test_disabled_quantizers_keep_their_pre_quant_scale_and_rotation(
             self._cleanup_registry(_SyntheticGroupedExperts)
         torch.testing.assert_close(outputs[1], outputs[0], rtol=1e-4, atol=1e-6)
 
-    def test_a_widening_pre_quant_scale_gives_one_grouped_forward_in_both_grad_modes(self):
-        """FP32 pre-quant scales on a BF16 model widen what the quantizers return. The grouped forward
-        computes from those FP32 slices with or without autograd, rather than rounding the weights to
-        BF16 only without it."""
+    @pytest.mark.parametrize("widened", ["every-expert", "a-later-expert"])
+    def test_a_widening_pre_quant_scale_gives_one_grouped_forward_in_both_grad_modes(self, widened):
+        """FP32 pre-quant scales on a BF16 model widen what the quantizers return, on every expert or
+        on one after the first. The grouped forward computes from the widened slices with or without
+        autograd, rather than rounding them to BF16 only without it."""
         self._cleanup_registry(_SyntheticGroupedExperts)
         torch.manual_seed(0)
         model = _GroupedMoEModel("grouped_mm").to(torch.bfloat16)
@@ -2022,8 +2023,9 @@ def test_a_widening_pre_quant_scale_gives_one_grouped_forward_in_both_grad_modes
                 getattr(experts, f"{name}_input_quantizer").pre_quant_scale = (
                     torch.rand(in_dim) + 0.5
                 )
-                for q in getattr(experts, f"{name}_weight_quantizers"):
-                    q.pre_quant_scale = torch.rand(in_dim) + 0.5
+                for idx, q in enumerate(getattr(experts, f"{name}_weight_quantizers")):
+                    if widened == "every-expert" or idx == 1:
+                        q.pre_quant_scale = torch.rand(in_dim) + 0.5
             for grad in (True, False):
                 with torch.set_grad_enabled(grad):
                     outputs.append(experts(hidden, top_k_index, top_k_weights).detach())

From 30d5750ec3adf77908babdec945f45a94517457d Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 17:42:07 +0000
Subject: [PATCH 21/26] Say the grouped_mm fallback does not call F.linear

Installing the F.linear and _grouped_linear patches together is safe only
because the grouped_mm fallback does not route through F.linear; the class
docstring now says so, and that the CPU tests running the fallback pin it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/quantization/plugins/huggingface.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 798ee49ac4a..607adf77a61 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -804,8 +804,9 @@ class _QuantFusedExperts(_QuantFunctionalMixin):
     forwards are intercepted. Both calibrate the same per-expert weight quantizers
     and shared input quantizers, and export the same keys; the activation scales
     they calibrate can differ. Only one backend's hook fires per forward
-    (selected by ``config._experts_implementation``), so installing both is safe.
-    Calibration still runs every model on ``eager``
+    (selected by ``config._experts_implementation``), so installing both is safe:
+    the ``grouped_mm`` fallback does not call ``F.linear``, which the CPU tests that
+    run it pin. Calibration still runs every model on ``eager``
     (``force_eager_experts_impl_on_the_fly``).
 
     Limitation: the ``grouped_mm`` path assumes the non-transposed weight layout

From 5f9739a27089bc97bb45a951f48a453298374d86 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 17:51:42 +0000
Subject: [PATCH 22/26] Say why calibration is still forced onto eager now that
 grouped_mm is intercepted

force_eager_experts_impl_on_the_fly's docstring said every non-eager backend
bypasses the quantizer hooks; after this change only batched_mm and sonicmoe
do. grouped_mm is intercepted, and calibration stays on eager until the
per-expert input captures follow it there.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/quantization/plugins/huggingface.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 607adf77a61..50ec6341628 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -1591,10 +1591,12 @@ def force_eager_experts_impl_on_the_fly(model):
 
     HF transformers 5.0+ decorates fused-experts forwards with
     ``@use_experts_implementation``, which may dispatch to ``torch._grouped_mm``
-    or ``torch.bmm`` backends. Those backends bypass ``F.linear`` and so bypass
-    ``_QuantFusedExperts``'s input/weight quantizer hooks — calibration silently
-    does nothing, no ``input_scale`` / ``amax`` is collected, and the exported
-    checkpoint produces garbage at inference.
+    or ``torch.bmm`` backends. ``batched_mm`` and ``sonicmoe`` bypass ``F.linear``
+    and so bypass ``_QuantFusedExperts``'s input/weight quantizer hooks — calibration
+    silently does nothing, no ``input_scale`` / ``amax`` is collected, and the
+    exported checkpoint produces garbage at inference. ``grouped_mm`` is intercepted
+    (see :class:`_QuantFusedExperts`), but calibration stays on eager until the
+    per-expert input captures follow it there.
 
     Sets ``config._experts_implementation = "eager"`` on the model config (and
     recursively on ``text_config`` / ``vision_config`` / ``audio_config`` /

From 689d3727febb444bdbf63c91772f31fa3ab30ddb Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 18:10:05 +0000
Subject: [PATCH 23/26] Say what keeps the FP8 experts wrapper on eager

Its _eager_only refusal sits in the grouped interceptor, which transformers'
FP8 grouped_mm forward never calls, so for this wrapper the flag cannot
fire; force_eager_experts_impl_on_the_fly, on convert and restore, is what
keeps it on eager. The comment now says so.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/quantization/plugins/huggingface.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 50ec6341628..f1ef96758c5 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -1118,7 +1118,9 @@ class _QuantFP8Experts(_QuantFusedExperts):
     """Scale-aware PTQ for native Transformers FP8 experts, dequantized per expert."""
 
     # The block scales apply in get_expert_weight, which only the eager forward reaches: transformers'
-    # FP8 grouped_mm forward runs its own kernels, which nothing here intercepts.
+    # FP8 grouped_mm forward runs its own kernels, which nothing here intercepts. So, unlike for the
+    # non-gated wrapper, the grouped interceptor's refusal never fires for these: what keeps them on
+    # eager is force_eager_experts_impl_on_the_fly, which runs on convert and on restore.
     _eager_only = True
 
     @property

From 1274887dc6c27849cba124ec2d89fcbf357a806a Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 18:10:14 +0000
Subject: [PATCH 24/26] Say the transposed check sees only a declared non-gated
 layout

The shape test separates the layouts of a gated container, whose first
projection is 2 * intermediate; a non-gated container's two layouts have the
same shapes, so only its own is_transposed gives it away.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/quantization/plugins/huggingface.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index f1ef96758c5..883daa62360 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -1565,7 +1565,9 @@ def register_fused_experts_on_the_fly(model):
         if wrapper is not None and (
             getattr(module, "is_transposed", False)
             # Stored `(num_experts, out_dim, in_dim)`, the first projection takes the dim down_proj
-            # puts out; a gated container stored transposed never does, whether it says so or not.
+            # puts out. That separates the layouts for a gated container, whose first projection is
+            # `2 * intermediate`, whether it says so or not; a non-gated one's two layouts have the
+            # same shapes, so only its own `is_transposed` gives it away.
             or getattr(module, wrapper._first_proj_attr).shape[-1] != module.down_proj.shape[1]
         ):
             # `_fused_experts_wrapper` matches on parameter SHAPE only, so it also claims a

From 3cfe6e048a690140e227a6b22ddcea85b7f9d3c1 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 21:36:52 +0000
Subject: [PATCH 25/26] Count the FSDP2 ownership scans themselves, not
 fully_shard.state calls

The two ownership tests (#3088) counted fully_shard.state calls to show
the weight-access window does not rescan a unit. #3091 made
_get_fsdp2_mesh call fully_shard.state on every window entry, so the
count no longer measured scans and both tests failed on main. They now
count _fsdp2_param_groups, the scan's own walk; removing the cache, or
rescanning on every miss, still fails them.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/unit/torch/quantization/test_utils.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py
index 9b06862e56c..a22703fc91e 100644
--- a/tests/unit/torch/quantization/test_utils.py
+++ b/tests/unit/torch/quantization/test_utils.py
@@ -24,6 +24,7 @@
 
 from modelopt.torch.quantization.utils import (
     convert_quantization_axis_to_reduce_axis,
+    core_utils,
     reduce_amax,
     reduce_block_amax,
 )
@@ -31,6 +32,7 @@
     _disable_fsdp_unshard_reshard,
     _fsdp2_owned_params,
     _fsdp2_owns,
+    _fsdp2_param_groups,
     fsdp2_weight_access_and_writeback_context,
     persistent_materialization,
 )
@@ -415,11 +417,11 @@ def _fsdp2_owns_its_params_in_every_state_job(rank, size):
     finally:
         layer.reshard()
     # The set is kept for the unit: opening the window again costs the module's parameters only.
-    with mock.patch.object(fully_shard, "state", wraps=fully_shard.state) as state:
+    with mock.patch.object(core_utils, "_fsdp2_param_groups", wraps=_fsdp2_param_groups) as scans:
         for _ in range(3):
             with fsdp2_weight_access_and_writeback_context(model.dense, model):
                 pass
-    assert state.call_count == 0
+    assert scans.call_count == 0
 
 
 def test_fsdp2_owns_its_parameters_in_every_shard_state(cpu_dist_workers):
@@ -447,11 +449,11 @@ def _window_leaves_what_fsdp2_does_not_own_job(rank, size):
             assert weight.to_local().shape[0] == 2
     # Every expert weight misses the kept set, and one scan of the unit answers them all.
     with (
-        mock.patch.object(fully_shard, "state", wraps=fully_shard.state) as state,
+        mock.patch.object(core_utils, "_fsdp2_param_groups", wraps=_fsdp2_param_groups) as scans,
         fsdp2_weight_access_and_writeback_context(layer, layer),
     ):
         pass
-    assert state.call_count == 1
+    assert scans.call_count == 1
     assert all(weight.to_local().shape[0] == 2 for weight in layer.experts.parameters())
 
 

From 75265cb385b9d511c4ec0acda69a5032e862033b Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 28 Oct 2026 22:01:24 +0000
Subject: [PATCH 26/26] Say why transposed experts stay in the eager scan and
 that their export is refused

The force-eager scan still matches transposed containers, which keeps
GPT-OSS's own wrapper on the eager forward its down_proj input
quantizer intercepts. The transposed warning says export_hf_checkpoint
will refuse their MoE block, and _fused_expert_projections' docstring
says the grouped forward clears the expert index.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/quantization/model_calib.py         |  6 ++++--
 modelopt/torch/quantization/plugins/huggingface.py | 10 +++++++---
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index 30f16ecb8a1..7e432cce9fb 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -1074,8 +1074,10 @@ def _fused_expert_projections(module: nn.Module):
     """Yield ``(weight_name, weight, weight_quantizers, input_quantizer)`` per fused-experts projection.
 
     ``weight`` is the 3-D ``(num_experts, out, in)`` parameter, ``weight_quantizers`` holds one
-    quantizer per expert, and ``input_quantizer`` is shared by the experts and runs once per expert
-    forward with ``module._current_expert_idx`` already set to that expert.
+    quantizer per expert, and ``input_quantizer`` is shared by the experts. Under eager it runs once
+    per expert forward with ``module._current_expert_idx`` set to that expert; under grouped_mm once
+    per projection, on every expert's rows, with the index cleared (``None``), so a capture keyed on
+    it skips that batch.
     """
     first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
     for weight_name, quantizers_name, input_q_name in (
diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 883daa62360..0e6b1ad7271 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -1537,7 +1537,10 @@ def _is_fused_experts_module(module):
     """Check if a module is a fused MoE expert container compatible with _QuantFusedExperts.
 
     See :func:`_fused_experts_wrapper` for the recognized layouts (gated
-    ``gate_up_proj`` and non-gated ``up_proj``).
+    ``gate_up_proj`` and non-gated ``up_proj``). Transposed containers match too, though
+    registration skips them: :func:`force_eager_experts_impl_on_the_fly` must still put them on
+    eager, the only forward on which GPT-OSS's own wrapper quantizes ``down_proj``'s input (it
+    intercepts the eager forward's ``bmm``/``matmul``).
     """
     return _fused_experts_wrapper(module) is not None
 
@@ -1578,8 +1581,9 @@ def register_fused_experts_on_the_fly(model):
                 f"Fused MoE experts '{name}' ({mod_type.__name__}) store their weights transposed "
                 "(num_experts, in_dim, out_dim), which the generic _QuantFusedExperts export path "
                 "does not support: it splits gate/up along dim 1, which is the input dim for this "
-                "layout. Leaving them unquantized; such a class needs its own quantized wrapper, as "
-                "GptOssExperts and Llama4TextExperts have."
+                "layout. Leaving them unquantized, so export_hf_checkpoint will refuse their MoE "
+                "block; such a class needs its own quantized wrapper, as GptOssExperts and "
+                "Llama4TextExperts have."
             )
             continue
         if wrapper is not None: