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..8f49bc9e6c2 100644 --- a/modelopt/torch/models/nemotron_h/modeling_ptq.py +++ b/modelopt/torch/models/nemotron_h/modeling_ptq.py @@ -16,13 +16,41 @@ """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_*``. 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 :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 + + +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..f63dc33b3b8 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 @@ -887,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( @@ -1210,28 +1217,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 +1609,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 +1633,23 @@ 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): @@ -1688,9 +1675,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 @@ -1755,15 +1741,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) } ) 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..d6b6f136991 --- /dev/null +++ b/tests/unit/torch/models/test_nemotron_h_modeling_ptq.py @@ -0,0 +1,266 @@ +# 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"} + + +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) + 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.""" + 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 its 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) + # 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) + + 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}]" + 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") + # 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(): + """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..defd3d5e1b1 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 @@ -71,7 +70,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, @@ -82,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) @@ -192,45 +203,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 +245,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 @@ -567,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 ) @@ -733,11 +678,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 +687,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 +1371,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,21 +1389,13 @@ 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" 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 @@ -2417,236 +2353,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 # ---------------------------------------------------------------------------