diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f299e566015..90aa9a89b4c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -50,6 +50,7 @@ Changelog - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm. - The Windows FVD example rejects I3D checkpoints requiring unrestricted pickle loading. Custom TorchScript archives now require ``--weights-sha256`` with a trusted checksum; tensor state dictionaries continue to load directly. +- A ``modelopt_state`` saved by an earlier release for a Hugging Face model whose fused MoE experts are stored transposed without a quantized wrapper of their own (for example Aria's) no longer restores, since those experts are now left unquantized. The checkpoints such a state exported had wrong expert weights, so there is nothing to migrate. - Layerwise-calibration resume (``checkpoint_dir``) now loads checkpoint files with ``weights_only=True``. A checkpoint directory written by an earlier release whose layer inputs hold objects other than tensors, plain containers or Hugging Face dynamic caches can no longer be resumed; recalibrate in a fresh directory, or disable checkpointing. @@ -61,6 +62,7 @@ Changelog **Bug Fixes** - ``hf_quant_config.json`` now declares ``pre_quant_scale`` only for INT4 / W4A8 layers that have the tensor; loading INT4 / W4A8 fused-MoE checkpoints in TensorRT-LLM also needs a release whose MoE loader treats it as optional. AWQ, SmoothQuant and SVDQuant now warn when they skip quantized fused-MoE experts. +- Hugging Face fused MoE experts stored transposed, ``(num_experts, in_dim, out_dim)``, without a quantized wrapper of their own (for example Aria's) are now left unquantized with a warning. They were quantized with gate and up split along the wrong axis, so their exported weights were wrong; the Hugging Face export now refuses such a model's MoE block instead. - Calibrating a Megatron-Core MLA model now gives ``linear_q_down_proj`` and ``linear_kv_down_proj`` one per-tensor input and weight range, as vLLM and TensorRT-LLM fuse them into one GEMM with one scale per tensor. Re-calibrate existing MLA checkpoints to get the shared ranges. - ``export_mcore_gpt_to_hf`` and ``capture_kv_cache_export_state`` now write NVFP4 KV-cache ``k_scale`` / ``v_scale`` as amax / 448, like the HF exporter and as vLLM and TensorRT-LLM expect. They wrote amax / (6 * 448), which overflows in the attention kernels. Re-export Megatron checkpoints with an NVFP4 KV cache. - Preserve accumulated histogram counts when computing an entropy calibration threshold with NumPy histograms. diff --git a/modelopt/torch/models/step3p5/modeling_ptq.py b/modelopt/torch/models/step3p5/modeling_ptq.py index 9e3d40d49b8..85860affe17 100644 --- a/modelopt/torch/models/step3p5/modeling_ptq.py +++ b/modelopt/torch/models/step3p5/modeling_ptq.py @@ -137,7 +137,7 @@ def _is_expert_indexed_moe_linear(module: nn.Module) -> bool: ``weight`` of shape ``[num_experts, out_features, in_features]``, whose ``forward(x, expert_id)`` runs ``F.linear`` against the selected expert's slice. The weights therefore live on the projection submodule rather than on the expert container, - which is what :func:`_fused_experts_wrapper_class` looks for, and the module is not an + which is what :func:`_fused_experts_wrapper` looks for, and the module is not an ``nn.Linear``, so neither the fused-experts path nor the plain linear path claims it. Detection is structural rather than keyed on class names so new Step revisions are picked diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 30f16ecb8a1..7e432cce9fb 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1074,8 +1074,10 @@ def _fused_expert_projections(module: nn.Module): """Yield ``(weight_name, weight, weight_quantizers, input_quantizer)`` per fused-experts projection. ``weight`` is the 3-D ``(num_experts, out, in)`` parameter, ``weight_quantizers`` holds one - quantizer per expert, and ``input_quantizer`` is shared by the experts and runs once per expert - forward with ``module._current_expert_idx`` already set to that expert. + quantizer per expert, and ``input_quantizer`` is shared by the experts. Under eager it runs once + per expert forward with ``module._current_expert_idx`` set to that expert; under grouped_mm once + per projection, on every expert's rows, with the index cleared (``None``), so a capture keyed on + it skips that batch. """ first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj") for weight_name, quantizers_name, input_q_name in ( diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index a0b896b7903..0e6b1ad7271 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -53,6 +53,11 @@ logger = logging.getLogger(__name__) +try: + import transformers.integrations.moe as _transformers_moe +except ImportError: + _transformers_moe = None + try: import kitchen from kitchen.fa import KitchenFlashAttentionModule @@ -793,9 +798,27 @@ class _QuantFusedExperts(_QuantFunctionalMixin): Verified compatible models: Mixtral, Qwen2-MoE, Qwen3-MoE, Qwen3.5-MoE, DeepSeek-V3, Jamba, OLMoE. - Limitation: only works when ``experts_implementation="eager"`` (default). - ``batched_mm`` / ``grouped_mm`` backends use ``torch.bmm`` / - ``torch._grouped_mm`` instead of ``F.linear`` and are not intercepted. + Backends: both the ``eager`` (per-expert ``F.linear``) and ``grouped_mm`` + (``transformers.integrations.moe._grouped_linear`` → + ``torch._grouped_mm`` / ``torch.nn.functional.grouped_mm`` / fallback) + forwards are intercepted. Both calibrate the same per-expert weight quantizers + and shared input quantizers, and export the same keys; the activation scales + they calibrate can differ. Only one backend's hook fires per forward + (selected by ``config._experts_implementation``), so installing both is safe: + the ``grouped_mm`` fallback does not call ``F.linear``, which the CPU tests that + run it pin. Calibration still runs every model on ``eager`` + (``force_eager_experts_impl_on_the_fly``). + + Limitation: the ``grouped_mm`` path assumes the non-transposed weight layout + ``(num_experts, out_dim, in_dim)`` used by all standard fused experts. The + transposed layout (e.g. GptOss, which has its own quantized wrapper) is not + handled here. ``batched_mm`` (``torch.bmm``) and ``sonicmoe`` (opaque CUDA + kernel) backends are not intercepted. + + On ``grouped_mm`` the weight-input captures keyed on one expert (``local_hessian``, + ``imatrix``, ``gptq``) skip the batch, which holds every expert's rows, so those + algorithms calibrate the experts' weights without their inputs; max calibration + is unaffected. The non-gated variant (``up_proj`` instead of ``gate_up_proj``, used by NemotronH) is handled by :class:`_QuantNonGatedFusedExperts` via the @@ -812,6 +835,22 @@ class _QuantFusedExperts(_QuantFunctionalMixin): # Under expert parallelism each rank's experts see only the tokens routed to them, so calibration # shares the layer's input amax over the expert-parallel group (see max_calibrate). _share_input_amax_across_ep = True + # Whether the wrapper quantizes only on the eager per-expert forward, so a model holding it must + # calibrate there. + _eager_only = False + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + # The grouped_mm forward quantizes the stored weight and never reaches get_expert_weight, so + # a wrapper that defines its weight domain there would calibrate the wrong weights on it. + if ( + cls.get_expert_weight is not _QuantFusedExperts.get_expert_weight + and not cls._eager_only + ): + raise TypeError( + f"{cls.__name__} overrides get_expert_weight, which the grouped_mm forward bypasses; " + "set _eager_only = True so that its models calibrate on the eager forward." + ) @property def _first_proj_input_quantizer_attr(self) -> str: @@ -822,7 +861,7 @@ def _first_proj_weight_quantizers_attr(self) -> str: return f"{self._first_proj_attr}_weight_quantizers" def _get_expert_idx_from_first_proj(self, weight: torch.Tensor) -> int: - """Recover expert index from a first-projection weight slice's storage offset. + """Recover expert index from the first-projection weight slice's storage offset. When HF indexes ``[idx]``, the result is a view sharing the same underlying storage. The offset delta divided by the stride along @@ -859,12 +898,58 @@ def _setup(self): # Per forward: each projection's experts, fake-quantized together (see _quantized_weight). self._register_temp_attribute("_stacked_weights", {}) + def _quantize_grouped_weight(self, weight: torch.Tensor, quantizers: nn.ModuleList): + """Fake-quantize a fused 3-D weight per expert. + + Each expert slice passes through its own weight quantizer (collecting + ``amax`` during calibration, fake-quantizing afterwards), matching the + eager per-expert ``F.linear`` semantics. Returns the weight untouched when every + quantizer would pass its slice through unchanged, so passthrough stays exact (and cheap). + """ + # A disabled quantizer still applies its pre-quant scale and rotation, as on the eager path. + if all( + isinstance(q, TensorQuantizer) + and not q.is_enabled + and q.pre_quant_scale is None + and not q.rotate_is_enabled + for q in quantizers + ): + return weight + # One kernel launch for every expert where the format has one. Unlike the eager path, which + # without autograd quantizes only the routed experts, this one quantizes every expert either + # way, and the kernel's one output is the buffer the slice-by-slice path below fills. + stacked = _try_stacked_fake_quant(weight, quantizers) + if stacked is not None: + return stacked + if torch.is_grad_enabled(): + # Slices written into one buffer would chain a full-size gradient copy per expert into + # the backward (QAT), so stack them there. + return torch.stack([q(w) for q, w in zip(quantizers, weight, strict=True)]) + # Each slice goes into the output as it is made: a stack would hold all of them beside it. + # Strict, so every row of the output is written or the call fails. The output takes the + # dtype the stack promotes the slices to: a pre-quant scale can widen one, and narrowing + # it here would compute from other weights than with autograd. + out: torch.Tensor | None = None + for idx, (q, w) in enumerate(zip(quantizers, weight, strict=True)): + quantized = q(w) + if out is None: + out = torch.empty_like(weight, dtype=quantized.dtype) + elif quantized.dtype != out.dtype: + out = out.to(torch.promote_types(out.dtype, quantized.dtype)) + out[idx] = quantized + del quantized # not alive beside the next slice + return out + @property def functionals_to_replace(self): _orig_linear = torch.nn.functional.linear + # First-projection quantizers, by layout: gate_up_proj_* (gated) or up_proj_* (non-gated). + # Bound once; the objects are stable across the forward. + first_input_q = getattr(self, self._first_proj_input_quantizer_attr) + first_weight_qs = getattr(self, self._first_proj_weight_quantizers_attr) # The HF fused expert forward calls F.linear exactly twice per expert - # in strict alternation: first for gate_up_proj, then for down_proj. + # in strict alternation: first the up/gate_up projection, then down_proj. # forward() resets the toggle before each call to super().forward(). def _quantized_linear(input, weight, bias=None): if self._down_proj_linear: @@ -876,19 +961,56 @@ def _quantized_linear(input, weight, bias=None): else: idx = self._get_expert_idx_from_first_proj(weight) self._current_expert_idx = idx - input = getattr(self, self._first_proj_input_quantizer_attr)(input) + input = first_input_q(input) weight = self._quantized_weight( self._first_proj_attr, - getattr(self, self._first_proj_weight_quantizers_attr), + first_weight_qs, idx, input.dtype, ) self._down_proj_linear = not self._down_proj_linear return _orig_linear(input, weight, bias) - return [ - (torch.nn.functional, "linear", _quantized_linear), - ] + replacements: list[tuple] = [(torch.nn.functional, "linear", _quantized_linear)] + + # grouped_mm backend: HF's grouped_mm_experts_forward calls the module-level + # _grouped_linear exactly twice (gate_up then down) with the FULL 3-D weight + # in the same (num_experts, out, in) layout as eager F.linear, so the same + # per-expert weight quantizers / shared input quantizers apply. Only one + # backend's hook fires per forward, so patching both F.linear and + # _grouped_linear is safe. + if _transformers_moe is not None and hasattr(_transformers_moe, "_grouped_linear"): + _orig_grouped_linear = _transformers_moe._grouped_linear + + def _quantized_grouped_linear(input, weight, offs, bias=None, is_transposed=False): + if self._eager_only: + # Quantizing the stored weight here would bypass get_expert_weight, which defines + # this wrapper's weights. + raise RuntimeError( + f"{type(self).__name__} quantizes on the eager experts forward only, but its " + 'grouped_mm forward ran. Load the model with experts_implementation="eager", ' + "and without expert parallelism, whose forward runs it too." + ) + # The batch holds every expert's rows: no single expert for the eager loop's index to + # name, so a weight-input capture keyed on it skips the batch rather than credit + # one expert with all of it. + self._current_expert_idx = None + quantizers = ( + self.down_proj_weight_quantizers if self._down_proj_linear else first_weight_qs + ) + if self._down_proj_linear: + input = self.down_proj_input_quantizer(input) + else: + input = first_input_q(input) + weight = self._quantize_grouped_weight(weight, quantizers) + self._down_proj_linear = not self._down_proj_linear + return _orig_grouped_linear( + input, weight, offs, bias=bias, is_transposed=is_transposed + ) + + replacements.append((_transformers_moe, "_grouped_linear", _quantized_grouped_linear)) + + return replacements def _quantized_weight(self, name, quantizers, idx, dtype): """Expert ``idx``'s ``name`` projection, fake-quantized by its own weight quantizer. @@ -987,11 +1109,20 @@ class _QuantNonGatedFusedExperts(_QuantFusedExperts): _first_proj_attr = "up_proj" _is_gated = False + # The grouped_mm interception is validated for the gated layout only, so these keep + # calibrating on the eager forward. + _eager_only = True class _QuantFP8Experts(_QuantFusedExperts): """Scale-aware PTQ for native Transformers FP8 experts, dequantized per expert.""" + # The block scales apply in get_expert_weight, which only the eager forward reaches: transformers' + # FP8 grouped_mm forward runs its own kernels, which nothing here intercepts. So, unlike for the + # non-gated wrapper, the grouped interceptor's refusal never fires for these: what keeps them on + # eager is force_eager_experts_impl_on_the_fly, which runs on convert and on restore. + _eager_only = True + @property def _first_proj_attr(self): return "gate_up_proj" if self.has_gate else "up_proj" @@ -1365,21 +1496,22 @@ def register_sparse_moe_on_the_fly(model): ) -def _fused_experts_wrapper_class(module): - """Return the _QuantFusedExperts subclass for a fused MoE expert container, or None. +def _fused_experts_wrapper(module): + """Return the ``_QuantFusedExperts`` subclass for a fused-MoE-experts module, else None. - Two 3-D fused layouts are recognized, both requiring ``num_experts`` and a - 3-D ``down_proj`` parameter: + Detects the standardized HuggingFace transformers 5.0+ fused expert pattern: ``num_experts``, + a 3-D ``down_proj`` and a 3-D first projection, either - * gated (``_QuantFusedExperts``): a 3-D ``gate_up_proj`` fusing gate+up. Matches + * gated (:class:`_QuantFusedExperts`): a 3-D ``gate_up_proj`` fusing gate+up. Matches ``MixtralExperts``, ``Qwen2MoeExperts``, ``Qwen3MoeExperts``, ``Qwen3_5MoeExperts``, ``DeepseekV3NaiveMoe``, ``JambaExperts``, ``OlmoeExperts``, ``MiniMaxM2Experts``, ``MiniMaxM3VLExperts``, etc. * non-gated (``_QuantNonGatedFusedExperts``): a 3-D ``up_proj`` with no ``gate_proj`` and no ``gate_up_proj``. Matches NemotronH ``NemotronHExperts``. - Returns ``None`` for non-standard layouts (GptOss, GraniteMoE, - Llama4TextExperts) which have their own explicit registrations. + Matching is by shape alone, so GptOss and Llama4TextExperts -- transposed layouts with their own + explicit registrations -- match too; :func:`register_fused_experts_on_the_fly` never reaches + them while those registrations stand. ``act_fn`` is not required: these wrappers only intercept the two ``F.linear`` calls, so modules with a custom gated activation (e.g. ``MiniMaxM3VLExperts``) @@ -1404,10 +1536,13 @@ def _fused_experts_wrapper_class(module): def _is_fused_experts_module(module): """Check if a module is a fused MoE expert container compatible with _QuantFusedExperts. - See :func:`_fused_experts_wrapper_class` for the recognized layouts (gated - ``gate_up_proj`` and non-gated ``up_proj``). + 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:`force_eager_experts_impl_on_the_fly` must still put them on + eager, the only forward on which GPT-OSS's own wrapper quantizes ``down_proj``'s input (it + intercepts the eager forward's ``bmm``/``matmul``). """ - return _fused_experts_wrapper_class(module) is not None + return _fused_experts_wrapper(module) is not None def register_fused_experts_on_the_fly(model): @@ -1429,13 +1564,34 @@ def register_fused_experts_on_the_fly(model): visited_types.add(mod_type) - wrapper_cls = _fused_experts_wrapper_class(module) - if wrapper_cls is not None: + wrapper = _fused_experts_wrapper(module) + if wrapper is not None and ( + getattr(module, "is_transposed", False) + # Stored `(num_experts, out_dim, in_dim)`, the first projection takes the dim down_proj + # puts out. That separates the layouts for a gated container, whose first projection is + # `2 * intermediate`, whether it says so or not; a non-gated one's two layouts have the + # same shapes, so only its own `is_transposed` gives it away. + or getattr(module, wrapper._first_proj_attr).shape[-1] != module.down_proj.shape[1] + ): + # `_fused_experts_wrapper` matches on parameter SHAPE only, so it also claims a + # transposed `(num_experts, in_dim, out_dim)` container. The fused-experts export + # (export/moe_utils.py) then splits gate/up by slicing dim 1 at `intermediate_dim`, the + # wrong axis for that layout, and writes silently wrong weights. Leave it unquantized. + warnings.warn( + f"Fused MoE experts '{name}' ({mod_type.__name__}) store their weights transposed " + "(num_experts, in_dim, out_dim), which the generic _QuantFusedExperts export path " + "does not support: it splits gate/up along dim 1, which is the input dim for this " + "layout. Leaving them unquantized, so export_hf_checkpoint will refuse their MoE " + "block; such a class needs its own quantized wrapper, as GptOssExperts and " + "Llama4TextExperts have." + ) + continue + if wrapper is not None: print( f"\033[1mDetected fused MoE experts '{name}' of type {mod_type.__name__}, " - f"registering with {wrapper_cls.__name__}.\033[0m" + f"registering with {wrapper.__name__}.\033[0m" ) - QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(wrapper_cls) + QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(wrapper) def force_eager_experts_impl_on_the_fly(model): @@ -1443,10 +1599,12 @@ def force_eager_experts_impl_on_the_fly(model): HF transformers 5.0+ decorates fused-experts forwards with ``@use_experts_implementation``, which may dispatch to ``torch._grouped_mm`` - or ``torch.bmm`` backends. Those backends bypass ``F.linear`` and so bypass - ``_QuantFusedExperts``'s input/weight quantizer hooks — calibration silently - does nothing, no ``input_scale`` / ``amax`` is collected, and the exported - checkpoint produces garbage at inference. + or ``torch.bmm`` backends. ``batched_mm`` and ``sonicmoe`` bypass ``F.linear`` + and so bypass ``_QuantFusedExperts``'s input/weight quantizer hooks — calibration + silently does nothing, no ``input_scale`` / ``amax`` is collected, and the + exported checkpoint produces garbage at inference. ``grouped_mm`` is intercepted + (see :class:`_QuantFusedExperts`), but calibration stays on eager until the + per-expert input captures follow it there. Sets ``config._experts_implementation = "eager"`` on the model config (and recursively on ``text_config`` / ``vision_config`` / ``audio_config`` / diff --git a/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py b/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py index a0669da1c1c..156532af3cc 100644 --- a/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py +++ b/tests/gpu/torch/quantization/plugins/test_fused_experts_cuda.py @@ -23,17 +23,39 @@ import modelopt.torch.quantization.plugins.huggingface as hf_plugin from modelopt.torch.quantization import tensor_quant +# As the plugin and the unit tests treat it: grouped_mm is intercepted where transformers has it. +try: + from transformers.integrations import moe as moe_integration + + HAS_GROUPED_MM = hasattr(moe_integration, "_grouped_linear") +except ImportError: + HAS_GROUPED_MM = False + pytestmark = pytest.mark.skipif( not hasattr(tensor_quant.triton_kernel, "fp4_fake_quant_block"), reason="fp4_fake_quant_block requires compute >= 8.9", ) +@pytest.mark.parametrize( + "impl", + [ + "eager", + pytest.param( + "grouped_mm", + marks=pytest.mark.skipif( + not HAS_GROUPED_MM, reason="transformers grouped_mm backend unavailable" + ), + ), + ], +) @pytest.mark.parametrize("pass_through_bwd", [True, False]) -def test_stacked_nvfp4_experts_match_per_expert_quantizers(monkeypatch, pass_through_bwd): +def test_stacked_nvfp4_experts_match_per_expert_quantizers(monkeypatch, pass_through_bwd, impl): model = get_tiny_qwen3_moe().cuda().train() input_ids = torch.randint(0, model.config.vocab_size, (2, 16), device="cuda") mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(input_ids)) + # Train on the backend under test; both fake-quantize each projection's experts in one launch. + model.config._experts_implementation = impl experts = [m for m in model.modules() if isinstance(m, hf_plugin._QuantFusedExperts)] assert experts for module in experts: @@ -62,14 +84,30 @@ def loss_and_grads(checkpointing): "stacked_nvfp4_fake_quant", lambda *args: calls.append(args) or stacked_fake_quant(*args), ) + grouped_calls = [] + if HAS_GROUPED_MM: + grouped_linear = moe_integration._grouped_linear + monkeypatch.setattr( + moe_integration, + "_grouped_linear", + lambda *args, **kwargs: grouped_calls.append(1) or grouped_linear(*args, **kwargs), + ) results = [] for checkpointing in (False, True): calls.clear() results.append(loss_and_grads(checkpointing)) # Checkpointing reruns the forward in backward. assert len(calls) == 2 * len(experts) * (1 + checkpointing) + assert bool(grouped_calls) == (impl == "grouped_mm"), "the backend under test did not run" + # Without autograd only the grouped forward, which quantizes every expert, takes the kernel. + calls.clear() + with torch.no_grad(): + logits = model(input_ids).logits + assert len(calls) == (2 * len(experts) if impl == "grouped_mm" else 0) monkeypatch.setattr(hf_plugin, "_try_stacked_fake_quant", lambda *args: None) + with torch.no_grad(): + assert torch.equal(model(input_ids).logits, logits) ref_loss, ref_grads = loss_and_grads(checkpointing=False) for loss, grads in results: assert torch.equal(loss, ref_loss) diff --git a/tests/unit/recipe/test_qwen3_8_recipe.py b/tests/unit/recipe/test_qwen3_8_recipe.py index 36abe30e14d..576406c508b 100644 --- a/tests/unit/recipe/test_qwen3_8_recipe.py +++ b/tests/unit/recipe/test_qwen3_8_recipe.py @@ -45,7 +45,7 @@ def __init__(self): super().__init__() self.num_experts = 2 self.intermediate_dim = _H - self.gate_up_proj = nn.Parameter(torch.randn(2, _H, 2 * _H)) + self.gate_up_proj = nn.Parameter(torch.randn(2, 2 * _H, _H)) self.down_proj = nn.Parameter(torch.randn(2, _H, _H)) def forward(self, hidden_states): diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 0847227c412..7c8dac3d500 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -17,6 +17,7 @@ import copy import warnings +import weakref from contextlib import nullcontext import pytest @@ -30,6 +31,7 @@ from transformers import PretrainedConfig from transformers.integrations.finegrained_fp8 import FP8Experts +import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module from modelopt.torch.export.model_utils import _release_exported_tensors @@ -40,12 +42,18 @@ get_quantization_format, process_layer_quant_config, ) +from modelopt.torch.export.unified_export_hf import _prepare_moe_inputs +from modelopt.torch.models import is_moe +from modelopt.torch.opt.conversion import ApplyModeError from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name -from modelopt.torch.quantization.model_calib import local_hessian_calibrate +from modelopt.torch.quantization.model_calib import ( + _register_weight_input_hooks, + local_hessian_calibrate, +) from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.plugins.huggingface import ( - _fused_experts_wrapper_class, + _fused_experts_wrapper, _is_fused_experts_module, _is_sparse_sequaential_moe_block, _QuantFusedExperts, @@ -54,7 +62,7 @@ register_fused_experts_on_the_fly, register_sparse_moe_on_the_fly, ) -from modelopt.torch.quantization.utils.core_utils import weight_attr_names +from modelopt.torch.quantization.utils.core_utils import module_name_maps, weight_attr_names def test_native_fp8_experts_dequantized_before_nvfp4_export(): @@ -285,6 +293,116 @@ def _route_once_to_each_expert(model): model.moe.experts(hidden_states, top_k_index, top_k_weights) +# --------------------------------------------------------------------------- +# grouped_mm backend fixtures: a synthetic fused-experts module that can run +# either the eager (F.linear) forward or the real transformers grouped_mm +# forward, selected by ``config._experts_implementation``. Used to verify +# grouped_mm calibration matches eager and that local-Hessian capture works in +# grouped mode. The class avoids the ``@use_experts_implementation`` decorator so +# it imports on transformers builds that predate the grouped_mm backend. +# --------------------------------------------------------------------------- +try: + import transformers.integrations.moe as _moe_integration + + HAS_GROUPED_MM = hasattr(_moe_integration, "_grouped_linear") +except ImportError: + HAS_GROUPED_MM = False + + +class _ExpertsConfig: + """Minimal config carrying the experts-implementation selector.""" + + def __init__(self, impl="grouped_mm"): + self._experts_implementation = impl + self.hidden_act = "silu" + + +class _SyntheticGroupedExperts(nn.Module): + """Fused experts dispatching to eager or the real grouped_mm forward. + + Mirrors the HF transformers 5.x fused-experts contract + (``has_gate``/``has_bias``/``is_transposed``/``_apply_gate`` + 3-D + ``gate_up_proj``/``down_proj``) that ``grouped_mm_experts_forward`` expects. + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.num_experts = NUM_EXPERTS + self.hidden_dim = HIDDEN_DIM + self.intermediate_dim = INTERMEDIATE_DIM + self.has_gate = True + self.has_bias = False + self.is_transposed = False + self.is_concatenated = True + # What transformers' @use_experts_implementation sets (5.18+ reads it in grouped_mm). + self._is_expert_parallel = False + self.gate_up_proj = nn.Parameter( + torch.randn(NUM_EXPERTS, 2 * INTERMEDIATE_DIM, HIDDEN_DIM) * 0.02 + ) + self.down_proj = nn.Parameter(torch.randn(NUM_EXPERTS, HIDDEN_DIM, INTERMEDIATE_DIM) * 0.02) + self.act_fn = nn.SiLU() + + def _apply_gate(self, gate_up_out): + gate, up = gate_up_out.chunk(2, dim=-1) + return self.act_fn(gate) * up + + def _eager_forward(self, hidden_states, top_k_index, top_k_weights): + final_hidden_states = torch.zeros_like(hidden_states) + with torch.no_grad(): + expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + for expert_idx in expert_hit: + expert_idx = expert_idx[0] + if expert_idx == self.num_experts: + continue + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + current_state = hidden_states[token_idx] + gate, up = F.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1) + current_hidden_states = self.act_fn(gate) * up + current_hidden_states = F.linear(current_hidden_states, self.down_proj[expert_idx]) + current_hidden_states = ( + current_hidden_states * top_k_weights[token_idx, top_k_pos, None] + ) + final_hidden_states.index_add_( + 0, token_idx, current_hidden_states.to(final_hidden_states.dtype) + ) + return final_hidden_states + + def forward(self, hidden_states, top_k_index, top_k_weights): + if self.config._experts_implementation == "grouped_mm": + return _moe_integration.grouped_mm_experts_forward( + self, hidden_states, top_k_index, top_k_weights + ) + return self._eager_forward(hidden_states, top_k_index, top_k_weights) + + +class _GroupedSparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.gate = _SyntheticTopKRouter() + self.experts = _SyntheticGroupedExperts(config) + + def forward(self, hidden_states): + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + _, top_k_weights, top_k_index = self.gate(hidden_states) + hidden_states = self.experts(hidden_states, top_k_index, top_k_weights) + return hidden_states.reshape(batch_size, sequence_length, hidden_dim) + + +class _GroupedMoEModel(nn.Module): + """Minimal MoE model whose experts run the selected backend.""" + + def __init__(self, impl): + super().__init__() + self.config = _ExpertsConfig(impl=impl) + self.moe = _GroupedSparseMoeBlock(self.config) + + def forward(self, x): + return self.moe(x) + + # --------------------------------------------------------------------------- # Tests for _is_fused_experts_module # --------------------------------------------------------------------------- @@ -355,6 +473,100 @@ def test_fused_experts_only_registration(self): assert QuantModuleRegistry.get(block_type) is None self._cleanup_registry(expert_type) + @pytest.mark.parametrize("declared", [True, False], ids=["declared", "by-shape"]) + def test_transposed_fused_experts_are_left_unquantized(self, declared): + """Matched by shape, a transposed container would export with gate/up split along its input + dim; registration leaves it unquantized, with a warning, rather than let export write wrong + weights or abort the rest of the model's quantization. A gated container's shapes give it + away when it does not say so.""" + + class _TransposedExperts(_SyntheticFusedExperts): + pass + + experts = _TransposedExperts() + if declared: + experts.is_transposed = True + else: + for name in ("gate_up_proj", "down_proj"): + weight = getattr(experts, name).detach().transpose(1, 2).contiguous() + setattr(experts, name, nn.Parameter(weight)) + model = nn.Module() + model.experts = experts + + with pytest.warns(UserWarning, match="store their weights transposed"): + register_fused_experts_on_the_fly(model) + assert QuantModuleRegistry.get(_TransposedExperts) is None + + def test_a_transposed_containers_moe_block_is_refused_at_export(self): + """Left unquantized, a transposed container reaches the export inside an MoE block it + recognizes, with no handler for its experts, so the export refuses the block rather than + write it.""" + + class _TransposedExperts(_SyntheticFusedExperts): + pass + + model = _TinyMoEModel() + model.moe.experts = _TransposedExperts() + model.moe.experts.is_transposed = True + with pytest.warns(UserWarning, match="store their weights transposed"): + register_fused_experts_on_the_fly(model) + assert is_moe(model.moe) + + with pytest.raises(NotImplementedError): + _prepare_moe_inputs(model, torch.float16, is_modelopt_qlora=False) + + def test_a_state_saved_with_transposed_experts_quantized_no_longer_restores(self): + """An earlier release quantized such a container, wrongly. Restoring that state now finds its + experts left unquantized and refuses the quantizers saved for them.""" + + class _TransposedExperts(_SyntheticFusedExperts): + pass + + def build(): + model = _TinyMoEModel() + model.moe.experts = _TransposedExperts() + model.moe.experts.is_transposed = True + return model + + model = build() + # As an earlier release did: matched by shape, with the gated wrapper. + QuantModuleRegistry.register({_TransposedExperts: "test.TransposedExperts"})( + _QuantFusedExperts + ) + try: + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(torch.randn(2, 4, HIDDEN_DIM))) + state = mto.modelopt_state(model) + finally: + self._cleanup_registry(_TransposedExperts) + + with ( + pytest.warns(UserWarning, match="store their weights transposed"), + pytest.raises(ApplyModeError), + ): + mto.restore_from_modelopt_state(build(), state) + + def test_a_wrapper_that_defines_its_weights_in_get_expert_weight_is_eager_only(self): + """The grouped_mm forward quantizes the stored weight without get_expert_weight, so a wrapper + that defines its weight domain there must calibrate on the eager forward. One that does not + say so is refused when it is defined, rather than calibrating the stored weights.""" + with pytest.raises(TypeError, match="_eager_only"): + + class _Scaled(_QuantFusedExperts): + def get_expert_weight(self, name, idx, dtype=None): + return super().get_expert_weight(name, idx, dtype) + + class _ScaledOnEager(_QuantFusedExperts): + _eager_only = True + + def get_expert_weight(self, name, idx, dtype=None): + return super().get_expert_weight(name, idx, dtype) + + # Inheriting the override does not take it off the eager forward either. + with pytest.raises(TypeError, match="_eager_only"): + + class _ScaledOnGroupedMM(_ScaledOnEager): + _eager_only = False + def test_convert_creates_quantizers(self): """After conversion, fused experts should have shared input and per-expert weight quantizers.""" model = _TinyMoEModel() @@ -1595,6 +1807,304 @@ def test_unrelated_dotted_number_unchanged(self): ) +# --------------------------------------------------------------------------- +# Tests for the grouped_mm backend: _QuantFusedExperts intercepts the real +# transformers grouped_mm forward (via _grouped_linear), and grouped calibration +# must produce the same per-expert quantizer state as the eager F.linear path. +# --------------------------------------------------------------------------- +@pytest.mark.skipif(not HAS_GROUPED_MM, reason="transformers grouped_mm backend unavailable") +class TestGroupedMMExperts: + @staticmethod + def _cleanup_registry(mod_type): + if QuantModuleRegistry.get(mod_type) is not None: + QuantModuleRegistry.unregister(mod_type) + + @staticmethod + def _quant_cfg(): + return { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*gate_up_proj_input_quantizer", + "cfg": {"num_bits": 8, "axis": None}, + }, + { + "quantizer_name": "*down_proj_input_quantizer", + "cfg": {"num_bits": 8, "axis": None}, + }, + { + "quantizer_name": "*gate_up_proj_weight_quantizer", + "cfg": {"num_bits": 8, "axis": 0}, + }, + { + "quantizer_name": "*down_proj_weight_quantizer", + "cfg": {"num_bits": 8, "axis": 0}, + }, + ], + "algorithm": "max", + } + + @staticmethod + def _forward_loop_all_experts(m): + # Route every token to all experts (round-robin) so eager and grouped both + # calibrate every expert's weight quantizer — eager only sees routed experts. + torch.manual_seed(0) + seq = 16 + hidden = torch.randn(seq, HIDDEN_DIM) + top_k_index = torch.stack( + [torch.arange(seq) % NUM_EXPERTS, (torch.arange(seq) + 1) % NUM_EXPERTS], dim=1 + ) + top_k_weights = torch.softmax(torch.randn(seq, TOP_K), dim=1) + with torch.no_grad(): + m.moe.experts(hidden, top_k_index, top_k_weights) + + def test_forward_passthrough_matches(self): + """Converted grouped experts with their quantizers disabled match the unconverted forward.""" + self._cleanup_registry(_SyntheticGroupedExperts) + torch.manual_seed(0) + model = _GroupedMoEModel("grouped_mm") + ref = _SyntheticGroupedExperts(_ExpertsConfig("grouped_mm")) + ref.load_state_dict(model.moe.experts.state_dict()) + + try: + register_fused_experts_on_the_fly(model) + converted = QuantModuleRegistry.convert(model.moe.experts) + for q in converted.modules(): + if isinstance(q, TensorQuantizer): + q.disable() + + seq = 8 + hidden = torch.randn(seq, HIDDEN_DIM) + top_k_index = torch.randint(0, NUM_EXPERTS, (seq, TOP_K)) + top_k_weights = torch.softmax(torch.randn(seq, TOP_K), dim=1) + with torch.no_grad(): + out_ref = ref(hidden, top_k_index, top_k_weights) + out_test = converted(hidden, top_k_index, top_k_weights) + assert torch.allclose(out_ref, out_test, atol=1e-6), ( + f"Max diff: {(out_ref - out_test).abs().max().item()}" + ) + finally: + self._cleanup_registry(_SyntheticGroupedExperts) + + def test_calibration_matches_eager(self): + """grouped_mm PTQ calibration yields the same quantizer state as eager PTQ. + + Identical weights + identical routing ⇒ the shared input quantizers see the + same activation set and each per-expert weight quantizer sees the same slice, + so amax must match (max abs is order-free; only the down-proj input flows + through a GEMM, hence the small tolerance there). + """ + self._cleanup_registry(_SyntheticGroupedExperts) + torch.manual_seed(0) + eager_model = _GroupedMoEModel("eager") + grouped_model = _GroupedMoEModel("grouped_mm") + grouped_model.load_state_dict(eager_model.state_dict()) + + try: + for model in (eager_model, grouped_model): + impl = model.config._experts_implementation + # quantize() calibrates every model on eager; calibrate on the backend under test. + mtq.quantize(model, self._quant_cfg()) + model.config._experts_implementation = impl + mtq.calibrate(model, "max", self._forward_loop_all_experts) + + e = eager_model.moe.experts + g = grouped_model.moe.experts + + # Shared input quantizers (down flows through a GEMM -> small tolerance). + for q_name in ("gate_up_proj_input_quantizer", "down_proj_input_quantizer"): + ea = getattr(e, q_name).amax + ga = getattr(g, q_name).amax + assert ea is not None and ga is not None, f"{q_name} not calibrated" + assert torch.allclose(ea, ga, atol=1e-3, rtol=1e-2), ( + f"{q_name} amax mismatch: eager={ea} grouped={ga}" + ) + + # Per-expert weight amax (weights identical -> tight match), all experts calibrated. + for idx in range(NUM_EXPERTS): + for q_list in ("gate_up_proj_weight_quantizers", "down_proj_weight_quantizers"): + eq = getattr(e, q_list)[idx].amax + gq = getattr(g, q_list)[idx].amax + assert eq is not None and gq is not None, f"{q_list}[{idx}] not calibrated" + assert torch.allclose(eq, gq, atol=1e-5), ( + f"{q_list}[{idx}] weight amax mismatch" + ) + finally: + self._cleanup_registry(_SyntheticGroupedExperts) + + def test_fake_quantizing_a_projection_keeps_one_experts_slice_alive(self): + """Without autograd each expert's fake-quantized slice goes into the output as it is made: + a projection's slices all alive beside the output would double its memory.""" + model = _GroupedMoEModel("grouped_mm") + try: + mtq.quantize(model, self._quant_cfg(), self._forward_loop_all_experts) + finally: + self._cleanup_registry(_SyntheticGroupedExperts) + experts = model.moe.experts + weight, quantizers = experts.gate_up_proj, experts.gate_up_proj_weight_quantizers + made = [] + + def one_alive(_quantizer, _args, output): + assert all(ref() is None for ref in made), "an earlier expert's slice is still alive" + made.append(weakref.ref(output)) + + handles = [q.register_forward_hook(one_alive) for q in quantizers] + try: + with torch.no_grad(): + out = experts._quantize_grouped_weight(weight, quantizers) + finally: + for handle in handles: + handle.remove() + + assert len(made) == NUM_EXPERTS + with torch.no_grad(): + expected = torch.stack([q(weight[idx]) for idx, q in enumerate(quantizers)]) + torch.testing.assert_close(out, expected, rtol=0, atol=0) + + @pytest.mark.parametrize("transform", ["pre_quant_scale", "rotate"]) + def test_disabled_quantizers_keep_their_pre_quant_scale_and_rotation( + self, monkeypatch, transform + ): + """A disabled weight quantizer still applies its pre-quant scale and rotation, as on the eager + path, so the grouped forward does not pass the stored weight through unchanged.""" + # A stand-in rotation that changes any weight it is applied to. + monkeypatch.setattr( + tensor_quantizer_module, + "normalized_hadamard_transform", + lambda inputs, rotate_fp32=False, block_size=None: inputs.flip(-1), + ) + self._cleanup_registry(_SyntheticGroupedExperts) + torch.manual_seed(0) + models = [_GroupedMoEModel("eager"), _GroupedMoEModel("grouped_mm")] + models[1].load_state_dict(models[0].state_dict()) + scale = torch.rand(HIDDEN_DIM) + 0.5 + hidden = torch.randn(8, HIDDEN_DIM) + top_k_index = torch.randint(0, NUM_EXPERTS, (8, TOP_K)) + top_k_weights = torch.softmax(torch.randn(8, TOP_K), dim=1) + outputs = [] + try: + register_fused_experts_on_the_fly(models[0]) + for model in models: + experts = QuantModuleRegistry.convert(model.moe.experts) + for q in experts.modules(): + if isinstance(q, TensorQuantizer): + q.disable() + for q in experts.gate_up_proj_weight_quantizers: + if transform == "rotate": + q.set_from_attribute_config({"rotate": True}) + else: + q.pre_quant_scale = scale + with torch.no_grad(): + outputs.append(experts(hidden, top_k_index, top_k_weights)) + finally: + self._cleanup_registry(_SyntheticGroupedExperts) + torch.testing.assert_close(outputs[1], outputs[0], rtol=1e-4, atol=1e-6) + + @pytest.mark.parametrize("widened", ["every-expert", "a-later-expert"]) + def test_a_widening_pre_quant_scale_gives_one_grouped_forward_in_both_grad_modes(self, widened): + """FP32 pre-quant scales on a BF16 model widen what the quantizers return, on every expert or + on one after the first. The grouped forward computes from the widened slices with or without + autograd, rather than rounding them to BF16 only without it.""" + self._cleanup_registry(_SyntheticGroupedExperts) + torch.manual_seed(0) + model = _GroupedMoEModel("grouped_mm").to(torch.bfloat16) + hidden = torch.randn(8, HIDDEN_DIM, dtype=torch.bfloat16) + top_k_index = torch.randint(0, NUM_EXPERTS, (8, TOP_K)) + top_k_weights = torch.softmax(torch.randn(8, TOP_K), dim=1).to(torch.bfloat16) + outputs = [] + try: + register_fused_experts_on_the_fly(model) + experts = QuantModuleRegistry.convert(model.moe.experts) + for q in experts.modules(): + if isinstance(q, TensorQuantizer): + q.disable() + for name in ("gate_up_proj", "down_proj"): + in_dim = getattr(experts, name).shape[-1] + getattr(experts, f"{name}_input_quantizer").pre_quant_scale = ( + torch.rand(in_dim) + 0.5 + ) + for idx, q in enumerate(getattr(experts, f"{name}_weight_quantizers")): + if widened == "every-expert" or idx == 1: + q.pre_quant_scale = torch.rand(in_dim) + 0.5 + for grad in (True, False): + with torch.set_grad_enabled(grad): + outputs.append(experts(hidden, top_k_index, top_k_weights).detach()) + finally: + self._cleanup_registry(_SyntheticGroupedExperts) + torch.testing.assert_close(outputs[1], outputs[0], rtol=0, atol=0) + + def test_fake_quantized_grouped_forward_and_backward_match_eager(self): + """Under autograd (QAT) the grouped forward fake-quantizes every expert's slice at once; its + output and gradients match the eager forward's on the same calibrated quantizers.""" + self._cleanup_registry(_SyntheticGroupedExperts) + torch.manual_seed(0) + models = [_GroupedMoEModel("eager"), _GroupedMoEModel("grouped_mm")] + models[1].load_state_dict(models[0].state_dict()) + try: + # quantize() calibrates every model on eager, so both get the same quantizers. + for model in models: + mtq.quantize(model, self._quant_cfg(), self._forward_loop_all_experts) + models[1].config._experts_implementation = "grouped_mm" + finally: + self._cleanup_registry(_SyntheticGroupedExperts) + hidden = torch.randn(16, HIDDEN_DIM) + top_k_index = torch.stack( + [torch.arange(16) % NUM_EXPERTS, (torch.arange(16) + 1) % NUM_EXPERTS], 1 + ) + top_k_weights = torch.softmax(torch.randn(16, TOP_K), dim=1) + results = [] + for model in models: + x = hidden.clone().requires_grad_() + out = model.moe.experts(x, top_k_index, top_k_weights) + out.square().sum().backward() + experts = model.moe.experts + results.append((out, x.grad, experts.gate_up_proj.grad, experts.down_proj.grad)) + for grouped, eager in zip(results[1], results[0]): + torch.testing.assert_close(grouped, eager, rtol=1e-4, atol=1e-6) + + def test_weight_input_captures_skip_the_grouped_forward(self): + """The weight-input captures key on the eager loop's expert index, which names no expert in the + grouped forward's batch of every expert's rows: they skip it rather than credit one expert + with all of it.""" + model = _GroupedMoEModel("eager") + rows = {} + + def capture(quantizer, _weight, inputs): + rows[id(quantizer)] = rows.get(id(quantizer), 0) + inputs.shape[0] + + try: + mtq.quantize(model, self._quant_cfg(), self._forward_loop_all_experts) + handles = _register_weight_input_hooks(model, module_name_maps(model), capture) + model.config._experts_implementation = "grouped_mm" + try: + self._forward_loop_all_experts(model) + finally: + for handle in handles: + handle.remove() + finally: + self._cleanup_registry(_SyntheticGroupedExperts) + assert rows == {} + + def test_an_eager_only_wrapper_refuses_the_grouped_forward(self): + """A wrapper whose weights are defined on the eager forward raises if the grouped one runs, + rather than quantize the stored weight.""" + + class _EagerOnlyExperts(_QuantFusedExperts): + _eager_only = True + + self._cleanup_registry(_SyntheticGroupedExperts) + QuantModuleRegistry.register({_SyntheticGroupedExperts: "test.SyntheticGroupedExperts"})( + _EagerOnlyExperts + ) + try: + model = _GroupedMoEModel("grouped_mm") + QuantModuleRegistry.convert(model.moe.experts) + with pytest.raises(RuntimeError, match="eager experts forward only"): + self._forward_loop_all_experts(model) + finally: + self._cleanup_registry(_SyntheticGroupedExperts) + + # --------------------------------------------------------------------------- # Tests for the non-gated fused-experts path (NemotronH NemotronHExperts): # single up_proj (no gate half) + down_proj. Quantizers are named after the @@ -1609,10 +2119,10 @@ def _cleanup_registry(mod_type): def test_detected_and_picks_nongated_wrapper(self): module = _SyntheticNonGatedFusedExperts() assert _is_fused_experts_module(module) is True - assert _fused_experts_wrapper_class(module) is _QuantNonGatedFusedExperts + assert _fused_experts_wrapper(module) is _QuantNonGatedFusedExperts def test_gated_still_picks_base_wrapper(self): - assert _fused_experts_wrapper_class(_SyntheticFusedExperts()) is _QuantFusedExperts + assert _fused_experts_wrapper(_SyntheticFusedExperts()) is _QuantFusedExperts def test_register_uses_nongated_wrapper(self): model = _TinyNonGatedMoEModel() @@ -1800,7 +2310,7 @@ def __init__(self): self.act_fn = nn.SiLU() module = _SplitGatedExperts() - assert _fused_experts_wrapper_class(module) is None + assert _fused_experts_wrapper(module) is None assert _is_fused_experts_module(module) is False def test_get_quant_config_resolves_nongated_experts(self): @@ -1885,7 +2395,7 @@ def test_left_to_generic_fused_experts_wrapper(self): """The experts must be left to the generic fused-experts wrapper.""" # A static registration would shadow on-the-fly detection. assert QuantModuleRegistry.get(self._experts_type()) is None - assert _fused_experts_wrapper_class(_make_qwen3_vl_moe_experts()) is _QuantFusedExperts + assert _fused_experts_wrapper(_make_qwen3_vl_moe_experts()) is _QuantFusedExperts def test_convert_and_forward_matches_reference(self): """Conversion must succeed and stay numerically transparent before calibration.""" diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py index 9b06862e56c..a22703fc91e 100644 --- a/tests/unit/torch/quantization/test_utils.py +++ b/tests/unit/torch/quantization/test_utils.py @@ -24,6 +24,7 @@ from modelopt.torch.quantization.utils import ( convert_quantization_axis_to_reduce_axis, + core_utils, reduce_amax, reduce_block_amax, ) @@ -31,6 +32,7 @@ _disable_fsdp_unshard_reshard, _fsdp2_owned_params, _fsdp2_owns, + _fsdp2_param_groups, fsdp2_weight_access_and_writeback_context, persistent_materialization, ) @@ -415,11 +417,11 @@ def _fsdp2_owns_its_params_in_every_state_job(rank, size): finally: layer.reshard() # The set is kept for the unit: opening the window again costs the module's parameters only. - with mock.patch.object(fully_shard, "state", wraps=fully_shard.state) as state: + with mock.patch.object(core_utils, "_fsdp2_param_groups", wraps=_fsdp2_param_groups) as scans: for _ in range(3): with fsdp2_weight_access_and_writeback_context(model.dense, model): pass - assert state.call_count == 0 + assert scans.call_count == 0 def test_fsdp2_owns_its_parameters_in_every_shard_state(cpu_dist_workers): @@ -447,11 +449,11 @@ def _window_leaves_what_fsdp2_does_not_own_job(rank, size): assert weight.to_local().shape[0] == 2 # Every expert weight misses the kept set, and one scan of the unit answers them all. with ( - mock.patch.object(fully_shard, "state", wraps=fully_shard.state) as state, + mock.patch.object(core_utils, "_fsdp2_param_groups", wraps=_fsdp2_param_groups) as scans, fsdp2_weight_access_and_writeback_context(layer, layer), ): pass - assert state.call_count == 1 + assert scans.call_count == 1 assert all(weight.to_local().shape[0] == 2 for weight in layer.experts.parameters())