From 4d746b251d9b643a7210e2bbdddc0cfedc90de62 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 20:18:18 +0000 Subject: [PATCH 01/16] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder Nemotron-H's non-gated experts get their own wrapper, _QuantNemotronHExperts, in modelopt/torch/models/nemotron_h/modeling_ptq.py, registered on transformers' NemotronHExperts (every supported release ships it) in place of the generic plugin's _QuantNonGatedFusedExperts. It is not eager-only: the grouped_mm interception reaches the non-gated layout through _first_proj_attr, so these experts calibrate on the backend transformers selects, as the gated ones do. The generic structural detection now claims only the gated layout. _is_fused_experts_module also counts a class registered with a _QuantFusedExperts subclass, so the backend normalization still finds Nemotron-H's experts, and transformers' non-gated FP8 experts, which stay on eager. The non-gated tests move to tests/unit/torch/models and run the real NemotronHExperts. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 2 +- .../torch/models/nemotron_h/modeling_ptq.py | 26 +- .../torch/quantization/plugins/huggingface.py | 60 ++-- .../models/test_modeling_ptq_registration.py | 7 + .../models/test_nemotron_h_modeling_ptq.py | 218 ++++++++++++ .../plugins/test_fused_experts.py | 321 +----------------- 6 files changed, 282 insertions(+), 352 deletions(-) create mode 100644 tests/unit/torch/models/test_nemotron_h_modeling_ptq.py diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 07e0ddd4129..c993dc1d746 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -409,7 +409,7 @@ def _export_fused_experts( """Split fused MoE expert weights and export per-expert quantization scales. Works with any module wrapped by ``_QuantFusedExperts`` (gated, with a fused - ``gate_up_proj``) or ``_QuantNonGatedFusedExperts`` (non-gated, with a single + ``gate_up_proj``) or by a non-gated subclass (``_is_gated = False``, with a single ``up_proj`` — e.g. NemotronH). Both store their projections as 3-D ``nn.Parameter`` tensors with per-expert quantizer ``nn.ModuleList`` s. diff --git a/modelopt/torch/models/nemotron_h/modeling_ptq.py b/modelopt/torch/models/nemotron_h/modeling_ptq.py index 46670b9b7eb..378bbaa4eb7 100644 --- a/modelopt/torch/models/nemotron_h/modeling_ptq.py +++ b/modelopt/torch/models/nemotron_h/modeling_ptq.py @@ -16,13 +16,37 @@ """Nemotron-H PTQ modeling (HF model type ``nemotron_h``).""" import torch.nn as nn +from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHExperts -from modelopt.torch.quantization.plugins.huggingface import _is_supported_hf_model +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins.huggingface import ( + _is_supported_hf_model, + _QuantFusedExperts, +) from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector __all__: list[str] = [] +class _QuantNemotronHExperts(_QuantFusedExperts): + """Quantized ``NemotronHExperts``: non-gated fused experts, ``down_proj(act_fn(up_proj(x)))``. + + The container holds 3-D ``up_proj`` and ``down_proj`` parameters instead of the gated + ``gate_up_proj``, and its forward runs ``up_proj`` then ``down_proj`` once per expert on the eager + path, or once per expert-sorted batch on ``grouped_mm``. So only the first projection's name + differs from :class:`_QuantFusedExperts`: its per-expert weight quantizers and shared input + quantizer live under ``up_proj_*``. + """ + + _first_proj_attr = "up_proj" + # Read by the export: a non-gated up_proj is exported whole, not split into gate|up halves. + _is_gated = False + + +if NemotronHExperts not in QuantModuleRegistry: + QuantModuleRegistry.register({NemotronHExperts: "hf.NemotronHExperts"})(_QuantNemotronHExperts) + + def is_nemotron_h_model(model: nn.Module) -> bool: return get_nemotron_h_decoder_layers(model) is not None diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index dd086013118..fe34fd12fa9 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -824,15 +824,14 @@ class _QuantFusedExperts(_QuantFunctionalMixin): kernel) backends are not intercepted; ``normalize_experts_impl_on_the_fly`` redirects them to ``grouped_mm`` before calibration. - The non-gated variant (``up_proj`` instead of ``gate_up_proj``, used by - NemotronH) is handled by :class:`_QuantNonGatedFusedExperts` via the + A non-gated layout (a single ``up_proj``) is handled by a subclass that overrides the ``_first_proj_attr`` / ``_is_gated`` hooks below; each layout names the first-projection quantizers after its backing parameter. """ # Name of the 3-D weight parameter feeding the first ``F.linear`` per expert. # Gated experts fuse gate+up into ``gate_up_proj``; non-gated experts use a - # single ``up_proj`` (see _QuantNonGatedFusedExperts). + # single ``up_proj``, which their subclass names here. _first_proj_attr = "gate_up_proj" # Whether the first projection packs a gate half that must be split on export. _is_gated = True @@ -1210,28 +1209,13 @@ def fold_weight(self, keep_attrs: bool = False): self._fold_weight_quantizer(q, (weight_slice,), keep_attrs) -class _QuantNonGatedFusedExperts(_QuantFusedExperts): - """Quantized wrapper for non-gated fused MoE experts. - - Used by NemotronH (transformers 5.5+ ``NemotronHExperts``), whose experts - are a *non-gated* MLP: a single ``up_proj`` (no gate half) and a ``down_proj``, - both stored as 3-D ``nn.Parameter`` s indexed per expert. - """ - - _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. So, unlike for the - # non-gated wrapper, the grouped interceptor's refusal never fires for these: what keeps them on - # eager is normalize_experts_impl_on_the_fly, which runs on convert and on restore. + # FP8 grouped_mm forward runs its own kernels, which nothing here intercepts. So the grouped + # interceptor's refusal never fires for these: what keeps them on eager is + # normalize_experts_impl_on_the_fly, which runs on convert and on restore. _eager_only = True @property @@ -1617,17 +1601,13 @@ def register_sparse_moe_on_the_fly(model): def _fused_experts_wrapper(module): - """Return the ``_QuantFusedExperts`` subclass for a fused-MoE-experts module, else None. + """Return :class:`_QuantFusedExperts` for a fused-MoE-experts module, else None. 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 (: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``. + a 3-D ``down_proj`` and a 3-D ``gate_up_proj`` fusing gate+up. Matches ``MixtralExperts``, + ``Qwen2MoeExperts``, ``Qwen3MoeExperts``, ``Qwen3_5MoeExperts``, ``DeepseekV3NaiveMoe``, + ``JambaExperts``, ``OlmoeExperts``, ``MiniMaxM2Experts``, ``MiniMaxM3VLExperts``, etc. Experts + of another layout register a wrapper of their own, as Nemotron-H's non-gated ones do. 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 @@ -1645,24 +1625,22 @@ def _fused_experts_wrapper(module): gate_up = getattr(module, "gate_up_proj", None) if isinstance(gate_up, (nn.Parameter, Tensor)) and gate_up.dim() == 3: return _QuantFusedExperts - up = getattr(module, "up_proj", None) - if isinstance(up, (nn.Parameter, Tensor)) and up.dim() == 3: - # Only claim non-gated experts that alternate up_proj then down_proj. - if getattr(module, "gate_proj", None) is None and gate_up is None: - return _QuantNonGatedFusedExperts return None 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``). Transposed containers match too, though - registration skips them: :func:`normalize_experts_impl_on_the_fly` must still find them, to keep - GPT-OSS's own wrapper on eager, the only forward on which it quantizes ``down_proj``'s input (it - intercepts the eager forward's ``bmm``/``matmul``). + That is one :func:`_fused_experts_wrapper` detects, or one whose class registered a subclass of + it, as Nemotron-H's non-gated experts and transformers' FP8 experts do. Transposed containers + match too, though registration skips them: :func:`normalize_experts_impl_on_the_fly` must still + find them, to keep GPT-OSS's own wrapper on eager, the only forward on which it quantizes + ``down_proj``'s input (it intercepts the eager forward's ``bmm``/``matmul``). """ - return _fused_experts_wrapper(module) is not None + if _fused_experts_wrapper(module) is not None: + return True + wrapper = QuantModuleRegistry.get(type(module)) if hasattr(module, "num_experts") else None + return wrapper is not None and issubclass(wrapper, _QuantFusedExperts) def register_fused_experts_on_the_fly(model): diff --git a/tests/gpu/torch/models/test_modeling_ptq_registration.py b/tests/gpu/torch/models/test_modeling_ptq_registration.py index 6782b64c344..60a230b2598 100644 --- a/tests/gpu/torch/models/test_modeling_ptq_registration.py +++ b/tests/gpu/torch/models/test_modeling_ptq_registration.py @@ -76,6 +76,13 @@ else: assert GptOssExperts in QuantModuleRegistry, "GptOssExperts not registered" +try: + from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHExperts +except ImportError: + pass +else: + assert NemotronHExperts in QuantModuleRegistry, "NemotronHExperts not registered" + assert register_falcon_linears_on_the_fly in CUSTOM_MODEL_PLUGINS, "Falcon callback missing" assert register_moe_linear_on_the_fly in CUSTOM_MODEL_PLUGINS, "Step callback missing" diff --git a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py new file mode 100644 index 00000000000..c2d6d20de38 --- /dev/null +++ b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for _QuantNemotronHExperts: Nemotron-H's non-gated fused experts, eager and grouped_mm.""" + +import copy + +import pytest +import torch +import torch.nn as nn + +pytest.importorskip("transformers") + +from transformers import NemotronHConfig +from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHExperts + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.moe_utils import _export_fused_experts +from modelopt.torch.export.quant_utils import get_quant_config, get_quantization_format +from modelopt.torch.models.nemotron_h.modeling_ptq import _QuantNemotronHExperts +from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer +from modelopt.torch.quantization.plugins.huggingface import ( + _is_fused_experts_module, + normalize_experts_impl_on_the_fly, +) +from modelopt.torch.quantization.utils.calib_utils import GPTQHelper +from modelopt.torch.quantization.utils.core_utils import weight_attr_names + +NUM_EXPERTS = 4 +HIDDEN_DIM = 32 +INTERMEDIATE_DIM = 16 +TOP_K = 2 + +# Per-tensor 8-bit inputs, per-channel 8-bit weights, on both projections. +QUANT_CFG = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*up_proj_input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + {"quantizer_name": "*down_proj_input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + {"quantizer_name": "*up_proj_weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*down_proj_weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": "max", +} + + +class _ExpertsModel(nn.Module): + """A tiny real ``NemotronHExperts`` on the ``impl`` experts backend, with random weights.""" + + def __init__(self, impl): + super().__init__() + config = NemotronHConfig( + n_routed_experts=NUM_EXPERTS, + num_experts_per_tok=TOP_K, + hidden_size=HIDDEN_DIM, + moe_intermediate_size=INTERMEDIATE_DIM, + ) + config._experts_implementation = impl + self.experts = NemotronHExperts(config) + with torch.no_grad(): + for param in self.experts.parameters(): + param.normal_(std=0.02) + + +def _batch(): + """Tokens routed round-robin, so every expert sees some.""" + 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) + return hidden, top_k_index, top_k_weights + + +def _forward_loop(model): + torch.manual_seed(0) + with torch.no_grad(): + model.experts(*_batch()) + + +def _eager_and_grouped(): + torch.manual_seed(0) + eager, grouped = _ExpertsModel("eager"), _ExpertsModel("grouped_mm") + grouped.load_state_dict(eager.state_dict()) + return eager, grouped + + +def test_registered_wrapper_quantizes_up_and_down_proj(): + experts = _ExpertsModel("eager").experts + assert _is_fused_experts_module(experts) + + converted = QuantModuleRegistry.convert(experts) + + assert isinstance(converted, _QuantNemotronHExperts) + assert len(converted.up_proj_weight_quantizers) == NUM_EXPERTS + assert len(converted.down_proj_weight_quantizers) == NUM_EXPERTS + assert not hasattr(converted, "gate_up_proj_input_quantizer") + assert set(weight_attr_names(converted)) == {"up_proj", "down_proj"} + + +@pytest.mark.parametrize("impl", ["eager", "grouped_mm"]) +def test_forward_passthrough_matches(impl): + """With its quantizers disabled, the wrapper computes what the unconverted experts do.""" + torch.manual_seed(0) + experts = _ExpertsModel(impl).experts + reference = copy.deepcopy(experts) + converted = QuantModuleRegistry.convert(experts) + for module in converted.modules(): + if isinstance(module, TensorQuantizer): + module.disable() + + with torch.no_grad(): + batch = _batch() + torch.testing.assert_close(converted(*batch), reference(*batch)) + + +def test_calibration_matches_eager(): + """grouped_mm calibrates every expert's quantizers to what eager does. + + The same weights and routing give the shared input quantizers the same activations and each + per-expert weight quantizer the same slice; only down_proj's input passes through a GEMM, hence + the tolerance on the input amax. + """ + eager, grouped = _eager_and_grouped() + for model in (eager, grouped): + mtq.quantize(model, QUANT_CFG, forward_loop=_forward_loop) + e, g = eager.experts, grouped.experts + assert (e._running_grouped_mm, g._running_grouped_mm) == (False, True) + + for name in ("up_proj_input_quantizer", "down_proj_input_quantizer"): + assert getattr(e, name).amax is not None, name + torch.testing.assert_close( + getattr(g, name).amax, getattr(e, name).amax, atol=1e-3, rtol=1e-2 + ) + for name in ("up_proj_weight_quantizers", "down_proj_weight_quantizers"): + for idx, (eq, gq) in enumerate(zip(getattr(e, name), getattr(g, name))): + assert eq.amax is not None, f"{name}[{idx}]" + torch.testing.assert_close(gq.amax, eq.amax) + + +def test_gptq_hessians_match_eager(monkeypatch): + """GPTQ gives each expert the Hessian of only its own tokens under grouped_mm, as eager.""" + hessians = {} + free = GPTQHelper.free + + def recording(helper): + hessians[helper.name] = (helper.n_samples, helper.hessian.clone()) + free(helper) + + monkeypatch.setattr(GPTQHelper, "free", recording) + cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": 4, "block_sizes": {-1: 16}}, + }, + ], + "algorithm": {"method": "gptq", "block_size": 16}, + } + found = [] + for model in _eager_and_grouped(): + hessians.clear() + mtq.quantize(model, copy.deepcopy(cfg), forward_loop=_forward_loop) + found.append(dict(hessians)) + + eager, grouped = found + assert eager.keys() == grouped.keys() and len(eager) == 2 * NUM_EXPERTS + for name, (n_samples, hessian) in eager.items(): + assert n_samples > 0 and grouped[name][0] == n_samples, name + # The down-proj input flows through a GEMM, hence the tolerance. + torch.testing.assert_close(grouped[name][1], hessian, rtol=1e-3, atol=1e-4) + + +def test_export_writes_up_and_down_proj_per_expert(): + """The export leaves up_proj whole: there is no gate half to split off.""" + model = _ExpertsModel("grouped_mm") + mtq.quantize(model, QUANT_CFG, forward_loop=_forward_loop) + experts = model.experts + assert get_quantization_format(experts) is not None + assert get_quant_config(model)["quantization"].get("quant_algo") is not None + up_scale, down_scale = ( + (q.amax / q.maxbound).squeeze() + for q in (experts.up_proj_input_quantizer, experts.down_proj_input_quantizer) + ) + + _export_fused_experts(experts, torch.float16) + + for idx in range(NUM_EXPERTS): + expert = getattr(experts, str(idx)) + assert not hasattr(expert, "gate_proj") + assert expert.up_proj.weight.shape == (INTERMEDIATE_DIM, HIDDEN_DIM) + assert expert.down_proj.weight.shape == (HIDDEN_DIM, INTERMEDIATE_DIM) + torch.testing.assert_close(expert.up_proj.input_scale, up_scale) + torch.testing.assert_close(expert.down_proj.input_scale, down_scale) + assert not hasattr(experts, "up_proj_weight_quantizers") + + +def test_unsupported_backend_is_redirected_to_grouped_mm(): + """Found through their registration, the experts leave a backend the quantizers cannot see.""" + model = _ExpertsModel("batched_mm") + + normalize_experts_impl_on_the_fly(model) + + assert model.experts.config._experts_implementation == "grouped_mm" diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index d7f03854837..76132e227cd 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -71,7 +71,6 @@ _is_fused_experts_module, _is_sparse_sequaential_moe_block, _QuantFusedExperts, - _QuantNonGatedFusedExperts, assert_fused_experts_intercepted, normalize_experts_impl_on_the_fly, register_fused_experts_on_the_fly, @@ -192,45 +191,6 @@ def forward(self, hidden_states, top_k_index, top_k_weights): return final_hidden_states -class _SyntheticNonGatedFusedExperts(nn.Module): - """Mimics NemotronHExperts (transformers 5.5+): non-gated fused experts. - - A single ``up_proj`` (no gate half) + ``down_proj``, both 3-D ``nn.Parameter`` s, - with the forward calling ``F.linear`` exactly twice per expert (up then down). - """ - - def __init__(self): - super().__init__() - self.num_experts = NUM_EXPERTS - self.hidden_dim = HIDDEN_DIM - self.intermediate_dim = INTERMEDIATE_DIM - self.up_proj = nn.Parameter(torch.randn(NUM_EXPERTS, 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 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] - current_hidden_states = F.linear(current_state, self.up_proj[expert_idx]) - current_hidden_states = self.act_fn(current_hidden_states) - 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 - - class _SyntheticTopKRouter(nn.Module): def __init__(self): super().__init__() @@ -273,33 +233,6 @@ def forward(self, x): return self.moe(x) -class _SyntheticNonGatedSparseMoeBlock(nn.Module): - """Mimics NemotronHMoE: a router + non-gated fused experts.""" - - def __init__(self): - super().__init__() - self.gate = _SyntheticTopKRouter() - self.experts = _SyntheticNonGatedFusedExperts() - - 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 _TinyNonGatedMoEModel(nn.Module): - """Minimal model containing a single non-gated MoE block.""" - - def __init__(self): - super().__init__() - self.moe = _SyntheticNonGatedSparseMoeBlock() - - def forward(self, x): - return self.moe(x) - - def _route_once_to_each_expert(model): """Call fused experts directly with deterministic routing that covers every expert.""" assert NUM_EXPERTS % TOP_K == 0 @@ -733,11 +666,8 @@ def _convert(model): TestQuantFusedExperts._cleanup_registry(expert_type) return converted - @pytest.mark.parametrize( - "model_cls", [_TinyMoEModel, _TinyNonGatedMoEModel], ids=["gated", "non_gated"] - ) - def test_yields_the_same_quantizers_in_the_same_order(self, model_cls): - experts = self._convert(model_cls()) + def test_yields_the_same_quantizers_in_the_same_order(self): + experts = self._convert(_TinyMoEModel()) from_weights = [q for _, q in experts.iter_weights_for_calibration()] quantizers_only = list(experts.iter_weight_quantizers_for_calibration()) @@ -745,12 +675,9 @@ def test_yields_the_same_quantizers_in_the_same_order(self, model_cls): assert quantizers_only, "expected per-expert weight quantizers" assert [id(q) for q in quantizers_only] == [id(q) for q in from_weights] - @pytest.mark.parametrize( - "model_cls", [_TinyMoEModel, _TinyNonGatedMoEModel], ids=["gated", "non_gated"] - ) - def test_does_not_index_the_fused_weight(self, model_cls): + def test_does_not_index_the_fused_weight(self): """The point of the override: no ``weight[idx]``, which is the per-expert collective.""" - experts = self._convert(model_cls()) + experts = self._convert(_TinyMoEModel()) class _NoIndexing(torch.Tensor): @staticmethod @@ -1432,11 +1359,16 @@ def test_no_crash_when_impl_attr_missing(self): normalize_experts_impl_on_the_fly(model) assert not hasattr(cfg, "_experts_implementation") + @pytest.mark.parametrize("has_gate", [True, False]) @pytest.mark.parametrize("impl", ["grouped_mm", None]) - def test_forces_eager_for_experts_quantized_on_the_eager_forward_only(self, impl, recwarn): + def test_forces_eager_for_experts_quantized_on_the_eager_forward_only( + self, impl, has_gate, recwarn + ): """Native FP8 experts apply their block scales on the eager forward; transformers' FP8 grouped_mm forward runs kernels nothing intercepts, so a model holding them calibrates on - eager, as main forced every MoE to. Moving it off grouped_mm warns; None is eager already.""" + eager, as main forced every MoE to. Moving it off grouped_mm warns; None is eager already. + Non-gated ones (an FP8 Nemotron-H's) are found through their registration, since the + structural detection claims only the gated layout.""" cfg = _StubConfig() cfg._experts_implementation = impl model = _NonMoEModelWithConfig(cfg) @@ -1445,6 +1377,7 @@ def test_forces_eager_for_experts_quantized_on_the_eager_forward_only(self, impl hidden_size=256, intermediate_size=256, num_experts=2, hidden_act="silu" ), block_size=(128, 128), + has_gate=has_gate, ) normalize_experts_impl_on_the_fly(model) assert model.config._experts_implementation == "eager" @@ -2417,236 +2350,6 @@ def test_local_hessian_captures_per_expert_in_grouped_mode(self, cpu_dist_worker cpu_dist_workers.run(_local_hessian_grouped_capture_job, impl=impl) -# --------------------------------------------------------------------------- -# Tests for the non-gated fused-experts path (NemotronH NemotronHExperts): -# single up_proj (no gate half) + down_proj. Quantizers are named after the -# backing weights: up_proj_* and down_proj_*. -# --------------------------------------------------------------------------- -class TestNonGatedFusedExperts: - @staticmethod - def _cleanup_registry(mod_type): - if QuantModuleRegistry.get(mod_type) is not None: - QuantModuleRegistry.unregister(mod_type) - - def test_detected_and_picks_nongated_wrapper(self): - module = _SyntheticNonGatedFusedExperts() - assert _is_fused_experts_module(module) is True - assert _fused_experts_wrapper(module) is _QuantNonGatedFusedExperts - - def test_gated_still_picks_base_wrapper(self): - assert _fused_experts_wrapper(_SyntheticFusedExperts()) is _QuantFusedExperts - - def test_register_uses_nongated_wrapper(self): - model = _TinyNonGatedMoEModel() - expert_type = type(model.moe.experts) - self._cleanup_registry(expert_type) - register_fused_experts_on_the_fly(model) - try: - converted = QuantModuleRegistry.convert(model.moe.experts) - assert isinstance(converted, _QuantNonGatedFusedExperts) - assert converted._first_proj_attr == "up_proj" - assert converted._is_gated is False - assert hasattr(converted, "up_proj_input_quantizer") - assert hasattr(converted, "up_proj_weight_quantizers") - assert not hasattr(converted, "gate_up_proj_input_quantizer") - assert not hasattr(converted, "gate_up_proj_weight_quantizers") - assert len(converted.up_proj_weight_quantizers) == NUM_EXPERTS - assert len(converted.down_proj_weight_quantizers) == NUM_EXPERTS - finally: - self._cleanup_registry(expert_type) - - def test_forward_passthrough_matches(self): - model = _TinyNonGatedMoEModel() - expert_type = type(model.moe.experts) - self._cleanup_registry(expert_type) - - ref_experts = _SyntheticNonGatedFusedExperts() - ref_experts.load_state_dict(model.moe.experts.state_dict()) - - register_fused_experts_on_the_fly(model) - try: - converted = QuantModuleRegistry.convert(model.moe.experts) - # Disable quantizers to isolate the wrapper's structural forward - # (the F.linear interception / per-expert index routing) from - # dynamic-quant noise — this is a passthrough equivalence check. - for q in converted.modules(): - if isinstance(q, TensorQuantizer): - q.disable() - seq_len = 8 - hidden_states = torch.randn(seq_len, HIDDEN_DIM) - top_k_index = torch.randint(0, NUM_EXPERTS, (seq_len, TOP_K)) - top_k_weights = torch.softmax(torch.randn(seq_len, TOP_K), dim=-1) - with torch.no_grad(): - out_ref = ref_experts(hidden_states, top_k_index, top_k_weights) - out_test = converted(hidden_states, top_k_index, top_k_weights) - assert torch.allclose(out_ref, out_test, atol=1e-5), ( - f"Max diff: {(out_ref - out_test).abs().max().item()}" - ) - finally: - self._cleanup_registry(expert_type) - - def test_expert_index_recovery(self): - experts = _SyntheticNonGatedFusedExperts() - expert_type = type(experts) - self._cleanup_registry(expert_type) - register_fused_experts_on_the_fly(_TinyNonGatedMoEModel()) - try: - converted = QuantModuleRegistry.convert(experts) - for idx in range(NUM_EXPERTS): - weight_slice = converted.up_proj[idx] - assert converted._get_expert_idx_from_first_proj(weight_slice) == idx - finally: - self._cleanup_registry(expert_type) - - def _nongated_fp8_cfg(self): - return { - "quant_cfg": [ - {"quantizer_name": "*", "enable": False}, - { - "quantizer_name": "*up_proj_input_quantizer", - "cfg": {"num_bits": 8, "axis": None}, - }, - { - "quantizer_name": "*down_proj_input_quantizer", - "cfg": {"num_bits": 8, "axis": None}, - }, - { - "quantizer_name": "*up_proj_weight_quantizer", - "cfg": {"num_bits": 8, "axis": 0}, - }, - { - "quantizer_name": "*down_proj_weight_quantizer", - "cfg": {"num_bits": 8, "axis": 0}, - }, - ], - "algorithm": "max", - } - - def test_calibration_populates_all_expert_quantizers(self): - model = _TinyNonGatedMoEModel() - expert_type = type(model.moe.experts) - self._cleanup_registry(expert_type) - - def forward_loop(m): - torch.manual_seed(0) - _route_once_to_each_expert(m) - - try: - mtq.quantize(model, self._nongated_fp8_cfg(), forward_loop=forward_loop) - experts = model.moe.experts - assert experts.up_proj_input_quantizer.amax is not None - assert experts.down_proj_input_quantizer.amax is not None - for idx in range(NUM_EXPERTS): - assert experts.up_proj_weight_quantizers[idx].amax is not None - assert experts.down_proj_weight_quantizers[idx].amax is not None - finally: - self._cleanup_registry(expert_type) - - def test_export_creates_per_expert_up_down_only(self): - model = _TinyNonGatedMoEModel() - expert_type = type(model.moe.experts) - self._cleanup_registry(expert_type) - - def forward_loop(m): - torch.manual_seed(0) - for _ in range(2): - m(torch.randn(1, 4, HIDDEN_DIM)) - - try: - mtq.quantize(model, self._nongated_fp8_cfg(), forward_loop=forward_loop) - converted = model.moe.experts - expected_up_input_scale = ( - converted.up_proj_input_quantizer.amax / converted.up_proj_input_quantizer.maxbound - ).squeeze() - expected_down_input_scale = ( - converted.down_proj_input_quantizer.amax - / converted.down_proj_input_quantizer.maxbound - ).squeeze() - _export_fused_experts(converted, torch.float16) - - for idx in range(NUM_EXPERTS): - expert_mod = getattr(converted, str(idx), None) - assert expert_mod is not None, f"Missing expert submodule {idx}" - # Non-gated: up_proj + down_proj, but NO gate_proj. - assert hasattr(expert_mod, "up_proj"), f"Expert {idx} missing up_proj" - assert hasattr(expert_mod, "down_proj"), f"Expert {idx} missing down_proj" - assert not hasattr(expert_mod, "gate_proj"), ( - f"Expert {idx} should NOT have gate_proj (non-gated MLP)" - ) - assert expert_mod.up_proj.weight.shape == (INTERMEDIATE_DIM, HIDDEN_DIM) - assert expert_mod.down_proj.weight.shape == (HIDDEN_DIM, INTERMEDIATE_DIM) - torch.testing.assert_close(expert_mod.up_proj.input_scale, expected_up_input_scale) - torch.testing.assert_close( - expert_mod.down_proj.input_scale, expected_down_input_scale - ) - - # Fused params and per-expert quantizer lists are removed. - assert not hasattr(converted, "up_proj") - assert not hasattr(converted, "down_proj") - assert not hasattr(converted, "up_proj_weight_quantizers") - assert not hasattr(converted, "down_proj_weight_quantizers") - finally: - self._cleanup_registry(expert_type) - - def test_enumeration_yields_up_and_down_proj(self): - """weight_attr_names must yield up_proj and down_proj for non-gated experts.""" - model = _TinyNonGatedMoEModel() - expert_type = type(model.moe.experts) - self._cleanup_registry(expert_type) - register_fused_experts_on_the_fly(model) - try: - converted = QuantModuleRegistry.convert(model.moe.experts) - assert set(weight_attr_names(converted)) == {"up_proj", "down_proj"} - finally: - self._cleanup_registry(expert_type) - - def test_split_gated_layout_not_claimed_as_nongated(self): - """A fused container with a separate 3-D gate_proj (split-gated: three - F.linear calls per expert) must NOT be claimed by the non-gated wrapper, - whose two-call toggle and up_proj-storage index recovery assume exactly - two projections. It is left unsupported (None) rather than mis-quantized.""" - - class _SplitGatedExperts(nn.Module): - def __init__(self): - super().__init__() - self.num_experts = NUM_EXPERTS - self.gate_proj = nn.Parameter( - torch.randn(NUM_EXPERTS, INTERMEDIATE_DIM, HIDDEN_DIM) * 0.02 - ) - self.up_proj = nn.Parameter( - torch.randn(NUM_EXPERTS, 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() - - module = _SplitGatedExperts() - assert _fused_experts_wrapper(module) is None - assert _is_fused_experts_module(module) is False - - def test_get_quant_config_resolves_nongated_experts(self): - """get_quant_config must detect the non-gated experts as quantized.""" - model = _TinyNonGatedMoEModel() - expert_type = type(model.moe.experts) - self._cleanup_registry(expert_type) - - def forward_loop(m): - torch.manual_seed(0) - for _ in range(2): - m(torch.randn(1, 4, HIDDEN_DIM)) - - try: - mtq.quantize(model, self._nongated_fp8_cfg(), forward_loop=forward_loop) - # Format resolves (via down_proj) instead of QUANTIZATION_NONE (None). - assert get_quantization_format(model.moe.experts) is not None - # The non-gated experts are reflected in the produced quant config. - quant = get_quant_config(model)["quantization"] - assert quant.get("quant_algo") is not None - finally: - self._cleanup_registry(expert_type) - - # --------------------------------------------------------------------------- # Regression tests for the real transformers Qwen3-VL MoE text experts # --------------------------------------------------------------------------- From 7c48966625164e3500e6a76c7e939f1eb2d3e0a0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:00:30 +0000 Subject: [PATCH 02/16] Drop the plugin-order test of the non-gated stand-in Nemotron-H's experts register at import here, so no class the on-the-fly registration claims is eager-only, and the stand-in the test needed is gone with the generic non-gated detection. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/test_fused_experts.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 76132e227cd..1b0da05da78 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -1384,15 +1384,6 @@ def test_forces_eager_for_experts_quantized_on_the_eager_forward_only( moved = [w for w in recwarn if "instead of experts_implementation" in str(w.message)] assert len(moved) == (impl is not None) - def test_forces_eager_for_eager_only_experts_not_registered_yet(self): - """The plugins run in no set order, so normalization can come before the on-the-fly - registration; the non-gated experts quantize on the eager forward only either way.""" - TestQuantFusedExperts._cleanup_registry(_SyntheticNonGatedFusedExperts) - model = _NonMoEModelWithConfig(_StubConfig(impl="grouped_mm")) - model.experts = _SyntheticNonGatedFusedExperts() - normalize_experts_impl_on_the_fly(model) - assert model.config._experts_implementation == "eager" - def test_refuses_experts_quantized_on_the_eager_forward_only_under_expert_parallelism(self): """Expert parallelism runs the experts on the torch_a2a forward, which cannot fall back to eager: such experts would calibrate on raw FP8 bytes, or not at all. Refused before any From 8514dbe96c598da8067b7ed1ca7c59b5ba89c61c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:43:08 +0000 Subject: [PATCH 03/16] Say Nemotron-H's experts wrapper relies on transformers' non-transposed layout Signed-off-by: Shengliang Xu --- modelopt/torch/models/nemotron_h/modeling_ptq.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/models/nemotron_h/modeling_ptq.py b/modelopt/torch/models/nemotron_h/modeling_ptq.py index 378bbaa4eb7..dc9d256250c 100644 --- a/modelopt/torch/models/nemotron_h/modeling_ptq.py +++ b/modelopt/torch/models/nemotron_h/modeling_ptq.py @@ -35,7 +35,9 @@ class _QuantNemotronHExperts(_QuantFusedExperts): ``gate_up_proj``, and its forward runs ``up_proj`` then ``down_proj`` once per expert on the eager path, or once per expert-sorted batch on ``grouped_mm``. So only the first projection's name differs from :class:`_QuantFusedExperts`: its per-expert weight quantizers and shared input - quantizer live under ``up_proj_*``. + quantizer live under ``up_proj_*``. Registered explicitly, these experts skip the transposed-layout + check of :func:`register_fused_experts_on_the_fly`: the wrapper relies on transformers storing their + weights ``(num_experts, out_dim, in_dim)``, as every supported release does. """ _first_proj_attr = "up_proj" From b86897895e1432851cb78657180cb09372f8df21 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:44:23 +0000 Subject: [PATCH 04/16] Wrap Nemotron-H's experts docstring at 100 columns Signed-off-by: Shengliang Xu --- modelopt/torch/models/nemotron_h/modeling_ptq.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/models/nemotron_h/modeling_ptq.py b/modelopt/torch/models/nemotron_h/modeling_ptq.py index dc9d256250c..e0f51b5a28c 100644 --- a/modelopt/torch/models/nemotron_h/modeling_ptq.py +++ b/modelopt/torch/models/nemotron_h/modeling_ptq.py @@ -32,12 +32,13 @@ class _QuantNemotronHExperts(_QuantFusedExperts): """Quantized ``NemotronHExperts``: non-gated fused experts, ``down_proj(act_fn(up_proj(x)))``. The container holds 3-D ``up_proj`` and ``down_proj`` parameters instead of the gated - ``gate_up_proj``, and its forward runs ``up_proj`` then ``down_proj`` once per expert on the eager - path, or once per expert-sorted batch on ``grouped_mm``. So only the first projection's name - differs from :class:`_QuantFusedExperts`: its per-expert weight quantizers and shared input - quantizer live under ``up_proj_*``. Registered explicitly, these experts skip the transposed-layout - check of :func:`register_fused_experts_on_the_fly`: the wrapper relies on transformers storing their - weights ``(num_experts, out_dim, in_dim)``, as every supported release does. + ``gate_up_proj``, and its forward runs ``up_proj`` then ``down_proj`` once per expert on the + eager path, or once per expert-sorted batch on ``grouped_mm``. So only the first projection's + name differs from :class:`_QuantFusedExperts`: its per-expert weight quantizers and shared input + quantizer live under ``up_proj_*``. Registered explicitly, these experts skip the + transposed-layout check of :func:`register_fused_experts_on_the_fly`: the wrapper relies on + transformers storing their weights ``(num_experts, out_dim, in_dim)``, as every supported + release does. """ _first_proj_attr = "up_proj" From 30986733bed889dd3d09ac52553b3aba18de41cd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 22:53:37 +0000 Subject: [PATCH 05/16] Pin that Nemotron-H's grouped_mm calibration goes through _grouped_linear The grouped-vs-eager comparisons would pass just as well if the grouped forward fell back to per-expert F.linear calls, which the eager hook intercepts with the same values. The interception count tells the two hooks apart: one _grouped_linear call per projection, against one F.linear call per projection and expert. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_nemotron_h_modeling_ptq.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py index c2d6d20de38..60bf34e4bc2 100644 --- a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py +++ b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py @@ -139,6 +139,9 @@ def test_calibration_matches_eager(): mtq.quantize(model, QUANT_CFG, forward_loop=_forward_loop) e, g = eager.experts, grouped.experts assert (e._running_grouped_mm, g._running_grouped_mm) == (False, True) + # grouped_mm went through `_grouped_linear`, one call per projection, not through per-expert + # `F.linear` calls as eager does: the amaxes below are the grouped hook's. + assert (e._intercept_count, g._intercept_count) == (2 * NUM_EXPERTS, 2) for name in ("up_proj_input_quantizer", "down_proj_input_quantizer"): assert getattr(e, name).amax is not None, name From f9d0e084d0f63d33b04b1f6d363083a7b5179a36 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 22:55:11 +0000 Subject: [PATCH 06/16] Pin that the Nemotron-H export releases the fused up_proj and down_proj The test this one replaced checked it: for the non-gated layout up_proj is the whole fused weight, so an export that kept it would hold every expert twice, the fused original beside the per-expert copies. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_nemotron_h_modeling_ptq.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py index 60bf34e4bc2..5e84a924f3a 100644 --- a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py +++ b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py @@ -210,6 +210,8 @@ def test_export_writes_up_and_down_proj_per_expert(): torch.testing.assert_close(expert.up_proj.input_scale, up_scale) torch.testing.assert_close(expert.down_proj.input_scale, down_scale) assert not hasattr(experts, "up_proj_weight_quantizers") + # The fused parameters are released, not kept beside the per-expert copies. + assert not hasattr(experts, "up_proj") and not hasattr(experts, "down_proj") def test_unsupported_backend_is_redirected_to_grouped_mm(): From cb466fdcf2bee78ef131b2106397fe9c0a6be192 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 00:15:37 +0000 Subject: [PATCH 07/16] Count already-converted experts as fused experts _is_fused_experts_module found Nemotron-H's experts through their class's registration, which the class quantize() converts them to does not carry, so a second quantize() over the converted model skipped the backend normalization for them. Gated experts are found by their shapes either way. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 13 +++++++------ .../torch/models/test_nemotron_h_modeling_ptq.py | 2 ++ 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index fe34fd12fa9..f95ba4127f6 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1631,13 +1631,14 @@ def _fused_experts_wrapper(module): def _is_fused_experts_module(module): """Check if a module is a fused MoE expert container compatible with _QuantFusedExperts. - That is one :func:`_fused_experts_wrapper` detects, or one whose class registered a subclass of - it, as Nemotron-H's non-gated experts and transformers' FP8 experts do. Transposed containers - match too, though registration skips them: :func:`normalize_experts_impl_on_the_fly` must still - find them, to keep GPT-OSS's own wrapper on eager, the only forward on which it quantizes - ``down_proj``'s input (it intercepts the eager forward's ``bmm``/``matmul``). + That is one :func:`_fused_experts_wrapper` detects, one whose class registered a subclass of it, + as Nemotron-H's non-gated experts and transformers' FP8 experts do, or one already converted. + Transposed containers match too, though registration skips them: + :func:`normalize_experts_impl_on_the_fly` must still find them, to keep GPT-OSS's own wrapper on + eager, the only forward on which it quantizes ``down_proj``'s input (it intercepts the eager + forward's ``bmm``/``matmul``). """ - if _fused_experts_wrapper(module) is not None: + if _fused_experts_wrapper(module) is not None or isinstance(module, _QuantFusedExperts): return True wrapper = QuantModuleRegistry.get(type(module)) if hasattr(module, "num_experts") else None return wrapper is not None and issubclass(wrapper, _QuantFusedExperts) diff --git a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py index 5e84a924f3a..439788079f6 100644 --- a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py +++ b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py @@ -105,6 +105,8 @@ def test_registered_wrapper_quantizes_up_and_down_proj(): converted = QuantModuleRegistry.convert(experts) assert isinstance(converted, _QuantNemotronHExperts) + # A second quantize() finds the converted experts too, as it does gated ones by their shapes. + assert _is_fused_experts_module(converted) assert len(converted.up_proj_weight_quantizers) == NUM_EXPERTS assert len(converted.down_proj_weight_quantizers) == NUM_EXPERTS assert not hasattr(converted, "gate_up_proj_input_quantizer") From 6d072defeaa9c6778473cdcd456476ff145a2b53 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 00:16:19 +0000 Subject: [PATCH 08/16] Compare up_proj's input amax exactly between grouped_mm and eager up_proj's input is the hidden states, which the grouped forward hands _grouped_linear gathered by expert: the same rows eager feeds per expert, so their amax matches exactly. Only down_proj's input passes through a GEMM and keeps the tolerance. Signed-off-by: Shengliang Xu --- .../torch/models/test_nemotron_h_modeling_ptq.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py index 439788079f6..d4824cad169 100644 --- a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py +++ b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py @@ -134,7 +134,7 @@ def test_calibration_matches_eager(): The same weights and routing give the shared input quantizers the same activations and each per-expert weight quantizer the same slice; only down_proj's input passes through a GEMM, hence - the tolerance on the input amax. + the tolerance on its amax. """ eager, grouped = _eager_and_grouped() for model in (eager, grouped): @@ -145,11 +145,12 @@ def test_calibration_matches_eager(): # `F.linear` calls as eager does: the amaxes below are the grouped hook's. assert (e._intercept_count, g._intercept_count) == (2 * NUM_EXPERTS, 2) - for name in ("up_proj_input_quantizer", "down_proj_input_quantizer"): - assert getattr(e, name).amax is not None, name - torch.testing.assert_close( - getattr(g, name).amax, getattr(e, name).amax, atol=1e-3, rtol=1e-2 - ) + assert e.up_proj_input_quantizer.amax is not None + torch.testing.assert_close(g.up_proj_input_quantizer.amax, e.up_proj_input_quantizer.amax) + assert e.down_proj_input_quantizer.amax is not None + torch.testing.assert_close( + g.down_proj_input_quantizer.amax, e.down_proj_input_quantizer.amax, atol=1e-3, rtol=1e-2 + ) for name in ("up_proj_weight_quantizers", "down_proj_weight_quantizers"): for idx, (eq, gq) in enumerate(zip(getattr(e, name), getattr(g, name))): assert eq.amax is not None, f"{name}[{idx}]" From 1e8f0a5346121e6fc70e00bcb7abf2afd97f97eb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 01:09:01 +0000 Subject: [PATCH 09/16] Reuse _is_quant_fused_experts_module for the converted-experts check Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index f95ba4127f6..d12c53d176e 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1638,7 +1638,7 @@ def _is_fused_experts_module(module): eager, the only forward on which it quantizes ``down_proj``'s input (it intercepts the eager forward's ``bmm``/``matmul``). """ - if _fused_experts_wrapper(module) is not None or isinstance(module, _QuantFusedExperts): + if _fused_experts_wrapper(module) is not None or _is_quant_fused_experts_module(module): return True wrapper = QuantModuleRegistry.get(type(module)) if hasattr(module, "num_experts") else None return wrapper is not None and issubclass(wrapper, _QuantFusedExperts) From 52a327489df54ed961fa62f218f83ecf26616dce Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 07:19:36 +0000 Subject: [PATCH 10/16] Drop the converted-experts case from _is_fused_experts_module Its one caller, normalize_experts_impl_on_the_fly, runs inside replace_quant_module, which refuses a model that is already quantized, so no converted experts module reaches it. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 6 +++--- tests/unit/torch/models/test_nemotron_h_modeling_ptq.py | 2 -- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index d12c53d176e..9976db4502f 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1631,14 +1631,14 @@ def _fused_experts_wrapper(module): def _is_fused_experts_module(module): """Check if a module is a fused MoE expert container compatible with _QuantFusedExperts. - That is one :func:`_fused_experts_wrapper` detects, one whose class registered a subclass of it, - as Nemotron-H's non-gated experts and transformers' FP8 experts do, or one already converted. + That is one :func:`_fused_experts_wrapper` detects, or one whose class registered a subclass of + it, as Nemotron-H's non-gated experts and transformers' FP8 experts do. Transposed containers match too, though registration skips them: :func:`normalize_experts_impl_on_the_fly` must still find them, to keep GPT-OSS's own wrapper on eager, the only forward on which it quantizes ``down_proj``'s input (it intercepts the eager forward's ``bmm``/``matmul``). """ - if _fused_experts_wrapper(module) is not None or _is_quant_fused_experts_module(module): + if _fused_experts_wrapper(module) is not None: return True wrapper = QuantModuleRegistry.get(type(module)) if hasattr(module, "num_experts") else None return wrapper is not None and issubclass(wrapper, _QuantFusedExperts) diff --git a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py index d4824cad169..c284496a9f0 100644 --- a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py +++ b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py @@ -105,8 +105,6 @@ def test_registered_wrapper_quantizes_up_and_down_proj(): converted = QuantModuleRegistry.convert(experts) assert isinstance(converted, _QuantNemotronHExperts) - # A second quantize() finds the converted experts too, as it does gated ones by their shapes. - assert _is_fused_experts_module(converted) assert len(converted.up_proj_weight_quantizers) == NUM_EXPERTS assert len(converted.down_proj_weight_quantizers) == NUM_EXPERTS assert not hasattr(converted, "gate_up_proj_input_quantizer") From be88b378e5b864b948b2888642fc0504d75a3c54 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:10:34 +0000 Subject: [PATCH 11/16] Drop the import the moved non-gated tests no longer need Rebased onto #3074 on main, the fused-experts tests import #2745's quant_utils helpers; get_quantization_format was used only by the non-gated tests that moved to the Nemotron-H test file. Signed-off-by: Shengliang Xu --- tests/unit/torch/quantization/plugins/test_fused_experts.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 1b0da05da78..4f217657d18 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -50,7 +50,6 @@ from modelopt.torch.export.quant_utils import ( _has_pre_quant_scale, get_quant_config, - get_quantization_format, process_layer_quant_config, ) from modelopt.torch.export.unified_export_hf import _prepare_moe_inputs From 1e35b2915ca008f4b289372ce00b55c0453f366b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:14:33 +0000 Subject: [PATCH 12/16] Test the calibration iterators and expert-index recovery on NemotronHExperts The non-gated cases of these fused-experts tests went with the stand-in when the non-gated tests moved here. They run on the real class again: the eager forward recovers each up_proj slice's expert, and the quantizer-only iterator yields the weight iterator's quantizers in its order without indexing the fused weight; that test fails if the iterator names the gated projection. Signed-off-by: Shengliang Xu --- .../models/test_nemotron_h_modeling_ptq.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py index c284496a9f0..e175cd08077 100644 --- a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py +++ b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py @@ -111,6 +111,39 @@ def test_registered_wrapper_quantizes_up_and_down_proj(): assert set(weight_attr_names(converted)) == {"up_proj", "down_proj"} +def test_expert_index_recovery(): + """The eager forward names an up_proj slice's expert from its storage offset.""" + experts = QuantModuleRegistry.convert(_ExpertsModel("eager").experts) + for idx in range(NUM_EXPERTS): + assert experts._get_expert_idx_from_first_proj(experts.up_proj[idx]) == idx + + +def test_weight_quantizers_come_in_the_weight_iterators_order_without_indexing(): + """The quantizer-only iterator yields the weight iterator's quantizers, in its order, and never + indexes the fused weight, which under FSDP2 is a collective per expert.""" + experts = QuantModuleRegistry.convert(_ExpertsModel("eager").experts) + from_weights = [q for _, q in experts.iter_weights_for_calibration()] + + class _NoIndexing(torch.Tensor): + @staticmethod + def __new__(cls, data): + return torch.Tensor._make_subclass(cls, data, False) + + def __getitem__(self, item): + raise AssertionError("indexed the weight") + + for name in ("up_proj", "down_proj"): + weight = getattr(experts, name) + setattr(experts, name, nn.Parameter(_NoIndexing(weight.data), requires_grad=False)) + + quantizers_only = list(experts.iter_weight_quantizers_for_calibration()) + assert len(quantizers_only) == 2 * NUM_EXPERTS + assert [id(q) for q in quantizers_only] == [id(q) for q in from_weights] + # The weight iterator still slices, so the guard above is live. + with pytest.raises(AssertionError, match="indexed the weight"): + list(experts.iter_weights_for_calibration()) + + @pytest.mark.parametrize("impl", ["eager", "grouped_mm"]) def test_forward_passthrough_matches(impl): """With its quantizers disabled, the wrapper computes what the unconverted experts do.""" From fbb41418a0583023420cef70eda981bb91940f33 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:31:18 +0000 Subject: [PATCH 13/16] Drop the on-the-fly fallback eager-only experts no longer need With the non-gated wrapper gone, register_fused_experts_on_the_fly registers only the generic wrapper, which is not eager-only, so looking it up for a class not registered yet cannot change which experts are forced onto eager: those wrappers (FP8's, GPT-OSS's, Nemotron-H's) register when their modules are imported. The transposed check's comment no longer speaks of non-gated layouts. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 9976db4502f..2223608c283 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1667,9 +1667,8 @@ 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. 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. + # puts out. A gated container's first projection is `2 * intermediate`, so its shapes + # give a transposed layout away 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 @@ -1734,15 +1733,14 @@ def normalize_experts_impl_on_the_fly(model): fallback_impl = "grouped_mm" if _HAS_GROUPED_LINEAR else "eager" # A class with its own quantized wrapper -- GPT-OSS's quantizes down_proj's input by intercepting # the eager forward's bmm/matmul -- only works on the eager forward, and so does a fused-experts - # wrapper that says so (FP8's), so keep main's forced eager for such a model. The plugins run in - # no set order, so a class register_fused_experts_on_the_fly registers may not be registered yet: - # the wrapper it will get decides then. + # wrapper that says so (FP8's), so keep main's forced eager for such a model. Those register + # when their module is imported; register_fused_experts_on_the_fly, which may run after this, + # registers only the generic wrapper, which is not eager-only. eager_only = sorted( { type(module).__name__ for module in fused_experts - if (wrapper := QuantModuleRegistry.get(type(module)) or _fused_experts_wrapper(module)) - is not None + if (wrapper := QuantModuleRegistry.get(type(module))) is not None and (not issubclass(wrapper, _QuantFusedExperts) or wrapper._eager_only) } ) From 64695a50a4fd2be6d07a2394c69c98f05b5d320e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:57:33 +0000 Subject: [PATCH 14/16] Refuse Nemotron-H experts transformers marks transposed Registered explicitly, the Nemotron-H experts skip the transposed-layout check register_fused_experts_on_the_fly applies, and the wrapper slices each expert's weight as (out_dim, in_dim). A transposed layout would slice every expert wrong with nothing in the shapes to catch it, so setup now refuses experts transformers marks is_transposed. A test marks real NemotronHExperts transposed and checks the conversion raises; it fails without the change. Signed-off-by: Shengliang Xu --- modelopt/torch/models/nemotron_h/modeling_ptq.py | 12 +++++++++++- .../torch/models/test_nemotron_h_modeling_ptq.py | 9 +++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/models/nemotron_h/modeling_ptq.py b/modelopt/torch/models/nemotron_h/modeling_ptq.py index e0f51b5a28c..574c2c26a02 100644 --- a/modelopt/torch/models/nemotron_h/modeling_ptq.py +++ b/modelopt/torch/models/nemotron_h/modeling_ptq.py @@ -38,13 +38,23 @@ class _QuantNemotronHExperts(_QuantFusedExperts): quantizer live under ``up_proj_*``. Registered explicitly, these experts skip the transposed-layout check of :func:`register_fused_experts_on_the_fly`: the wrapper relies on transformers storing their weights ``(num_experts, out_dim, in_dim)``, as every supported - release does. + release does, and refuses experts transformers marks ``is_transposed``. """ _first_proj_attr = "up_proj" # Read by the export: a non-gated up_proj is exported whole, not split into gate|up halves. _is_gated = False + def _setup(self): + # The per-expert slicing reads (out_dim, in_dim); a transposed layout would slice every + # expert's weight wrong without any shape mismatch to catch it. + if getattr(self, "is_transposed", False): + raise NotImplementedError( + f"{type(self).__name__} quantizes expert weights stored (num_experts, out_dim, " + "in_dim), but these experts are marked is_transposed." + ) + super()._setup() + if NemotronHExperts not in QuantModuleRegistry: QuantModuleRegistry.register({NemotronHExperts: "hf.NemotronHExperts"})(_QuantNemotronHExperts) diff --git a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py index e175cd08077..d6b6f136991 100644 --- a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py +++ b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py @@ -111,6 +111,15 @@ def test_registered_wrapper_quantizes_up_and_down_proj(): assert set(weight_attr_names(converted)) == {"up_proj", "down_proj"} +def test_experts_marked_transposed_are_refused(): + """The wrapper slices each expert's weight as (out_dim, in_dim), with nothing in the shapes to + tell a transposed layout apart, so experts transformers marks is_transposed are refused.""" + experts = _ExpertsModel("eager").experts + experts.is_transposed = True + with pytest.raises(NotImplementedError, match="is_transposed"): + QuantModuleRegistry.convert(experts) + + def test_expert_index_recovery(): """The eager forward names an up_proj slice's expert from its storage offset.""" experts = QuantModuleRegistry.convert(_ExpertsModel("eager").experts) From 0815e7a863777f389bb85e9fd6cc61c840f0fbcf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 04:30:31 +0000 Subject: [PATCH 15/16] Refuse transposed experts in the fused-experts wrapper itself The is_transposed refusal lived in the Nemotron-H wrapper only, but the native FP8 experts wrapper is registered explicitly too, past the on-the-fly check that skips transposed containers, and reads each expert's weights and block scales as (out_dim, in_dim) just the same. The check now runs in _QuantFusedExperts._setup, before any quantizer is built, and the Nemotron-H override is gone. A test marks native FP8 experts transposed and checks their conversion raises; it and the Nemotron-H test fail without the check. The test that rebuilds an earlier release's state quantizes its experts unmarked, as that release read no flag. Signed-off-by: Shengliang Xu --- .../torch/models/nemotron_h/modeling_ptq.py | 13 ++---------- .../torch/quantization/plugins/huggingface.py | 8 +++++++ .../plugins/test_fused_experts.py | 21 +++++++++++++++---- 3 files changed, 27 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/models/nemotron_h/modeling_ptq.py b/modelopt/torch/models/nemotron_h/modeling_ptq.py index 574c2c26a02..8f49bc9e6c2 100644 --- a/modelopt/torch/models/nemotron_h/modeling_ptq.py +++ b/modelopt/torch/models/nemotron_h/modeling_ptq.py @@ -38,23 +38,14 @@ class _QuantNemotronHExperts(_QuantFusedExperts): quantizer live under ``up_proj_*``. Registered explicitly, these experts skip the transposed-layout check of :func:`register_fused_experts_on_the_fly`: the wrapper relies on transformers storing their weights ``(num_experts, out_dim, in_dim)``, as every supported - release does, and refuses experts transformers marks ``is_transposed``. + release does, and :class:`_QuantFusedExperts` refuses experts transformers marks + ``is_transposed``. """ _first_proj_attr = "up_proj" # Read by the export: a non-gated up_proj is exported whole, not split into gate|up halves. _is_gated = False - def _setup(self): - # The per-expert slicing reads (out_dim, in_dim); a transposed layout would slice every - # expert's weight wrong without any shape mismatch to catch it. - if getattr(self, "is_transposed", False): - raise NotImplementedError( - f"{type(self).__name__} quantizes expert weights stored (num_experts, out_dim, " - "in_dim), but these experts are marked is_transposed." - ) - super()._setup() - if NemotronHExperts not in QuantModuleRegistry: QuantModuleRegistry.register({NemotronHExperts: "hf.NemotronHExperts"})(_QuantNemotronHExperts) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 2223608c283..f63dc33b3b8 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -886,6 +886,14 @@ def _get_expert_idx_from_first_proj(self, weight: torch.Tensor) -> int: return idx def _setup(self): + # Each expert's weight is sliced as (out_dim, in_dim), and a transposed container would be + # sliced wrong with nothing in the shapes to catch it. register_fused_experts_on_the_fly skips + # such containers; this refuses them for the wrappers registered explicitly (FP8, Nemotron-H). + if getattr(self, "is_transposed", False): + raise NotImplementedError( + f"{type(self).__name__} quantizes expert weights stored (num_experts, out_dim, " + "in_dim), but these experts are marked is_transposed." + ) n = self.num_experts setattr(self, self._first_proj_input_quantizer_attr, TensorQuantizer()) setattr( diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 4f217657d18..defd3d5e1b1 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -80,6 +80,19 @@ from modelopt.torch.utils.moe import torch_a2a_experts_forward +def test_native_fp8_experts_marked_transposed_are_refused(): + """Native FP8 experts are registered explicitly, past the check that skips transposed + containers, and their weights and block scales are read (out_dim, in_dim): experts marked + is_transposed are refused at conversion rather than calibrated and exported wrong.""" + experts = FP8Experts( + PretrainedConfig(hidden_size=32, intermediate_size=32, num_experts=2, hidden_act="silu"), + block_size=None, + ) + experts.is_transposed = True + with pytest.raises(NotImplementedError, match="is_transposed"): + QuantModuleRegistry.convert(experts) + + def test_native_fp8_experts_dequantized_before_nvfp4_export(): """Calibrate and pack scaled FP8 gate/up and down weights, preserving their BF16 values.""" generator = torch.Generator().manual_seed(42) @@ -499,14 +512,14 @@ def test_a_state_saved_with_transposed_experts_quantized_no_longer_restores(self class _TransposedExperts(_SyntheticFusedExperts): pass - def build(): + def build(marked=True): model = _TinyMoEModel() model.moe.experts = _TransposedExperts() - model.moe.experts.is_transposed = True + model.moe.experts.is_transposed = marked return model - model = build() - # As an earlier release did: matched by shape, with the gated wrapper. + # As an earlier release did: matched by shape, with the gated wrapper, which read no flag. + model = build(marked=False) QuantModuleRegistry.register({_TransposedExperts: "test.TransposedExperts"})( _QuantFusedExperts ) From 495f0bee52f2dc2d29c5c004f5c7eb7f53a7abff Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 21:02:06 +0000 Subject: [PATCH 16/16] Name the Nemotron-H backend tests as comparisons, not matches to eager grouped_mm and eager compute the same expert products in different orders, and neither is the reference for the other. The calibration and GPTQ tests now say that the two backends agree, and the calibration test says where its down_proj tolerance comes from. Signed-off-by: Shengliang Xu --- .../torch/models/test_nemotron_h_modeling_ptq.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py index d6b6f136991..18751577938 100644 --- a/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py +++ b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py @@ -169,12 +169,13 @@ def test_forward_passthrough_matches(impl): torch.testing.assert_close(converted(*batch), reference(*batch)) -def test_calibration_matches_eager(): - """grouped_mm calibrates every expert's quantizers to what eager does. +def test_grouped_mm_and_eager_calibrate_each_expert_alike(): + """grouped_mm and eager calibrate every expert's quantizers alike. The same weights and routing give the shared input quantizers the same activations and each - per-expert weight quantizer the same slice; only down_proj's input passes through a GEMM, hence - the tolerance on its amax. + per-expert weight quantizer the same slice. Only down_proj's input comes out of a GEMM, which + the two backends compute in different orders that round differently: hence the tolerance on + its amax. """ eager, grouped = _eager_and_grouped() for model in (eager, grouped): @@ -197,8 +198,8 @@ def test_calibration_matches_eager(): torch.testing.assert_close(gq.amax, eq.amax) -def test_gptq_hessians_match_eager(monkeypatch): - """GPTQ gives each expert the Hessian of only its own tokens under grouped_mm, as eager.""" +def test_gptq_hessians_agree_across_backends(monkeypatch): + """GPTQ gives each expert the Hessian of only its own tokens, on both backends alike.""" hessians = {} free = GPTQHelper.free