diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 371bea5f656..89fb2a8b8d7 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -55,6 +55,8 @@ Changelog - ``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. +- ``mtq.quantize`` and ``mto.restore`` no longer force a Hugging Face model's fused MoE experts onto the ``eager`` backend: calibration and every later forward (QAT/QAD training, AutoQuantize scoring, fake-quant evaluation) run on the ``grouped_mm`` backend Transformers selects by default, so calibrated activation scales (some of Qwen3-30B-A3B's moved by up to about 10%), step time and memory can differ from earlier releases. Load the model with ``experts_implementation="eager"`` (``--experts_implementation eager`` in ``hf_ptq.py``) to run as earlier releases did. +- ``mtq.calibrate`` now raises an error, naming the modules, when a Hugging Face fused-experts forward reaches neither quantizer hook (a hand-written ``bmm`` forward that transformers cannot redirect, say), instead of exporting weight-derived input scales with a warning. - 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. diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index cca908f4553..0ef0c4fa389 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -600,6 +600,7 @@ def get_model( trust_remote_code=False, use_seq_device_map=False, attn_implementation=None, + experts_implementation=None, offload_folder=None, max_cpu_memory_gb=None, max_gpu_memory_gb=None, @@ -641,6 +642,11 @@ def get_model( # Note: Forcibly converting the model precision between bf16 and fp16 may introduce accuracy drop model_kwargs = config_kwargs.copy() model_kwargs.setdefault("dtype", "auto") + # Only when asked, on every load below: an explicit None is a value transformers stores. + experts_kwargs = ( + {} if experts_implementation is None else {"experts_implementation": experts_implementation} + ) + model_kwargs.update(experts_kwargs) # DiffusionGemma ties encoder/decoder weights. device_map "auto" (balanced) can split # a tied pair across GPUs, leaving one side on the meta device and breaking generation. @@ -720,6 +726,7 @@ def has_pack_quantized_config(config): device_map="auto", trust_remote_code=trust_remote_code, dtype="auto", + **experts_kwargs, ) elif get_original_hf_quant_method(hf_config) == "mxfp4": # Native MXFP4 checkpoints (e.g. openai/gpt-oss-*) must be dequantized to diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 85bb980ea38..47d6a8b2c19 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -323,6 +323,11 @@ def load_model(args: argparse.Namespace): hf_config = AutoConfig.from_pretrained( args.local_checkpoint_path, trust_remote_code=args.trust_remote_code ) + if args.experts_implementation is not None: + # Set on the config, as build_meta_causal_lm sets a pre-fetched config's + # attn_implementation: the meta model is built from it, with no kwarg of its own. The + # setter carries it to the sub-configs, such as the text config a VLM's experts read. + hf_config._experts_implementation = args.experts_implementation validate_fsdp2_supported(args, hf_config) full_model = parallel_load_and_prepare_fsdp2( args.local_checkpoint_path, @@ -342,6 +347,7 @@ def load_model(args: argparse.Namespace): trust_remote_code=args.trust_remote_code, use_seq_device_map=args.use_seq_device_map, attn_implementation=args.attn_implementation, + experts_implementation=args.experts_implementation, offload_folder=args.offload_folder, max_cpu_memory_gb=args.max_cpu_memory_gb, max_gpu_memory_gb=args.max_gpu_memory_gb, @@ -367,6 +373,8 @@ def load_model(args: argparse.Namespace): model_kwargs = {"trust_remote_code": args.trust_remote_code} if args.attn_implementation is not None: model_kwargs["attn_implementation"] = args.attn_implementation + if args.experts_implementation is not None: + model_kwargs["experts_implementation"] = args.experts_implementation full_model = load_pretrained( args.local_checkpoint_path, AutoModelForCausalLM, @@ -1346,6 +1354,17 @@ def parse_args() -> argparse.Namespace: default=None, type=str, ) + parser.add_argument( + "--experts_implementation", + help=( + "Fused MoE experts backend to load the model with, passed to the HF model loading as " + "experts_implementation. By default transformers picks grouped_mm where the model " + "supports it; eager calibrates exactly as earlier releases did. Quantization redirects " + "every other backend to one of these two." + ), + default=None, + choices=["eager", "grouped_mm"], + ) parser.add_argument( "--auto_quantize_checkpoint", type=str, diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 3a210a3545a..756b8d96c7d 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1075,9 +1075,8 @@ def _fused_expert_projections(module: nn.Module): ``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. 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. + per expert forward with ``module._current_expert_idx`` set to that expert; under grouped_mm + and torch_a2a once per projection, on a batch sorted by expert (``module._grouped_offs``). """ first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj") for weight_name, quantizers_name, input_q_name in ( @@ -1131,6 +1130,31 @@ def _refuse_packed_fused_expert_weights(name, module, method, updates): ) +def _expert_slices(offs: torch.Tensor, num_experts: int, rows: int): + """``(expert, start, end)`` for the rows of each expert in a batch sorted by expert. + + ``offs`` is the cumulative token count per expert (``_grouped_offs``): exclusive ends, one per + expert, as transformers' grouped_mm forward and ``a2a_moe`` build them. An expert without rows is + skipped, and rows beyond ``offs[-1]`` (expert-parallel sentinels) belong to no expert. Offsets of + another form raise (the inclusive ``[0, c0, c0 + c1, ...]`` form, which carries one end too many, + or ends out of order), since they would credit experts with other experts' rows. + """ + ends = offs.tolist() + if ( + len(ends) != num_experts + or (ends and ends[-1] > rows) + or any(end < prev for prev, end in zip([0, *ends], ends)) + ): + raise RuntimeError( + f"Grouped experts offsets {ends} do not split {rows} rows over {num_experts} experts." + ) + start = 0 + for idx, end in enumerate(ends): + if end > start: + yield idx, start, end + start = end + + class _WeightInputHook: """Identify active input collectors so auxiliary modules can request a calibration forward.""" @@ -1146,23 +1170,41 @@ def _register_weight_input_hooks(model, names, capture, on_weight=None): ``capture(weight_quantizer, weight, inputs)`` runs on every forward of a weight whose quantizer is enabled: dense quantized linears hook the layer input; HF fused-MoE experts hook the shared - input quantizers, keyed by the active expert (``_current_expert_idx``). Weights without a hook - (conv, SequentialQuantizer, non-eager experts) are never captured. ``on_weight(name, weight, - weight_quantizer)``, if given, is called once per hooked weight. Returns removable handles. + input quantizers, keyed by the active expert (``_current_expert_idx``) when eager, or split per + expert by ``_grouped_offs`` under ``grouped_mm``. Weights without a hook (conv, + SequentialQuantizer) are never captured. ``on_weight(name, weight, weight_quantizer)``, if + given, is called once per hooked weight. Returns removable handles. """ handles: list = [] def _make_expert_hook(expert_module, weight_name, quantizers, enabled): + """Per-expert capture for whichever forward ran. + + The eager forward feeds the input quantizer one call per expert (``_current_expert_idx``); + the grouped ones (``grouped_mm``, ``torch_a2a``) feed it one expert-sorted batch per + projection, which ``_grouped_offs`` (set by the grouped interceptor, cleared by the eager one) + splits back per expert. Rows beyond ``offs[-1]`` (EP sentinels) belong to no expert. + """ + def _expert_hook(_input_quantizer, args): if not args: return - idx = expert_module._current_expert_idx - if idx in enabled: - # Read the weight fresh (valid under accelerate/FSDP re-materialization), indexed by - # this rank's own experts: an expert-parallel weight is a DTensor whose index names a - # global expert (see ``_local_experts``). - weight = expert_module._local_experts(getattr(expert_module, weight_name)) - capture(quantizers[idx], weight[idx], args[0]) + # Read the weight fresh (valid under accelerate/FSDP re-materialization), indexed by this + # rank's own experts: an expert-parallel weight is a DTensor whose index names a global + # expert (see ``_local_experts``). + offs = expert_module._grouped_offs + if offs is None: + idx = expert_module._current_expert_idx + if idx in enabled: + weight = expert_module._local_experts(getattr(expert_module, weight_name)) + capture(quantizers[idx], weight[idx], args[0]) + return + weight = expert_module._local_experts(getattr(expert_module, weight_name)) + for idx, start, end in _expert_slices( + offs, expert_module.num_experts, args[0].shape[0] + ): + if idx in enabled: + capture(quantizers[idx], weight[idx], args[0][start:end]) return _expert_hook @@ -1198,6 +1240,7 @@ def _dense_hook(linear, args): enabled = {i for i, q in enumerate(quantizers) if q.is_enabled} if not enabled: continue + module._weight_input_capture = True handles.append( input_quantizer.register_forward_pre_hook( _WeightInputHook( @@ -2588,6 +2631,16 @@ def _make_gptq_handle(name, m, weight_quantizer=None, weight=None, log_mse=True) def _make_expert_hook(experts, per_expert): def _expert_hook(_input_quantizer, _args, output): + if experts._grouped_offs is not None: + # grouped_mm and torch_a2a quantize one expert-sorted batch per projection: split + # it back per expert by the cumulative token counts, as the eager loop feeds it. + for idx, start, end in _expert_slices( + experts._grouped_offs, experts.num_experts, output.shape[0] + ): + handle = per_expert.get(idx) + if handle is not None: + handle.accumulate(output[start:end]) + return handle = per_expert.get(experts._current_expert_idx) if handle is not None: handle.accumulate(output) @@ -2611,6 +2664,7 @@ def _expert_hook(_input_quantizer, _args, output): if q.is_enabled } if per_expert: + module._weight_input_capture = True expert_handles.update((id(quantizers[i]), h) for i, h in per_expert.items()) hooks.append( input_quantizer.register_forward_hook(_make_expert_hook(module, per_expert)) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 0a69ae9ede9..87c8677985c 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -18,6 +18,7 @@ import fnmatch import inspect import os +import sys import warnings from collections.abc import Callable, Iterable, Sequence from contextlib import contextmanager @@ -143,6 +144,25 @@ def calibrate( def forward_loop(model): return original_forward_loop() # type: ignore[call-arg] + # Whether the algorithm actually invoked the loop. `max_calibrate` may skip it + # (`skip_forward_without_activation_calib`) when no enabled quantizer needs activation + # statistics, and a post-calibration check that assumed it ran would fire on every + # weight-only calibration. + forward_loop_ran = False + if forward_loop is not None: + _requested_forward_loop = forward_loop + + def forward_loop(model): + nonlocal forward_loop_ran + forward_loop_ran = True + return _requested_forward_loop(model) + + # The fused-experts interception check below judges this calibration's forwards alone. + for module in model.modules(): + if hasattr(module, "_intercept_count"): + module._intercept_count = module._forward_count = 0 + module._weight_input_capture = False + # move the model to eval mode, remembering each module's mode; quantizers, and plain containers of them, # follow the module they quantize instead training_modes = { @@ -169,6 +189,20 @@ def forward_loop(model): _restore_training_modes(model, training_modes) + if forward_loop_ran: + # A fused-experts module whose forward ran past both interception hooks collected + # nothing. Surface it here rather than let the export emit weight-derived fallback scales + # under a warning that blames the size of the calibration set. Raised once the model is + # back in its mode, for a caller that catches it to calibrate again. Only a loaded HF plugin + # installs fused-experts wrappers; importing it here would rerun one that failed to load, + # which import_plugin only warns about, and fail every calibration with its error. Both + # conditions must come out the same on every rank: under expert parallelism the check + # all-reduces over the experts' group, and a rank that skipped it would leave its peers + # waiting there. + huggingface = sys.modules.get("modelopt.torch.quantization.plugins.huggingface") + if huggingface is not None: + huggingface.assert_fused_experts_intercepted(model) + return model diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 8a2d0001086..dd086013118 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -24,6 +24,7 @@ from typing import TYPE_CHECKING import torch +import torch.distributed as dist import torch.nn as nn import transformers from torch import Tensor @@ -39,6 +40,7 @@ ) from modelopt.torch.kernels.quantization.gemm import IS_AVAILABLE as IS_TRITON_AVAILABLE from modelopt.torch.opt.dynamic import DynamicModule +from modelopt.torch.utils import warn_rank_0 from modelopt.torch.utils.distributed import ParallelState from ..algorithms import AutoQuantizeGradientSearcher @@ -58,6 +60,12 @@ except ImportError: _transformers_moe = None +# Whether transformers has the grouped experts forward's `_grouped_linear`, the call _QuantFusedExperts +# intercepts on it: without it nothing intercepts grouped_mm, so calibration cannot use that backend. +_HAS_GROUPED_LINEAR = _transformers_moe is not None and hasattr( + _transformers_moe, "_grouped_linear" +) + try: import kitchen from kitchen.fa import KitchenFlashAttentionModule @@ -803,22 +811,18 @@ class _QuantFusedExperts(_QuantFunctionalMixin): ``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 + they calibrate can differ, so ``eager`` is the way to calibrate as before + ``grouped_mm`` was intercepted. 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``). + run it pin. 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. + 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 @@ -836,7 +840,7 @@ class _QuantFusedExperts(_QuantFunctionalMixin): # 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. + # calibrate there (see normalize_experts_impl_on_the_fly). _eager_only = False def __init_subclass__(cls, **kwargs): @@ -897,6 +901,32 @@ def _setup(self): self._register_temp_attribute("_current_expert_idx", 0) # Per forward: each projection's experts, fake-quantized together (see _quantized_weight). self._register_temp_attribute("_stacked_weights", {}) + # Cumulative tokens-per-expert offsets of the grouped forward in flight (grouped_mm or + # torch_a2a), None on the eager path: tells a weight-input capture (local Hessian, imatrix, + # GPTQ) which forward ran and how to split the expert-sorted batch. Cleared when forward + # returns, so a capture reached outside one skips rather than splitting by stale offsets. + self._register_temp_attribute("_grouped_offs", None) + # How many times an interception hook fired for this module. Still zero after a + # calibration forward means the active experts backend reached neither `F.linear` nor + # `_grouped_linear`, so no quantizer ever saw data -- `assert_fused_experts_intercepted` + # turns that into an error instead of an export full of fallback scales. + self._register_temp_attribute("_intercept_count", 0) + # How many times this module's forward ran: an experts module that no calibration token + # reached (e.g. vision experts under text-only calibration) has nothing to intercept. + self._register_temp_attribute("_forward_count", 0) + # Whether this calibration captures the experts' inputs to calibrate their weights (GPTQ, + # local Hessian, imatrix), so that the weight quantizers need the interception too. + self._register_temp_attribute("_weight_input_capture", False) + + @property + def _running_grouped_mm(self) -> bool: + """Whether this module's forward runs its expert matmuls through ``_grouped_linear``. + + The ``grouped_mm`` backend and the ``torch_a2a`` expert-parallel forward both do, handing each + projection one expert-sorted batch rather than one ``F.linear`` call per expert. + """ + cfg = getattr(self, "config", None) + return getattr(cfg, "_experts_implementation", None) in ("grouped_mm", "torch_a2a") def _local_experts(self, weight: torch.Tensor) -> torch.Tensor: """``weight`` indexed the way the per-expert quantizers are: by this module's own experts. @@ -925,16 +955,30 @@ def _quantize_grouped_weight(self, weight: torch.Tensor, quantizers: nn.ModuleLi Each expert slice passes through its own weight quantizer (collecting ``amax`` during calibration, fake-quantizing afterwards), matching the eager per-expert ``F.linear`` semantics. Returns the weight untouched when every - quantizer would pass its slice through unchanged, so passthrough stays exact (and cheap). + quantizer would pass its slice through unchanged, so passthrough stays exact (and cheap): + a disabled quantizer, or one that only collects statistics, as in calibration. + + Reached only inside the experts forward, around which transformers swaps an expert-parallel + DTensor weight for this rank's local shard, so the weight is already indexed by local expert. """ - # A disabled quantizer still applies its pre-quant scale and rotation, as on the eager path. + # A quantizer that is disabled, or that collects statistics without quantizing, hands its + # slice back unchanged, unless it applies a pre-quant scale or a rotation (a disabled one + # still does, as on the eager path). Then the result is the weight itself: building it + # slice by slice would copy every expert's weight on each calibration forward. if all( isinstance(q, TensorQuantizer) - and not q.is_enabled + and not (q.is_enabled and q._if_quant) and q.pre_quant_scale is None and not q.rotate_is_enabled for q in quantizers ): + # Only a calibrating quantizer does anything with its slice: one whose amax is constant + # (enable_stats_collection leaves it neither quantizing nor calibrating) would not. + collecting = [q.is_enabled and q._if_calib for q in quantizers] + if any(collecting): + for q, w, collects in zip(quantizers, weight, collecting, strict=True): + if collects: + q(w) # collects this expert's statistics 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 @@ -973,6 +1017,9 @@ def functionals_to_replace(self): # 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): + self._intercept_count += 1 + # The eager path: no expert-sorted batch for a capture hook to split. + self._grouped_offs = None if self._down_proj_linear: idx = self._current_expert_idx input = self.down_proj_input_quantizer(input) @@ -1000,29 +1047,31 @@ def _quantized_linear(input, weight, bias=None): # 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"): + if _HAS_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. + # this wrapper's weights; normalize_experts_impl_on_the_fly pins its model to eager. 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._intercept_count += 1 + # The batch holds every expert's rows, so no expert for the eager loop's index to name. + # The offsets, recorded before quantizing input, let a capture hook on the input + # quantizer split the expert-sorted batch back per expert. self._current_expert_idx = None + self._grouped_offs = offs 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) + input_quantizer = ( + self.down_proj_input_quantizer if self._down_proj_linear else first_input_q + ) + input = self._quantize_grouped_input(input, offs, input_quantizer) weight = self._quantize_grouped_weight(weight, quantizers) self._down_proj_linear = not self._down_proj_linear return _orig_grouped_linear( @@ -1033,6 +1082,24 @@ def _quantized_grouped_linear(input, weight, offs, bias=None, is_transposed=Fals return replacements + @staticmethod + def _quantize_grouped_input(input, offs, quantizer): + """Quantize a grouped projection's input, collecting statistics on the rows it computes only. + + Rows past ``offs[-1]`` belong to no expert of this module: under transformers' own expert + parallelism they are the rows routed to other ranks' experts, sorted to the tail, which the + grouped GEMM neither reads nor writes. Before transformers 5.17 the down projection's input + holds them uninitialized, so a calibrating quantizer sees only the rows before them; the tail + passes through unread. Finding where it starts costs a host sync, so it is skipped while the + quantizer collects nothing. + """ + if getattr(quantizer, "is_enabled", False) and getattr(quantizer, "_if_calib", False): + routed = int(offs[-1]) if offs.numel() else 0 + if routed < input.shape[0]: + computed = quantizer(input[:routed]) + return torch.cat([computed, input[routed:].to(computed.dtype)]) + return quantizer(input) + def _quantized_weight(self, name, quantizers, idx, dtype): """Expert ``idx``'s ``name`` projection, fake-quantized by its own weight quantizer. @@ -1059,12 +1126,31 @@ def _quantized_weight(self, name, quantizers, idx, dtype): def forward(self, *args, **kwargs): self._down_proj_linear = False + # Only a forward that makes a projection call is work the interception check expects. The + # eager path makes one per expert some row was routed to: none for a forward with no rows, nor + # for one whose rows transformers' expert parallelism routed to other ranks' experts (masked + # to num_experts). The grouped paths call _grouped_linear whatever rows they are handed. + # Once a hook has fired the check passes, so stop counting: the routing test syncs the host. + if self._intercept_count == 0: + hidden_states = args[0] if args else kwargs.get("hidden_states") + top_k_index = args[1] if len(args) > 1 else kwargs.get("top_k_index") + if hidden_states is None or ( + hidden_states.numel() + and ( + self._running_grouped_mm + or not isinstance(top_k_index, torch.Tensor) + or bool((top_k_index < self.num_experts).any()) + ) + ): + self._forward_count += 1 self._stacked_weights = {} try: return super().forward(*args, **kwargs) finally: # Autograd keeps what backward needs; holding more would pin every expert's copy. self._stacked_weights = {} + # A capture reached outside a forward must not split its input by this forward's offsets. + self._grouped_offs = None def get_expert_weight(self, name, idx, dtype=None): """Return one expert's projection in the calibration/export weight domain. @@ -1145,7 +1231,7 @@ class _QuantFP8Experts(_QuantFusedExperts): # 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 is normalize_experts_impl_on_the_fly, which runs on convert and on restore. _eager_only = True @property @@ -1572,8 +1658,8 @@ def _is_fused_experts_module(module): 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 + 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 @@ -1628,45 +1714,208 @@ def register_fused_experts_on_the_fly(model): QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(wrapper) -def force_eager_experts_impl_on_the_fly(model): - """Force HF fused-experts modules onto the eager ``F.linear``-based forward. +def normalize_experts_impl_on_the_fly(model): + """Redirect HF fused-experts modules onto a quantization-supported backend. HF transformers 5.0+ decorates fused-experts forwards with - ``@use_experts_implementation``, which may dispatch to ``torch._grouped_mm`` - 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`` / - ``speech_config``) whenever a fused-experts module is present. + ``@use_experts_implementation``, which dispatches on + ``config._experts_implementation``. ``_QuantFusedExperts`` intercepts the + ``eager`` (per-expert ``F.linear``) and ``grouped_mm`` + (``transformers.integrations.moe._grouped_linear``) forwards. The remaining + backends — ``batched_mm`` (``torch.bmm``), ``sonicmoe`` (opaque CUDA kernel), + ``kernels``/hub, etc. — bypass those hooks, so calibration silently collects + nothing and the exported checkpoint produces garbage at inference. + + To avoid that while keeping the fast grouped path, this redirects any + unsupported backend to ``grouped_mm`` (which has a CPU / old-HW fallback, so + it is universally runnable), or to ``eager`` on a transformers without the + ``_grouped_linear`` grouped_mm is intercepted at. ``eager``, ``grouped_mm`` and ``torch_a2a`` are + left untouched, and so is ``None``, which dispatches to the eager ``F.linear`` + forward -- but transformers resolves an unrequested backend to ``grouped_mm`` + wherever that can dispatch, when it builds the model, so a model calibrates on + the grouped path unless ``eager`` was asked for. Applied to the model config + and recursively to ``text_config`` / ``vision_config`` / ``audio_config`` / + ``speech_config`` whenever a fused-experts module is present. + + Experts that quantize only on the eager forward are refused on ``torch_a2a``, the + expert-parallel forward, which cannot fall back to eager; no config is changed then. """ fused_experts = [module for module in model.modules() if _is_fused_experts_module(module)] if not fused_experts: return nested_cfg_attrs = ("text_config", "vision_config", "audio_config", "speech_config") + # Backends already intercepted by _QuantFusedExperts (None => eager F.linear forward). `torch_a2a` + # (all-to-all DP=EP) runs its local experts through the same `_grouped_linear`, so _QuantFusedExperts + # intercepts it too -- don't clobber it back to grouped_mm. Without `_grouped_linear` grouped_mm is + # intercepted nowhere, so it is neither kept nor redirected to: main's eager is. + supported_impls = (None, "eager", "torch_a2a") + ( + ("grouped_mm",) if _HAS_GROUPED_LINEAR else () + ) + 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. + 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 + and (not issubclass(wrapper, _QuantFusedExperts) or wrapper._eager_only) + } + ) - def _force(cfg): - if cfg is None: - return - if hasattr(cfg, "_experts_implementation"): - cfg._experts_implementation = "eager" - for sub in nested_cfg_attrs: - if hasattr(cfg, sub): - _force(getattr(cfg, sub)) - - if hasattr(model, "config"): - _force(model.config) + def _configs(cfg): + if cfg is not None: + yield cfg + for sub in nested_cfg_attrs: + yield from _configs(getattr(cfg, sub, None)) # Some models construct auxiliary blocks from a deep copy of the root config. # Nemotron-H MTP does this, so updating only ``model.config`` leaves its fused - # experts on grouped_mm/batched_mm and bypasses the F.linear quantizer hooks. - for module in fused_experts: - _force(getattr(module, "config", None)) + # experts on an unsupported backend that bypasses the quantizer hooks. + roots = [getattr(model, "config", None)] + [getattr(m, "config", None) for m in fused_experts] + configs = [c for root in roots for c in _configs(root) if hasattr(c, "_experts_implementation")] + if eager_only and any(c._experts_implementation == "torch_a2a" for c in configs): + raise NotImplementedError( + f"{', '.join(eager_only)} quantize only on the eager experts forward, but this model runs " + "its experts on `torch_a2a`, the expert-parallel forward, which cannot fall back to " + "eager. Quantize the model without expert parallelism." + ) + moved = {} # (was, now), once each: the root config and its sub-configs often agree + for cfg in configs: + was = cfg._experts_implementation + if eager_only: + cfg._experts_implementation = "eager" + elif was not in supported_impls: + cfg._experts_implementation = fallback_impl + # None already dispatches to the eager F.linear forward: pinning it to eager moves nothing. + if cfg._experts_implementation != was and was is not None: + moved[(was, cfg._experts_implementation)] = None + reason = ( + f"{', '.join(eager_only)} quantize on the eager forward only" + if eager_only + else "nothing intercepts that backend" + ) + for was, now in moved: + # The caller may have asked for it (a load's experts_implementation, hf_ptq's flag). Every + # rank redirects alike, so one says so. + warn_rank_0( + f"Quantizing the fused experts on {now!r} instead of experts_implementation={was!r}: " + f"{reason}." + ) + + +def _bypassed_both_hooks(module) -> bool: + """Whether ``module``'s forward ran past both hooks, with a quantizer that needed them. + + See :func:`assert_fused_experts_intercepted`. + """ + # A missing attribute (rather than 0) means it was already cleaned up -- say nothing. + if getattr(module, "_intercept_count", None) != 0 or not getattr(module, "_forward_count", 0): + return False + first = module._first_proj_attr + quantizers = [getattr(module, f"{first}_input_quantizer", None)] + quantizers.append(getattr(module, "down_proj_input_quantizer", None)) + if getattr(module, "_weight_input_capture", False): + for attr in (f"{first}_weight_quantizers", "down_proj_weight_quantizers"): + quantizers.extend(getattr(module, attr, None) or []) + return any(getattr(q, "is_enabled", False) for q in quantizers) + + +def _expert_parallel_group(module): + """The process group ``module``'s experts are sharded over (expert parallelism), else ``None``. + + Read from the experts' weight, a DTensor sharded over the experts under the expert parallelism of + transformers 5.16 and later. Earlier versions keep each rank's experts as a plain tensor, which + names no group. + """ + weight = getattr(module, module._first_proj_attr, None) + if not isinstance(weight, DTensor) or weight.shape[0] == module.num_experts: + return None + dims = [dim for dim, placement in enumerate(weight.placements) if placement == Shard(0)] + return weight.device_mesh.get_group(dims[0]) if dims else None + + +def assert_fused_experts_intercepted(model): + """Raise if a quantized fused-experts module was never reached by an interception hook. + + :class:`_QuantFusedExperts` collects statistics by intercepting the calls the experts + forward makes -- ``F.linear`` per expert on the eager path, ``_grouped_linear`` once per + projection on the grouped_mm path. A backend that uses neither bypasses both hooks + *silently*: no quantizer is invoked, no amax is collected, and the export then falls back + to weight-derived scales while warning about calibration size -- which points at the + dataset when the fault is in the dispatch. + + The mismatch is structural. :func:`register_fused_experts_on_the_fly` claims a module on + the *shape* of its parameters, but the hooks depend on the *behavior* of a forward that + nothing enforces, and "eager" is not a shared implementation -- it is whatever the model + author wrote. :func:`normalize_experts_impl_on_the_fly` can redirect a class that carries + ``@use_experts_implementation``, but a class without it ignores + ``config._experts_implementation`` entirely, so no config change can rescue it. Checking + the hit count is a direct test of the thing actually depended on. + + Only modules whose forward ran and that have a quantizer needing it are reported: an enabled + input quantizer, or an enabled weight quantizer when the calibration captured the experts' + inputs to calibrate their weights (GPTQ, local Hessian, imatrix). Otherwise a weight quantizer + calibrates on its weight alone (:func:`weight_only_quantize`), so weight-only experts lose + nothing by not being intercepted, and neither do experts the recipe leaves unquantized. An + experts module no calibration token reached (vision experts under text-only calibration) falls + back to weight-derived scales as before. + + Under the expert parallelism of transformers 5.16 and later, a module's verdict is reduced over + the group its experts are sharded over, which every member already joins in the experts' + forward, so the whole group raises together. Otherwise the check is rank-local, also under + earlier versions' expert parallelism, whose plain-tensor shards name no group: it raises on the + ranks whose experts saw rows, and a rank whose experts saw none goes on to its next collective + and waits there until the launcher stops it or the collective times out, so the error is in + another rank's log. + """ + experts = [ + (name, module, _bypassed_both_hooks(module), _expert_parallel_group(module)) + for name, module in model.named_modules() + if isinstance(module, _QuantFusedExperts) + ] + bypassed = [here for _, _, here, _ in experts] + # Each member of an expert-parallel group holds the same experts modules in the same order, so + # one reduction per group, whatever this rank found, tells it what any member found. + for group in dict.fromkeys(group for *_, group in experts if group is not None): + members = [i for i, (*_, g) in enumerate(experts) if g is group] + device = experts[members[0]][1].down_proj.device # where the group's backend reads it + flags = torch.tensor([bypassed[i] for i in members], dtype=torch.int32, device=device) + dist.all_reduce(flags, op=dist.ReduceOp.MAX, group=group) + for i, flag in zip(members, flags.tolist()): + bypassed[i] = bool(flag) + + missed = [] + for (name, module, here, _), anywhere in zip(experts, bypassed): + if anywhere: + impl = getattr(getattr(module, "config", None), "_experts_implementation", None) + where = "" if here else ", on another rank of its expert-parallel group" + missed.append( + f" {name} ({type(module).__name__}, experts_implementation={impl!r}{where})" + ) + + if not missed: + return + + raise RuntimeError( + "Fused-experts calibration collected nothing: the forward of the module(s) below reached " + "neither `F.linear` nor `_grouped_linear`, so no calibration batch reached their quantizers: " + "their input scales would fall back to weight-derived values, and a weight calibration that " + "reads the experts' inputs got none.\n" + + "\n".join(missed) + + "\n\nTheir forward ran, but the active experts backend bypasses both hooks -- " + "`batched_mm`, `sonicmoe`, or a hand-written `bmm`/`matmul` forward on a class without " + "`@use_experts_implementation`, which `normalize_experts_impl_on_the_fly` cannot " + "redirect. Check the module's `forward` and `config._experts_implementation`. A class no " + "backend change redirects needs its own quantized wrapper, as GptOssExperts has; to " + "quantize the rest of the model meanwhile, disable these modules' quantizers in the quant " + "config." + ) def _is_supported_hf_model(model): @@ -1784,7 +2033,7 @@ def _is_param_grad_enabled_for_auto_quantize(pname, model): CUSTOM_MODEL_PLUGINS.update( [ register_fused_experts_on_the_fly, - force_eager_experts_impl_on_the_fly, + normalize_experts_impl_on_the_fly, register_sparse_moe_on_the_fly, register_hf_attentions_on_the_fly, convert_hf_parallel_linears_on_the_fly, diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 563c28fea2d..2aff08e6891 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -19,7 +19,7 @@ all-to-all expert transport (``torch_a2a``). The all-to-all transport is plain ``torch.distributed.all_to_all_single`` (NCCL) wrapped around -transformers' ``_grouped_linear``. ``register_torch_a2a()`` registers it into +transformers' quantizable ``_grouped_linear``. ``register_torch_a2a()`` registers it into transformers' ``ExpertsInterface`` lazily -- never at import time, so importing this module has no transformers-``ExpertsInterface`` side effect. """ @@ -147,14 +147,14 @@ def _routed_expert_modules(model: "nn.Module", ep_plan: dict[str, str]) -> "dict # # The throughput-correct MoE parallelism for calibration: dense+router run data-parallel (each rank its # own batch), experts are EP-sharded, and tokens are routed to the expert-owning rank by an all-to-all -# (dispatch) / reduce (combine). The expert matmul stays on transformers' `_grouped_linear`, so each -# expert sees the tokens routed to it from every rank. +# (dispatch) / reduce (combine). The expert matmul stays on transformers' `_grouped_linear`, so +# ModelOpt's per-expert quantizers attach and collect amax over the *global* (dispatched) token pool. # Transport is plain `torch.distributed.all_to_all_single` (NCCL) -- works cross-node at any rank count. # --------------------------------------------------------------------------------------------------- -# all_to_all_single addresses its flattened buffers with int32 element offsets, so a hop larger than -# this overflows them; a2a_moe chunks such transfers along the hidden dim. +# all_to_all_single addresses its flattened buffers with int32 element offsets; a2a_moe sends a larger +# hop as hidden-column slices, so every hop stays within them. _A2A_MAX_HOP_ELEMENTS = 2**31 - 1 # The combine weights and sums each token's top_k results a slice of tokens at a time, so its # temporaries stay within this many elements however many tokens the layer routes. @@ -170,7 +170,7 @@ def a2a_moe(group, hidden_states, topk_idx, topk_weights, num_local_experts, exp Works cross-node at any rank count, with explicit token permute/unpermute. Flow: expand (token, expert) pairs -> route each to dest_rank = global_expert // num_local -> - all_to_all the tokens (+ local-expert id) -> local grouped_mm -> all_to_all back -> + all_to_all the tokens (+ local-expert id) -> local grouped_mm (quantizable) -> all_to_all back -> weight by the routing weights and sum each token's top_k results. Tokens travel and the result returns in ``hidden_states``' dtype -- the model's, so its residual @@ -204,9 +204,9 @@ def _one_all_to_all(output, input_, out_split, in_split): def _all_to_all(output, input_, out_split, in_split, step=None): """All-to-all in hidden-column slices ``step`` wide; in one hop if ``step`` is None or covers H. - all_to_all_single addresses the flattened ``[rows, H]`` buffers with int32 element offsets, so a - hop past 2**31 elements overflows them and hangs. Every slice has the same row splits, and - concatenating the slices is autograd-safe. + all_to_all_single addresses the flattened ``[rows, H]`` buffers with int32 element offsets; + slicing keeps every hop within them. Every slice has the same row splits, and concatenating the + slices is autograd-safe. """ hdim = input_.shape[-1] if step is None or step >= hdim: @@ -286,7 +286,7 @@ def _all_to_all(output, input_, out_split, in_split, step=None): recv_eid = torch.empty(Q, device=dev, dtype=torch.int64) dist.all_to_all_single(recv_eid, local_eid_s.contiguous(), out_list, in_list, group=group) - # group received tokens by local expert -> grouped_mm; weighted back home + # group received tokens by local expert -> grouped_mm (quantizable); weighted back home rorder = torch.argsort(recv_eid, stable=True) counts = torch.bincount(recv_eid[rorder], minlength=num_local_experts) offsets = torch.cumsum(counts, dim=0).to(torch.int32) @@ -327,10 +327,10 @@ def torch_a2a_experts_forward(self, hidden_states, top_k_index, top_k_weights, p Registered as the `"torch_a2a"` experts implementation on transformers' `ExpertsInterface`. Dense+router run data-parallel; experts are EP-sharded; tokens are routed to the expert-owning - rank by the NCCL all-to-all above. Local experts run through transformers' `_grouped_linear` on - the tokens dispatched to them from every rank. `process_group` (the EP group) is injected by the - `torch_a2a_experts` TP style's `_prepare_input_fn`; the router must emit GLOBAL ids unmasked - (`torch_a2a_router`). + rank by the NCCL all-to-all above. Local experts run through the ModelOpt-quantizable + `_grouped_linear` so per-expert weight + shared input quantizers attach over the dispatched + (global) token pool. `process_group` (the EP group) is injected by the `torch_a2a_experts` TP style's + `_prepare_input_fn`; the router must emit GLOBAL ids unmasked (`torch_a2a_router`). ``self.num_experts`` is this rank's share of the experts: the expert-parallel loader loads them through transformers' expert-parallel sharding (the ``grouped_gemm`` / packed styles), which slices @@ -347,7 +347,8 @@ def torch_a2a_experts_forward(self, hidden_states, top_k_index, top_k_weights, p # Per-expert biases (e.g. GPT-OSS ``gate_up_proj_bias`` / ``down_proj_bias``) are indexed by local # expert id below, so the plan must shard them over dim 0 like the weights (grouped_gemm): left # whole, every rank but the first would add other experts' biases. transformers' - # ``_grouped_linear`` adds them after each expert's grouped GEMM. + # ``_grouped_linear`` (and ModelOpt's quantized override) add them after each expert's grouped GEMM, + # unquantized. has_bias = getattr(self, "has_bias", False) if has_bias: for name in ("gate_up_proj_bias" if self.has_gate else "up_proj_bias", "down_proj_bias"): diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index c2435e96973..392b3d2807e 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -28,7 +28,7 @@ import pytest import torch from _test_utils.examples.hf_ptq_example_utils import example_utils -from _test_utils.torch.transformers_models import get_tiny_nemotron_h +from _test_utils.torch.transformers_models import create_tiny_qwen3_moe_dir, get_tiny_nemotron_h from safetensors.torch import save_file from transformers import AutoModelForCausalLM, LlamaConfig @@ -137,6 +137,14 @@ def test_mtp_preparation_matches_loader(tmp_path, loader, trust_remote_code): assert not any("MTP weights may remain unquantized" in str(w.message) for w in caught) +def test_get_model_loads_with_the_requested_experts_implementation(tmp_path): + """hf_ptq's --experts_implementation reaches the fused experts: eager calibrates them as earlier + releases did, where transformers would pick grouped_mm for this model.""" + path = create_tiny_qwen3_moe_dir(tmp_path) + model = example_utils.get_model(str(path), device="cpu", experts_implementation="eager") + assert model.config._experts_implementation == "eager" + + @pytest.mark.parametrize("failure_at", [1, 2]) def test_get_model_preserves_preparation_error(monkeypatch, tmp_path, failure_at): config = LlamaConfig(architectures=["LlamaForCausalLM"]) diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 75b611b68ff..2f0b425de40 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -26,6 +26,7 @@ import yaml from _test_utils.mlflow import clean_env # noqa: F401 from _test_utils.torch.transformers_models import get_tiny_qwen3 +from transformers import Qwen3_5MoeConfig, Qwen3MoeConfig from modelopt.recipe import load_recipe from modelopt.recipe.config import ( @@ -414,6 +415,40 @@ def distributed_load(checkpoint, *_args, **_kwargs): hf_ptq.load_model(args) +@pytest.mark.parametrize( + "config", + [ + Qwen3MoeConfig(architectures=["Qwen3MoeForCausalLM"]), + Qwen3_5MoeConfig(architectures=["Qwen3_5MoeForConditionalGeneration"]), + ], + ids=["causal-lm", "vlm"], +) +def test_fsdp2_load_builds_with_the_requested_experts_implementation(monkeypatch, tmp_path, config): + """--experts_implementation reaches the config the FSDP2 load builds its meta model from, and + the text config a VLM's experts dispatch on.""" + config.save_pretrained(tmp_path) + monkeypatch.setenv("RANK", "0") + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + str(tmp_path), + "--use_fsdp2", + "--experts_implementation", + "eager", + ) + args.dist_state = SimpleNamespace(device="cpu", rank=0, world_size=1) + args.hub_model_id, args.local_checkpoint_path = None, str(tmp_path) + + def distributed_load(*_args, hf_config, **_kwargs): + assert hf_config._experts_implementation == "eager" + assert hf_config.get_text_config()._experts_implementation == "eager" + raise RuntimeError("distributed loader reached") + + monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", distributed_load) + with pytest.raises(RuntimeError, match="distributed loader reached"): + hf_ptq.load_model(args) + + @pytest.mark.parametrize("has_mtp", [False, True]) @pytest.mark.parametrize("declared", [False, True]) def test_fsdp2_rejects_mtp_before_distributed_loading(monkeypatch, tmp_path, has_mtp, declared): diff --git a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py index d37d3aba3b1..48165090687 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py +++ b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py @@ -56,6 +56,7 @@ def test_image_calibration_model_target_follows_recipe( trust_remote_code=False, use_seq_device_map=False, attn_implementation=None, + experts_implementation=None, offload_folder=None, max_cpu_memory_gb=None, max_gpu_memory_gb=None, diff --git a/tests/gpu/torch/quantization/test_fsdp2_autoquant.py b/tests/gpu/torch/quantization/test_fsdp2_autoquant.py index e587b9a7db9..e2f45ee9507 100644 --- a/tests/gpu/torch/quantization/test_fsdp2_autoquant.py +++ b/tests/gpu/torch/quantization/test_fsdp2_autoquant.py @@ -109,6 +109,7 @@ def _run_kv_autoquant_example(rank, size, checkpoint_dir, search_dir, composed): trust_remote_code=False, cpu_offload=False, attn_implementation="eager", + experts_implementation=None, calib_with_images=False, specdec_offline_dataset=None, low_memory_mode=False, diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 22c91454332..d7f03854837 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -15,14 +15,18 @@ """Tests for _QuantFusedExperts: generic fused MoE quantization and export.""" +import ast import copy +import sys import warnings import weakref from contextlib import nullcontext +from pathlib import Path from unittest import mock import pytest import torch +import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F from torch.distributed.device_mesh import init_device_mesh @@ -32,6 +36,7 @@ pytest.importorskip("transformers") +import transformers from _test_utils.torch.quantization.tied_modules import tie_fused_experts_3d_params from transformers import PretrainedConfig, Qwen3MoeConfig from transformers.integrations.finegrained_fp8 import FP8Experts @@ -54,7 +59,10 @@ from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name from modelopt.torch.quantization.model_calib import ( + _expert_slices, _register_weight_input_hooks, + enable_stats_collection, + finish_stats_collection, local_hessian_calibrate, ) from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer @@ -64,12 +72,14 @@ _is_sparse_sequaential_moe_block, _QuantFusedExperts, _QuantNonGatedFusedExperts, - force_eager_experts_impl_on_the_fly, + assert_fused_experts_intercepted, + normalize_experts_impl_on_the_fly, register_fused_experts_on_the_fly, register_sparse_moe_on_the_fly, ) from modelopt.torch.quantization.utils.calib_utils import GPTQHelper from modelopt.torch.quantization.utils.core_utils import module_name_maps, weight_attr_names +from modelopt.torch.utils.moe import torch_a2a_experts_forward def test_native_fp8_experts_dequantized_before_nvfp4_export(): @@ -162,7 +172,7 @@ def __init__(self): 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_mask = F.one_hot(top_k_index, num_classes=self.num_experts + 1).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] @@ -322,6 +332,7 @@ class _ExpertsConfig: def __init__(self, impl="grouped_mm"): self._experts_implementation = impl self.hidden_act = "silu" + self.ep_group = None # the expert-parallel group a torch_a2a forward dispatches over class _SyntheticGroupedExperts(nn.Module): @@ -381,9 +392,36 @@ def forward(self, hidden_states, top_k_index, top_k_weights): return _moe_integration.grouped_mm_experts_forward( self, hidden_states, top_k_index, top_k_weights ) + if self.config._experts_implementation == "torch_a2a": + return torch_a2a_experts_forward( + self, hidden_states, top_k_index, top_k_weights, self.config.ep_group + ) return self._eager_forward(hidden_states, top_k_index, top_k_weights) +class _SentinelTailExperts(_SyntheticGroupedExperts): + """A grouped forward as transformers' own expert parallelism runs it before 5.17. + + The rows routed to other ranks' experts are sorted past ``offs[-1]``, where the grouped GEMM + neither reads nor writes: its input carries ``tail`` such rows, and the down projection's input + keeps the up projection's tail rows uninitialized, here ``fill`` (the CPU fallback zeroes them). + """ + + tail = 0 + fill = float("nan") + + def forward(self, hidden_states, top_k_index, top_k_weights): + rows = hidden_states.shape[0] + offs = torch.arange(1, self.num_experts + 1, dtype=torch.int32) * (rows // self.num_experts) + garbage = hidden_states.new_full((self.tail, hidden_states.shape[1]), self.fill) + up = _moe_integration._grouped_linear( + torch.cat([hidden_states, garbage]), self.gate_up_proj, offs + ) + up[rows:] = self.fill + down = _moe_integration._grouped_linear(self._apply_gate(up), self.down_proj, offs) + return down[:rows] + + class _GroupedSparseMoeBlock(nn.Module): def __init__(self, config): super().__init__() @@ -552,28 +590,6 @@ def build(): ): 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() @@ -1338,7 +1354,7 @@ def test_untied_fused_experts_have_independent_buffers(self): # --------------------------------------------------------------------------- -# Tests for force_eager_experts_impl_on_the_fly +# Tests for normalize_experts_impl_on_the_fly # --------------------------------------------------------------------------- class _StubConfig: """Minimal stand-in for HF PretrainedConfig with optional nested sub-configs.""" @@ -1363,49 +1379,135 @@ def __init__(self, config): self.config = config -class TestForceEagerExpertsImpl: - def test_sets_eager_on_moe_model(self): - """Non-eager backend on an MoE model gets flipped to eager.""" +class TestNormalizeExpertsImpl: + def test_redirects_unsupported_backend_to_grouped_mm(self): + """An unsupported backend on an MoE model gets redirected to grouped_mm.""" cfg = _StubConfig(impl="kernels") model = _TinyMoEModelWithConfig(cfg) - force_eager_experts_impl_on_the_fly(model) - assert cfg._experts_implementation == "eager" + with pytest.warns( + UserWarning, match="on 'grouped_mm' instead of experts_implementation='kernels'" + ): + normalize_experts_impl_on_the_fly(model) + assert cfg._experts_implementation == "grouped_mm" - def test_recurses_into_nested_configs(self): - """VLM-style nested text_config / vision_config are also flipped.""" + def test_falls_back_to_eager_without_the_grouped_forward(self, monkeypatch): + """A transformers without `_grouped_linear` leaves grouped_mm intercepted nowhere, so it is + neither kept nor redirected to: every backend but eager goes to eager, as main forced.""" + monkeypatch.setattr( + "modelopt.torch.quantization.plugins.huggingface._HAS_GROUPED_LINEAR", False + ) text_cfg = _StubConfig(impl="grouped_mm") - vision_cfg = _StubConfig(impl="bmm") + root_cfg = _StubConfig(impl="batched_mm", text_config=text_cfg) + normalize_experts_impl_on_the_fly(_TinyMoEModelWithConfig(root_cfg)) + assert root_cfg._experts_implementation == text_cfg._experts_implementation == "eager" + + def test_recurses_into_nested_configs(self): + """VLM-style nested text_config / vision_config are also normalized.""" + text_cfg = _StubConfig(impl="grouped_mm") # already supported -> unchanged + vision_cfg = _StubConfig(impl="batched_mm") # unsupported -> grouped_mm root_cfg = _StubConfig(text_config=text_cfg, vision_config=vision_cfg) model = _TinyMoEModelWithConfig(root_cfg) - model.moe.experts.config = _StubConfig(impl="grouped_mm") - force_eager_experts_impl_on_the_fly(model) - assert text_cfg._experts_implementation == "eager" - assert vision_cfg._experts_implementation == "eager" - assert model.moe.experts.config._experts_implementation == "eager" + model.moe.experts.config = _StubConfig(impl="batched_mm") + normalize_experts_impl_on_the_fly(model) + assert text_cfg._experts_implementation == "grouped_mm" + assert vision_cfg._experts_implementation == "grouped_mm" + # A module's own (deep-copied) config is redirected too, as Nemotron-H MTP needs. + assert model.moe.experts.config._experts_implementation == "grouped_mm" def test_skips_model_without_fused_experts(self): """Non-MoE models must not have their config silently mutated.""" cfg = _StubConfig(impl="kernels") model = _NonMoEModelWithConfig(cfg) - force_eager_experts_impl_on_the_fly(model) + normalize_experts_impl_on_the_fly(model) assert cfg._experts_implementation == "kernels" def test_no_crash_when_config_missing(self): """Model without a ``config`` attribute must not raise.""" - force_eager_experts_impl_on_the_fly(_TinyMoEModel()) # no-op, no error + normalize_experts_impl_on_the_fly(_TinyMoEModel()) # no-op, no error def test_no_crash_when_impl_attr_missing(self): """Config without ``_experts_implementation`` must not raise.""" cfg = _StubConfig() # no impl attr model = _TinyMoEModelWithConfig(cfg) - force_eager_experts_impl_on_the_fly(model) + normalize_experts_impl_on_the_fly(model) assert not hasattr(cfg, "_experts_implementation") - def test_leaves_eager_value_unchanged(self): - cfg = _StubConfig(impl="eager") + @pytest.mark.parametrize("impl", ["grouped_mm", None]) + def test_forces_eager_for_experts_quantized_on_the_eager_forward_only(self, impl, 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.""" + cfg = _StubConfig() + cfg._experts_implementation = impl + model = _NonMoEModelWithConfig(cfg) + model.experts = FP8Experts( + PretrainedConfig( + hidden_size=256, intermediate_size=256, num_experts=2, hidden_act="silu" + ), + block_size=(128, 128), + ) + 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 + config changes.""" + text_config = _StubConfig(impl="torch_a2a") + model = _NonMoEModelWithConfig(_StubConfig(impl="grouped_mm", text_config=text_config)) + model.experts = FP8Experts( + PretrainedConfig( + hidden_size=256, intermediate_size=256, num_experts=2, hidden_act="silu" + ), + block_size=(128, 128), + ) + with pytest.raises(NotImplementedError, match="FP8Experts"): + normalize_experts_impl_on_the_fly(model) + assert model.config._experts_implementation == "grouped_mm" + + 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 + + @pytest.mark.parametrize("impl", ["eager", "grouped_mm", None]) + def test_leaves_a_supported_value_unchanged(self, impl, recwarn): + """``eager``, ``grouped_mm`` and an explicit None (the model's eager ``F.linear`` forward) are + supported, so they are left as they are, without a word.""" + cfg = _StubConfig() + cfg._experts_implementation = impl model = _TinyMoEModelWithConfig(cfg) - force_eager_experts_impl_on_the_fly(model) - assert cfg._experts_implementation == "eager" + normalize_experts_impl_on_the_fly(model) + assert cfg._experts_implementation == impl + assert not [w for w in recwarn if "instead of experts_implementation" in str(w.message)] # --------------------------------------------------------------------------- @@ -1814,6 +1916,74 @@ def test_unrelated_dotted_number_unchanged(self): ) +def _local_hessian_grouped_capture_job(rank, size, impl): + """Local-Hessian capture under a grouped dispatcher; torch_a2a gets a one-rank EP group.""" + ep_group = dist.new_subgroups(1)[0] if impl == "torch_a2a" else None + try: + torch.manual_seed(0) + model = _GroupedMoEModel(impl) + model.config.ep_group = ep_group + + weight_quant = {"num_bits": 8, "axis": 0} + quant_cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*gate_up_proj_weight_quantizer", "cfg": weight_quant}, + {"quantizer_name": "*down_proj_weight_quantizer", "cfg": weight_quant}, + ], + "algorithm": "max", + } + batches = [torch.randn(1, 8, HIDDEN_DIM) for _ in range(3)] + + def forward_loop(m): + for batch in batches: + m(batch) + + mtq.quantize(model, quant_cfg, forward_loop=forward_loop) + experts = model.moe.experts + assert experts._running_grouped_mm is True + + expert_quantizers = list(experts.gate_up_proj_weight_quantizers) + list( + experts.down_proj_weight_quantizers + ) + max_amax = {id(q): q.amax.clone() for q in expert_quantizers if q.amax is not None} + expected_shape = {} + for quantizers, weight in ( + (experts.gate_up_proj_weight_quantizers, experts.gate_up_proj), + (experts.down_proj_weight_quantizers, experts.down_proj), + ): + for i, q in enumerate(quantizers): + expected_shape[id(q)] = (weight[i].shape[0], weight[i].shape[1]) + + local_hessian_calibrate(model, forward_loop, fp8_scale_sweep=False, debug=True) + + # Each expert's Hessian, in both projections, holds exactly the rows routed to it. + routed = torch.cat([model.moe.gate(b.view(-1, HIDDEN_DIM))[2] for b in batches]) + rows = torch.bincount(routed.flatten(), minlength=NUM_EXPERTS).tolist() + accumulators = model._local_hessian_accumulators + for quantizers in ( + experts.gate_up_proj_weight_quantizers, + experts.down_proj_weight_quantizers, + ): + captured = [ + accumulators[id(q)].num_samples if id(q) in accumulators else 0 for q in quantizers + ] + assert captured == rows, (captured, rows) + for qid, acc in accumulators.items(): + if acc.num_samples: + assert (acc.cout, acc.cin) == expected_shape[qid] + assert all(q.amax is not None and torch.isfinite(q.amax).all() for q in expert_quantizers) + assert any( + id(q) in max_amax and not torch.allclose(q.amax, max_amax[id(q)]) + for q in expert_quantizers + ) + finally: + if QuantModuleRegistry.get(_SyntheticGroupedExperts) is not None: + QuantModuleRegistry.unregister(_SyntheticGroupedExperts) + if ep_group is not None: + dist.destroy_process_group(ep_group) + + # --------------------------------------------------------------------------- # Tests for the grouped_mm backend: _QuantFusedExperts intercepts the real # transformers grouped_mm forward (via _grouped_linear), and grouped calibration @@ -1908,16 +2078,19 @@ def test_calibration_matches_eager(self): 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) + mtq.quantize( + eager_model, self._quant_cfg(), forward_loop=self._forward_loop_all_experts + ) + mtq.quantize( + grouped_model, self._quant_cfg(), forward_loop=self._forward_loop_all_experts + ) e = eager_model.moe.experts g = grouped_model.moe.experts + assert g._running_grouped_mm is True + assert e._running_grouped_mm is False + # 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 @@ -1968,6 +2141,64 @@ def one_alive(_quantizer, _args, output): expected = torch.stack([q(weight[idx]) for idx, q in enumerate(quantizers)]) torch.testing.assert_close(out, expected, rtol=0, atol=0) + @pytest.mark.parametrize("constant_amax", [False, True], ids=["calibrating", "constant-amax"]) + def test_calibrating_a_projection_returns_its_weight_without_a_copy(self, constant_amax): + """Weight quantizers that only collect statistics hand every slice back unchanged, so a + calibration forward returns the stored weight itself rather than a copy of every expert's + weight, and still collects each expert's statistics. One whose amax is constant collects + nothing, so it is not called.""" + 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 + calls = [] + for q in quantizers: + q.reset_amax() + q._use_constant_amax = constant_amax + q.register_forward_hook(lambda *_: calls.append(1)) + enable_stats_collection(experts) + with torch.no_grad(): + out = experts._quantize_grouped_weight(weight, quantizers) + + assert out is weight + assert len(calls) == (0 if constant_amax else NUM_EXPERTS) + if not constant_amax: + finish_stats_collection(experts) + for idx, q in enumerate(quantizers): + torch.testing.assert_close( + q.amax.flatten(), weight[idx].detach().abs().amax(dim=1), rtol=0, atol=0 + ) + + @pytest.mark.parametrize("fill", [float("nan"), 1e6], ids=["nan", "finite"]) + def test_rows_past_the_offsets_do_not_reach_the_input_statistics(self, fill): + """Under transformers' own expert parallelism before 5.17, rows past ``offs[-1]`` reach the + grouped projections uncomputed, the down projection's uninitialized: calibrating on them + would set its input scale from garbage, silently when it is finite. The input quantizers + collect statistics on the computed rows only, so such a tail changes no amax.""" + torch.manual_seed(0) + base = _SentinelTailExperts(_ExpertsConfig("grouped_mm")) + hidden = torch.randn(8, HIDDEN_DIM) + routing = (torch.zeros(8, TOP_K, dtype=torch.long), torch.ones(8, TOP_K)) + amaxes = [] + try: + for tail in (0, 3): + model = _ExpertsHolder(copy.deepcopy(base)) + model.experts.tail, model.experts.fill = tail, fill + mtq.quantize(model, self._quant_cfg(), lambda m: m(hidden, *routing)) + amaxes.append( + [ + model.experts.gate_up_proj_input_quantizer.amax, + model.experts.down_proj_input_quantizer.amax, + ] + ) + finally: + self._cleanup_registry(_SentinelTailExperts) + for with_tail, without in zip(amaxes[1], amaxes[0], strict=True): + torch.testing.assert_close(with_tail, without, 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 @@ -2045,10 +2276,10 @@ def test_fake_quantized_grouped_forward_and_backward_match_eager(self): 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 = [_GroupedMoEModel("eager"), _GroupedMoEModel("eager")] models[1].load_state_dict(models[0].state_dict()) try: - # quantize() calibrates every model on eager, so both get the same quantizers. + # Both calibrate on eager, so they get the same quantizers; then one runs grouped. for model in models: mtq.quantize(model, self._quant_cfg(), self._forward_loop_all_experts) models[1].config._experts_implementation = "grouped_mm" @@ -2069,28 +2300,49 @@ def test_fake_quantized_grouped_forward_and_backward_match_eager(self): for grouped, eager in zip(results[1], results[0]): torch.testing.assert_close(grouped, eager, rtol=1e-4, atol=1e-6) - def test_weight_input_captures_skip_the_grouped_forward(self): - """The weight-input captures key on the eager loop's expert index, which names no expert in the - grouped forward's batch of every expert's rows: they skip it rather than credit one expert - with all of it.""" - model = _GroupedMoEModel("eager") - rows = {} + def test_gptq_hessians_match_eager(self, monkeypatch): + """GPTQ gives each expert the Hessian of only its own tokens under grouped_mm, as eager. - def capture(quantizer, _weight, inputs): - rows[id(quantizer)] = rows.get(id(quantizer), 0) + inputs.shape[0] + grouped_mm quantizes one expert-sorted batch per projection rather than one call per + expert, so GPTQ's per-expert hook splits it back; keyed by the eager loop's expert index + alone, the first expert took every token and the others none. + """ + hessians = {} + free = GPTQHelper.free - 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 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}, + } + 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()) + found = [] + for model in (eager_model, grouped_model): + hessians.clear() + mtq.quantize(model, copy.deepcopy(cfg), forward_loop=self._forward_loop_all_experts) + found.append(dict(hessians)) + self._cleanup_registry(_SyntheticGroupedExperts) + + 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_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, @@ -2111,6 +2363,59 @@ class _EagerOnlyExperts(_QuantFusedExperts): finally: self._cleanup_registry(_SyntheticGroupedExperts) + @pytest.mark.parametrize( + ("selected", "ran"), [("grouped_mm", "eager"), ("eager", "grouped_mm")] + ) + def test_a_capture_follows_the_path_that_ran(self, selected, ran): + """A capture hook reads the path that ran, not the one the config named when it was + registered, so each expert still receives exactly the rows routed to it. Once a grouped + forward has returned, a capture reached outside it does not split its input by that + forward's offsets.""" + model = _GroupedMoEModel(selected) + 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 = ran + try: + self._forward_loop_all_experts(model) + if ran == "grouped_mm": + captured = dict(rows) + model.moe.experts.gate_up_proj_input_quantizer(torch.randn(3, HIDDEN_DIM)) + assert rows == captured + finally: + for handle in handles: + handle.remove() + finally: + self._cleanup_registry(_SyntheticGroupedExperts) + experts = model.moe.experts + # The loop routes 16 rows round-robin, two experts each: 8 rows per expert and projection. + for q in [*experts.gate_up_proj_weight_quantizers, *experts.down_proj_weight_quantizers]: + assert rows.get(id(q)) == 16 * TOP_K // NUM_EXPERTS + + @pytest.mark.parametrize( + "offs", + [[0, 2, 2, 5], [2, 5], [2, 2, 7], [2, 1, 5]], + ids=["inclusive-ends", "short", "past-the-rows", "out-of-order"], + ) + def test_grouped_offsets_of_another_form_are_refused(self, offs): + """The per-expert split reads the offsets as exclusive cumulative ends, one per expert: any + other form would credit experts with other experts' rows, so it raises instead.""" + assert list(_expert_slices(torch.tensor([2, 2, 5]), 3, 6)) == [(0, 0, 2), (2, 2, 5)] + with pytest.raises(RuntimeError, match="do not split 6 rows over 3 experts"): + list(_expert_slices(torch.tensor(offs), 3, 6)) + + @pytest.mark.parametrize("impl", ["grouped_mm", "torch_a2a"]) + def test_local_hessian_captures_per_expert_in_grouped_mode(self, cpu_dist_workers, impl): + """local_hessian splits the expert-sorted grouped batch by offsets and refines + each expert's weight amax — the grouped analogue of the eager per-expert hook. + The torch_a2a forward hands ``_grouped_linear`` the same batch, so it splits the same.""" + cpu_dist_workers.run(_local_hessian_grouped_capture_job, impl=impl) + # --------------------------------------------------------------------------- # Tests for the non-gated fused-experts path (NemotronH NemotronHExperts): @@ -2690,6 +2995,239 @@ def test_native_fp8_experts_reject_packed_weight_updates(cfg, algorithm): assert not torch.equal(getattr(experts, name), weight) +# --------------------------------------------------------------------------- +# A class with its own wrapper keeps eager; the interception check spares unreached experts +# --------------------------------------------------------------------------- +def test_gpt_oss_calibrates_down_proj_input_on_the_eager_forward(): + """GPT-OSS's own wrapper quantizes down_proj's input by intercepting the eager forward's + bmm/matmul, which the grouped path never calls: the model keeps main's forced eager.""" + transformers = pytest.importorskip("transformers") + config = transformers.GptOssConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + num_local_experts=4, + num_experts_per_tok=2, + max_position_embeddings=64, + layer_types=["full_attention"], + ) + torch.manual_seed(0) + model = transformers.GptOssForCausalLM(config).eval() + + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randint(0, 64, (2, 8)))) + + assert model.config._experts_implementation == "eager" + assert model.model.layers[0].mlp.experts.down_proj_input_quantizer.amax is not None + + +_INT8_EXPERTS_CFG = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*proj_input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + {"quantizer_name": "*proj_weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": "max", +} + + +class _TextAndVisionMoEModel(nn.Module): + """Two MoE blocks, as Ernie-4.5-VL has: text-only calibration never reaches the vision one.""" + + def __init__(self): + super().__init__() + self.text_moe = _SyntheticSparseMoeBlock() + self.vision_moe = _SyntheticSparseMoeBlock() + + def forward(self, x): + return self.text_moe(x) + + +class _BypassingExperts(_SyntheticFusedExperts): + """A forward that, while ``bypass``, reaches neither ``F.linear`` nor ``_grouped_linear``.""" + + bypass = True + + def forward(self, hidden_states, top_k_index, top_k_weights): + if self.bypass: + return torch.zeros_like(hidden_states) + return super().forward(hidden_states, top_k_index, top_k_weights) + + +class TestFusedExpertsInterceptionCheck: + def test_experts_no_calibration_token_reached_are_not_an_error(self): + model = _TextAndVisionMoEModel() + TestFusedExpertsCalibration._cleanup_registry(_SyntheticFusedExperts) + try: + mtq.quantize(model, _INT8_EXPERTS_CFG, lambda m: m(torch.randn(2, 4, HIDDEN_DIM))) + finally: + TestFusedExpertsCalibration._cleanup_registry(_SyntheticFusedExperts) + + assert model.text_moe.experts.gate_up_proj_input_quantizer.amax is not None + assert model.vision_moe.experts.gate_up_proj_input_quantizer.amax is None + + @pytest.mark.parametrize("rows", [0, 4], ids=["no-rows", "rows-routed-to-other-ranks"]) + def test_experts_that_make_no_projection_call_are_not_an_error(self, rows): + """The eager path calls F.linear only for the experts some row was routed to, which is not a + bypassed backend: none for a forward handed no rows (a modality's experts on its empty token + subset), nor, under transformers' expert parallelism, for one whose rows were all routed to + other ranks' experts, which the router masks to num_experts.""" + model = _TinyMoEModel() + TestFusedExpertsCalibration._cleanup_registry(_SyntheticFusedExperts) + + def forward_loop(m): + other_ranks = torch.full((rows, TOP_K), NUM_EXPERTS) + m.moe.experts(torch.randn(rows, HIDDEN_DIM), other_ranks, torch.ones(rows, TOP_K)) + + try: + mtq.quantize(model, _INT8_EXPERTS_CFG, forward_loop) + finally: + TestFusedExpertsCalibration._cleanup_registry(_SyntheticFusedExperts) + + def test_a_grouped_forward_counts_whatever_its_routing_says(self): + """The grouped forwards call `_grouped_linear` whatever rows they are handed, and torch_a2a's + routing indices are global expert ids, so the eager routing test does not apply there: a + grouped experts forward that reached neither hook fails loudly even when every row's index + is past this module's experts.""" + model = _TinyMoEModel() + model.moe.experts = _BypassingExperts() + model.moe.experts.config = _ExpertsConfig("grouped_mm") + + def forward_loop(m): + beyond = torch.full((4, TOP_K), NUM_EXPERTS) + m.moe.experts(torch.randn(4, HIDDEN_DIM), beyond, torch.ones(4, TOP_K)) + + try: + with pytest.raises(RuntimeError, match="collected nothing"): + mtq.quantize(model, _INT8_EXPERTS_CFG, forward_loop) + finally: + TestFusedExpertsCalibration._cleanup_registry(_BypassingExperts) + + def test_forwards_after_a_hook_fired_are_not_counted(self): + """Once a hook has fired the check passes, so later forwards skip the routing test, which + syncs the host.""" + model = _TinyMoEModel() + counts = [] + + def forward_loop(m): + for _ in range(2): + m(torch.randn(2, 4, HIDDEN_DIM)) + counts.append(m.moe.experts._forward_count) + + TestFusedExpertsCalibration._cleanup_registry(_SyntheticFusedExperts) + try: + mtq.quantize(model, _INT8_EXPERTS_CFG, forward_loop) + finally: + TestFusedExpertsCalibration._cleanup_registry(_SyntheticFusedExperts) + assert counts == [1, 1] + + def test_transformers_experts_take_the_routing_indices_second(self): + """forward() reads the routing indices as the second positional argument, after the hidden + states, to count only the forwards that route rows to this rank's experts: an experts class + taking them in another place would have its routing weights compared with num_experts. + GPT-OSS-style experts name them router_indices.""" + checked = 0 + for path in sorted((Path(transformers.__file__).parent / "models").glob("*/modeling_*.py")): + source = path.read_text(encoding="utf-8") + if "use_experts_implementation" not in source: + continue + for node in ast.parse(source).body: + if not isinstance(node, ast.ClassDef) or not any( + "use_experts_implementation" in ast.unparse(d) for d in node.decorator_list + ): + continue + methods = {n.name: n for n in node.body if isinstance(n, ast.FunctionDef)} + if "forward" not in methods: # inherited: checked on the class that defines it + continue + params = [arg.arg for arg in methods["forward"].args.args] + assert params[1] == "hidden_states", (node.name, params) + assert params[2] in ("top_k_index", "router_indices"), (node.name, params) + checked += 1 + assert checked + + @pytest.mark.parametrize("reached_before", [False, True], ids=["first-pass", "later-pass"]) + def test_experts_whose_forward_ran_past_both_hooks_still_fail_loudly(self, reached_before): + """Also when an earlier calibration reached them: each judges its own forwards. The model + comes back in the mode it went in, for a caller that catches the error.""" + model = _TinyMoEModel() + model.moe.experts = _BypassingExperts() + + def forward_loop(m): + m(torch.randn(2, 4, HIDDEN_DIM)) + + try: + if reached_before: + model.moe.experts.bypass = False + mtq.quantize(model, _INT8_EXPERTS_CFG, forward_loop) + model.moe.experts.bypass = True + model.train() + with pytest.raises(RuntimeError, match="collected nothing"): + if reached_before: + mtq.calibrate(model, "max", forward_loop) + else: + mtq.quantize(model, _INT8_EXPERTS_CFG, forward_loop) + assert model.training + finally: + TestFusedExpertsCalibration._cleanup_registry(_BypassingExperts) + + def test_a_plugin_that_failed_to_load_leaves_calibration_alone(self, monkeypatch): + """import_plugin only warns when the HF plugin's body raises, and such a plugin installed no + fused-experts wrappers: calibration has nothing to check, rather than raising its error.""" + name = "modelopt.torch.quantization.plugins.huggingface" + + class _BrokenPlugin: + @staticmethod + def find_spec(fullname, path=None, target=None): + if fullname == name: + raise AttributeError("a transformers version skew") + + monkeypatch.delitem(sys.modules, name) + monkeypatch.setattr(sys, "meta_path", [_BrokenPlugin(), *sys.meta_path]) + model = nn.Sequential(nn.Linear(HIDDEN_DIM, HIDDEN_DIM)) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(torch.randn(2, HIDDEN_DIM))) + + @pytest.mark.parametrize("algorithm", ["max", "local_hessian", "gptq"]) + def test_weight_only_experts_need_interception_when_their_weights_read_inputs(self, algorithm): + """Under max a weight quantizer calibrates on its weight alone, so weight-only experts whose + forward ran past both hooks (a dense activation quantizer runs the forward) are calibrated; + local Hessian and GPTQ calibrate their weights from their inputs, which they never saw.""" + model = _MixedPrecisionModel() + model.moe.experts = _BypassingExperts() + cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*proj_weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*dense.input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + ], + "algorithm": algorithm, + } + + def forward_loop(m): + m(torch.randn(2, 4, HIDDEN_DIM)) + + try: + if algorithm == "max": + mtq.quantize(model, cfg, forward_loop) + else: + with pytest.raises(RuntimeError, match="collected nothing"): + mtq.quantize(model, cfg, forward_loop) + finally: + TestFusedExpertsCalibration._cleanup_registry(_BypassingExperts) + if algorithm == "max": + experts = model.moe.experts + for quantizers, weight in ( + (experts.gate_up_proj_weight_quantizers, experts.gate_up_proj), + (experts.down_proj_weight_quantizers, experts.down_proj), + ): + for idx, quantizer in enumerate(quantizers): + torch.testing.assert_close( + quantizer.amax.flatten(), weight[idx].abs().amax(dim=1) + ) + + # --------------------------------------------------------------------------- # Fused experts under expert parallelism # --------------------------------------------------------------------------- @@ -2924,3 +3462,26 @@ def recording(helper, module, name, *args, weight=None, **kwargs): def test_ep_gptq_builds_its_expert_helpers_from_this_ranks_experts(cpu_dist_workers): """A GPTQ helper only needs its expert's shape, so it must not gather other ranks' experts.""" cpu_dist_workers.run(_ep_gptq_helpers_take_local_experts_job) + + +def _ep_interception_check_job(rank, size): + """Only rank 0's experts ran past both hooks: every rank of their group raises, not just rank 0, + and none does once every rank's experts were reached.""" + config = _ep_experts_config() + mesh = init_device_mesh("cpu", (size,)) + whole = _experts_with_distinct_scales(config) + model = _ExpertsHolder(_sharded_like_transformers(whole, config, mesh, size)) + mtq.quantize(model, _INT8_EXPERTS_CFG) + experts = model.experts + experts._forward_count, experts._intercept_count = 1, int(rank != 0) + with pytest.raises(RuntimeError, match="collected nothing") as raised: + assert_fused_experts_intercepted(model) + assert ("on another rank" in str(raised.value)) == (rank != 0) + experts._intercept_count = 1 + assert_fused_experts_intercepted(model) + + +def test_ep_interception_check_raises_on_every_rank_of_the_group(cpu_dist_workers): + """A rank whose experts saw no rows must not go on to wait in its next collective while the + others raise: the check reduces its verdict over the experts' expert-parallel group.""" + cpu_dist_workers.run(_ep_interception_check_job) diff --git a/tests/unit/torch/utils/test_moe_a2a.py b/tests/unit/torch/utils/test_moe_a2a.py index 5d2498ae86d..9dea6d86c67 100644 --- a/tests/unit/torch/utils/test_moe_a2a.py +++ b/tests/unit/torch/utils/test_moe_a2a.py @@ -41,6 +41,7 @@ from packaging.version import Version from torch import nn +import modelopt.torch.quantization as mtq import modelopt.torch.utils.moe as moe from modelopt.torch.utils.moe import ( _choose_ep_degree, @@ -99,10 +100,10 @@ def test_a2a_moe_keeps_the_dtype_it_is_given(cpu_dist_workers, dtype, tol): cpu_dist_workers.run(_keeps_dtype_job, dtype=dtype, tol=tol) -def _chunked_hops_job(rank, size): +def _chunked_hops_job(rank, size, grad): """Shrinking the per-hop element limit splits every 2-D hop along the hidden dim into slices within it -- the guard against int32 offset overflow, at a size a unit test can afford -- - without changing the result.""" + without changing the result or, with a graph, the gradients.""" gen = torch.Generator().manual_seed(rank) x = torch.randn(T, H, generator=gen) # Every rank sends and receives T * TOP_K = 12 rows. A limit of 34 holds 2 of their H = 8 @@ -110,7 +111,7 @@ def _chunked_hops_job(rank, size): ti = (torch.arange(T * TOP_K) % E).reshape(T, TOP_K) tw = torch.rand(T, TOP_K, generator=gen) limit = 34 - hops = [] # the elements each 2-D hop moves + hops = [] # the elements each 2-D hop moves, the backward's included real = dist.all_to_all_single def recording(output, input_, *args, **kwargs): @@ -118,19 +119,27 @@ def recording(output, input_, *args, **kwargs): hops.append(max(output.numel(), input_.numel())) return real(output, input_, *args, **kwargs) + def run(): + x_ = x.clone().requires_grad_(grad) + with torch.set_grad_enabled(grad): + out = a2a_moe(dist.group.WORLD, x_, ti, tw, E // size, _identity_experts) + if grad: + out.square().sum().backward() + return out.detach(), x_.grad + dist.all_to_all_single = recording try: - with torch.no_grad(): - whole = a2a_moe(dist.group.WORLD, x, ti, tw, E // size, _identity_experts) - whole_hops = hops.copy() - hops.clear() - moe._A2A_MAX_HOP_ELEMENTS = limit - chunked = a2a_moe(dist.group.WORLD, x, ti, tw, E // size, _identity_experts) + whole, whole_grad = run() + whole_hops = hops.copy() + hops.clear() + moe._A2A_MAX_HOP_ELEMENTS = limit + chunked, chunked_grad = run() finally: dist.all_to_all_single = real moe._A2A_MAX_HOP_ELEMENTS = 2**31 - 1 assert max(hops) <= limit < max(whole_hops) torch.testing.assert_close(chunked, whole, atol=0, rtol=0) + torch.testing.assert_close(chunked_grad, whole_grad, atol=0, rtol=0) def _bounded_combine_job(rank, size): @@ -175,8 +184,9 @@ def test_the_combine_works_a_slice_of_tokens_at_a_time(cpu_dist_workers): cpu_dist_workers.run(_bounded_combine_job) -def test_a2a_moe_chunks_hops_past_the_element_limit(cpu_dist_workers): - cpu_dist_workers.run(_chunked_hops_job) +@pytest.mark.parametrize("grad", [False, True], ids=["forward", "backward"]) +def test_a2a_moe_chunks_hops_past_the_element_limit(cpu_dist_workers, grad): + cpu_dist_workers.run(_chunked_hops_job, grad=grad) def _routing_skew_job(rank, size, routing): @@ -374,20 +384,56 @@ def _assert_experts_backward_matches(rank, whole, local, owned, size, to_rank0=F ) -def _real_experts_backward_job(rank, size): +def _fake_quantized_like_the_whole_module(rank, size, whole, local, owned): + """Both with ModelOpt's fake quantizers, the local ones set to the whole module's calibrated + scales so the two compute the same quantized function; returns them converted. Calibrating + through the torch_a2a forward sees only the tokens dispatched to this rank's experts.""" + cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*proj_input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + {"quantizer_name": "*proj_weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": "max", + } + x, weights, index = _routed_tokens(size) + tok = slice(rank * T, (rank + 1) * T) + whole_holder, local_holder = torch.nn.ModuleList([whole]), torch.nn.ModuleList([local]) + whole.config._experts_implementation = "grouped_mm" + mtq.quantize(whole_holder, cfg, lambda m: m[0](x, index, weights)) + local.config._experts_implementation = "torch_a2a" + mtq.quantize( + local_holder, cfg, lambda m: m[0](x[tok], index[tok], weights[tok], dist.group.WORLD) + ) + whole, local = whole_holder[0], local_holder[0] + for name in ("gate_up_proj_input_quantizer", "down_proj_input_quantizer"): + getattr(local, name).amax = getattr(whole, name).amax + for name in ("gate_up_proj_weight_quantizers", "down_proj_weight_quantizers"): + for i, quantizer in enumerate(getattr(local, name)): + quantizer.amax = getattr(whole, name)[owned.start + i].amax + return whole, local + + +def _real_experts_backward_job(rank, size, quantized): moe.register_torch_a2a() whole, local, owned = _split_qwen3_moe_experts(rank, size) + if quantized: + whole, local = _fake_quantized_like_the_whole_module(rank, size, whole, local, owned) _assert_experts_backward_matches(rank, whole, local, owned, size) -def test_fused_experts_backward_through_torch_a2a_matches_the_whole_module(cpu_dist_workers): +@pytest.mark.parametrize("quantized", [False, True], ids=["unquantized", "fake-quantized"]) +def test_fused_experts_backward_through_torch_a2a_matches_the_whole_module( + cpu_dist_workers, quantized +): """transformers' fused experts, split over two ranks, train through the torch_a2a forward exactly as the whole module does on one process: same output, and same gradients for the hidden states, - the routing weights and each rank's experts.""" + the routing weights and each rank's experts. With ModelOpt's fake quantizers inserted, that is + QAT under expert parallelism.""" transformers = pytest.importorskip("transformers") if Version(transformers.__version__) < Version("5.6"): pytest.skip("expert parallelism needs transformers >= 5.6") - cpu_dist_workers.run(_real_experts_backward_job) + cpu_dist_workers.run(_real_experts_backward_job, quantized=quantized) def _fp32_routing_weights_job(rank, size):