From 3d9f398f5901e060ed22d490a032ace057199f2e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 00:10:05 +0000 Subject: [PATCH 01/42] Bring in the synthetic experts' torch_a2a dispatch with the tests that drive it The expert-parallel captures and transport tests this PR adds run the grouped test experts on torch_a2a over an expert-parallel group. Signed-off-by: Shengliang Xu --- tests/unit/torch/quantization/plugins/test_fused_experts.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 22c91454332..2aa3402a91d 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -70,6 +70,7 @@ ) 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(): @@ -322,6 +323,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,6 +383,10 @@ def forward(self, hidden_states, top_k_index, top_k_weights): return _moe_integration.grouped_mm_experts_forward( self, hidden_states, top_k_index, top_k_weights ) + if self.config._experts_implementation == "torch_a2a": + return torch_a2a_experts_forward( + self, hidden_states, top_k_index, top_k_weights, self.config.ep_group + ) return self._eager_forward(hidden_states, top_k_index, top_k_weights) From 08297f8eb2560a864a70eb79e9e44fa87f388839 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 23:15:49 +0000 Subject: [PATCH 02/42] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager normalize_experts_impl_on_the_fly replaces force_eager_experts_impl_on_the_fly: eager, grouped_mm (which transformers selects by default) and torch_a2a, the expert-parallel forward, are kept, and a backend nothing intercepts (batched_mm, sonicmoe, kernels) is redirected to grouped_mm. A model holding experts that quantize on the eager forward only keeps main's forced eager and is refused under torch_a2a, as is a class with its own wrapper (GPT-OSS). The weight-input captures (local Hessian, imatrix, GPTQ) split the grouped forward's expert-sorted batch back per expert by its offsets, and refuse a forward on the other path than the one they were set up for. Under expert parallelism a fused expert weight is a DTensor sharded over the experts, so the quantizers, captures and folds index this rank's own experts, and the FSDP2 weight access gathers only the parameters FSDP2 manages. Calibrated activation scales can move with the backend; the changelog says how to calibrate as before. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + modelopt/torch/quantization/model_calib.py | 91 ++++- .../torch/quantization/plugins/huggingface.py | 130 ++++--- modelopt/torch/utils/moe.py | 31 +- .../plugins/test_fused_experts.py | 336 ++++++++++++++---- tests/unit/torch/utils/test_moe_a2a.py | 76 +++- 6 files changed, 512 insertions(+), 153 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 371bea5f656..8b0c2ddef37 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -55,6 +55,7 @@ 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. +- Hugging Face MoE calibration now runs fused experts on the ``grouped_mm`` backend Transformers selects by default instead of forcing ``eager``, so calibrated activation scales can differ from earlier releases: some of Qwen3-30B-A3B's moved by up to about 10%. Load the model with ``experts_implementation="eager"`` to calibrate exactly as earlier releases did. - 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/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 3a210a3545a..33cff29d026 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,36 @@ def _refuse_packed_fused_expert_weights(name, module, method, updates): ) +def _check_experts_path(experts, grouped: bool) -> None: + """Raise if ``experts`` ran on the other path than the one its capture hook was set up for. + + A capture hook is chosen by the path ``config._experts_implementation`` names; the interceptors + record the path that ran (``_grouped_offs`` set on the grouped one, cleared on the eager one), + and a capture fed by the other path would attribute inputs to the wrong experts. + """ + ran_grouped = getattr(experts, "_grouped_offs", None) is not None + if ran_grouped != grouped: + impl = getattr(getattr(experts, "config", None), "_experts_implementation", None) + raise RuntimeError( + f"{type(experts).__name__} ran its experts on the {'grouped' if ran_grouped else 'eager'} " + f"path while its config selects {impl!r}; the weight-input capture would attribute " + "its inputs to the wrong experts." + ) + + +def _expert_slices(offs: torch.Tensor): + """``(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``): an expert without rows is + skipped, and rows beyond ``offs[-1]`` (expert-parallel sentinels) belong to no expert. + """ + start = 0 + for idx, end in enumerate(offs.tolist()): + if end > start: + yield idx, start, end + start = end + + class _WeightInputHook: """Identify active input collectors so auxiliary modules can request a calibration forward.""" @@ -1146,9 +1175,10 @@ 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 = [] @@ -1156,16 +1186,40 @@ def _make_expert_hook(expert_module, weight_name, quantizers, enabled): def _expert_hook(_input_quantizer, args): if not args: return + _check_experts_path(expert_module, grouped=False) 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``). + # this rank's own experts as the grouped hook below is (see ``_local_experts``). weight = expert_module._local_experts(getattr(expert_module, weight_name)) capture(quantizers[idx], weight[idx], args[0]) return _expert_hook + def _make_grouped_expert_hook(expert_module, weight_name, quantizers, enabled): + """Per-expert capture for the grouped dispatchers (``grouped_mm``, ``torch_a2a``). + + Both feed the input quantizer a single expert-sorted batch ``(S, cin)`` rather than + one call per expert, so we split it back into per-expert chunks using ``_grouped_offs`` + (cumulative tokens-per-expert, stashed by the _QuantFusedExperts grouped interceptor). + Rows beyond ``offs[-1]`` (EP sentinels) belong to no expert and are skipped. + """ + + def _expert_hook(_input_quantizer, args): + if not args: + return + _check_experts_path(expert_module, grouped=True) + offs = expert_module._grouped_offs + x = args[0] + # This rank's experts, in quantizer order: 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)) + for idx, start, end in _expert_slices(offs): + if idx in enabled: + capture(quantizers[idx], weight[idx], x[start:end]) + + return _expert_hook + for name, module in names.name_to_module.items(): if is_quantized_linear(module) and isinstance(module.weight_quantizer, TensorQuantizer): with enable_weight_access_and_writeback(module, model, names): @@ -1198,11 +1252,16 @@ def _dense_hook(linear, args): enabled = {i for i, q in enumerate(quantizers) if q.is_enabled} if not enabled: continue + # grouped_mm and torch_a2a feed one expert-sorted batch per projection; eager + # feeds one call per expert keyed by _current_expert_idx. + make_hook = ( + _make_grouped_expert_hook + if getattr(module, "_running_grouped_mm", False) + else _make_expert_hook + ) handles.append( input_quantizer.register_forward_pre_hook( - _WeightInputHook( - _make_expert_hook(module, weight_name, quantizers, enabled) - ) + _WeightInputHook(make_hook(module, weight_name, quantizers, enabled)) ) ) return handles @@ -2588,6 +2647,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): + grouped = getattr(experts, "_running_grouped_mm", False) + _check_experts_path(experts, grouped) + if grouped: + # 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): + 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) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 8a2d0001086..652270b4e02 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -803,17 +803,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. + kernel) backends are not intercepted; ``normalize_experts_impl_on_the_fly`` + redirects them to ``grouped_mm`` before calibration. 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 @@ -836,7 +837,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 +898,19 @@ 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 current grouped_mm forward, + # stashed so local-Hessian capture can split the expert-sorted batch. + self._register_temp_attribute("_grouped_offs", None) + + @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. @@ -924,9 +938,10 @@ 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 + eager per-expert ``F.linear`` semantics. Returns this rank's experts untouched when every quantizer would pass its slice through unchanged, so passthrough stays exact (and cheap). """ + weight = self._local_experts(weight) # A disabled quantizer still applies its pre-quant scale and rotation, as on the eager path. if all( isinstance(q, TensorQuantizer) @@ -973,6 +988,8 @@ 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): + # 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) @@ -1006,16 +1023,17 @@ def _quantized_linear(input, weight, bias=None): 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. + # 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 ) @@ -1572,8 +1590,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 +1646,77 @@ 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). ``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. + supported_impls = (None, "eager", "grouped_mm", "torch_a2a") + # 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." + ) + for cfg in configs: + if eager_only: + cfg._experts_implementation = "eager" + elif cfg._experts_implementation not in supported_impls: + cfg._experts_implementation = "grouped_mm" def _is_supported_hf_model(model): @@ -1784,7 +1834,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/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 2aa3402a91d..1f32689682d 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -23,6 +23,7 @@ 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 @@ -64,7 +65,7 @@ _is_sparse_sequaential_moe_block, _QuantFusedExperts, _QuantNonGatedFusedExperts, - force_eager_experts_impl_on_the_fly, + normalize_experts_impl_on_the_fly, register_fused_experts_on_the_fly, register_sparse_moe_on_the_fly, ) @@ -558,28 +559,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() @@ -1344,7 +1323,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.""" @@ -1369,49 +1348,115 @@ 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" + 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.""" - text_cfg = _StubConfig(impl="grouped_mm") - vision_cfg = _StubConfig(impl="bmm") + """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") + def test_forces_eager_for_experts_quantized_on_the_eager_forward_only(self): + """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.""" + model = _NonMoEModelWithConfig(_StubConfig(impl="grouped_mm")) + 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" + + 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): + """``eager``, ``grouped_mm`` and an explicit None (the model's eager ``F.linear`` forward) are + supported, so they are left as they are.""" + 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 # --------------------------------------------------------------------------- @@ -1820,6 +1865,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 @@ -1914,16 +2027,15 @@ 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 @@ -2051,10 +2163,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" @@ -2075,28 +2187,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, @@ -2117,6 +2250,36 @@ class _EagerOnlyExperts(_QuantFusedExperts): finally: self._cleanup_registry(_SyntheticGroupedExperts) + @pytest.mark.parametrize( + ("selected", "ran"), [("grouped_mm", "eager"), ("eager", "grouped_mm")] + ) + def test_a_capture_refuses_a_forward_on_the_other_path(self, selected, ran): + """The capture hooks are chosen by the path the config selects; fed by the other path they + would attribute inputs to the wrong experts, so they raise instead.""" + model = _GroupedMoEModel(selected) + try: + mtq.quantize(model, self._quant_cfg(), self._forward_loop_all_experts) + handles = _register_weight_input_hooks(model, module_name_maps(model), lambda *a: None) + model.config._experts_implementation = ran + linear, grouped_linear = F.linear, _moe_integration._grouped_linear + try: + with pytest.raises(RuntimeError, match="wrong experts"): + self._forward_loop_all_experts(model) + # The forward's interceptors are undone even though it raised. + assert (F.linear, _moe_integration._grouped_linear) == (linear, grouped_linear) + finally: + for handle in handles: + handle.remove() + finally: + self._cleanup_registry(_SyntheticGroupedExperts) + + @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): @@ -2696,6 +2859,35 @@ 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 +# --------------------------------------------------------------------------- +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 + + # --------------------------------------------------------------------------- # Fused experts under expert parallelism # --------------------------------------------------------------------------- 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): From e474021c7ab712c39aea22f4f256d676ce7e56e4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 23:17:07 +0000 Subject: [PATCH 03/42] Raise when a fused-experts forward bypasses both interception hooks register_fused_experts_on_the_fly claims experts on their parameters' shapes, but the quantizers see data only if the forward reaches F.linear or _grouped_linear. A backend that reaches neither (a hand-written bmm forward on a class without @use_experts_implementation, which no config change redirects) left every quantizer without data, and the export fell back to weight-derived scales under a warning that blamed the calibration set's size. Each wrapper now counts its forwards and its hooks' calls, and calibrate() raises, naming the modules, when one ran a forward its hooks never saw. Experts no calibration token reached, or that the recipe leaves unquantized, are not reported. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 30 +++++++ .../torch/quantization/plugins/huggingface.py | 77 ++++++++++++++++ .../plugins/test_fused_experts.py | 88 ++++++++++++++++++- 3 files changed, 194 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 0a69ae9ede9..834c2e65efc 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -143,6 +143,24 @@ 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 + # 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 +187,18 @@ 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. + try: + from .plugins.huggingface import assert_fused_experts_intercepted + except ImportError: + pass # no transformers -> no HF fused-experts modules to check + else: + assert_fused_experts_intercepted(model) + return model diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 652270b4e02..45a32d7dc52 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -901,6 +901,14 @@ def _setup(self): # Cumulative tokens-per-expert offsets of the current grouped_mm forward, # stashed so local-Hessian capture can split the expert-sorted batch. 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) @property def _running_grouped_mm(self) -> bool: @@ -988,6 +996,7 @@ 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: @@ -1029,6 +1038,7 @@ def _quantized_grouped_linear(input, weight, offs, bias=None, is_transposed=Fals 'grouped_mm forward ran. Load the model with experts_implementation="eager", ' "and without expert parallelism, whose forward runs it too." ) + 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. @@ -1077,6 +1087,11 @@ def _quantized_weight(self, name, quantizers, idx, dtype): def forward(self, *args, **kwargs): self._down_proj_linear = False + hidden_states = args[0] if args else kwargs.get("hidden_states") + # A forward with no rows makes no projection call (the eager path loops over the experts + # some token reached), so only one with rows is work the interception check expects. + if hidden_states is None or hidden_states.numel(): + self._forward_count += 1 self._stacked_weights = {} try: return super().forward(*args, **kwargs) @@ -1719,6 +1734,68 @@ def _configs(cfg): cfg._experts_implementation = "grouped_mm" +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 at least one enabled quantizer are reported: + an experts module no calibration token reached (vision experts under text-only calibration) + falls back to weight-derived scales as before, and one the recipe leaves unquantized loses + nothing by not being intercepted. + """ + missed = [] + for name, module in model.named_modules(): + if not isinstance(module, _QuantFusedExperts): + continue + # 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 + ): + continue + + first = module._first_proj_attr + quantizers = [] + for attr in (f"{first}_weight_quantizers", "down_proj_weight_quantizers"): + quantizers.extend(getattr(module, attr, None) or []) + for attr in (f"{first}_input_quantizer", "down_proj_input_quantizer"): + quantizer = getattr(module, attr, None) + if quantizer is not None: + quantizers.append(quantizer) + if not any(getattr(q, "is_enabled", False) for q in quantizers): + continue + + impl = getattr(getattr(module, "config", None), "_experts_implementation", None) + missed.append(f" {name} ({type(module).__name__}, experts_implementation={impl!r})") + + 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 not one of their quantizers was invoked and " + "every expert would be exported with a weight-derived fallback scale.\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`." + ) + + def _is_supported_hf_model(model): """Check if the model a valid model for transformers quantization specific support.""" supported_models = [transformers.PreTrainedModel] diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 1f32689682d..4842af200cf 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -2860,7 +2860,7 @@ def test_native_fp8_experts_reject_packed_weight_updates(cfg, algorithm): # --------------------------------------------------------------------------- -# A class with its own wrapper keeps eager +# 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 @@ -2888,6 +2888,92 @@ def test_gpt_oss_calibrates_down_proj_input_on_the_eager_forward(): 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 + + def test_experts_called_with_no_rows_are_not_an_error(self): + """A forward handed no tokens makes no projection call on the eager path, which is not a + bypassed backend: a modality's experts called on its empty token subset, say.""" + model = _TinyMoEModel() + TestFusedExpertsCalibration._cleanup_registry(_SyntheticFusedExperts) + + def forward_loop(m): + no_rows = torch.zeros(0, HIDDEN_DIM) + m.moe.experts(no_rows, torch.zeros(0, TOP_K, dtype=torch.long), torch.zeros(0, TOP_K)) + + try: + mtq.quantize(model, _INT8_EXPERTS_CFG, forward_loop) + finally: + TestFusedExpertsCalibration._cleanup_registry(_SyntheticFusedExperts) + + @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) + + # --------------------------------------------------------------------------- # Fused experts under expert parallelism # --------------------------------------------------------------------------- From d169ad5913bef3002f000f3881a25c1acfddf1e9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 22:54:54 +0000 Subject: [PATCH 04/42] Keep grouped_mm, and redirect to it, only where _grouped_linear can be intercepted The grouped_mm interception is installed only when transformers has _grouped_linear, but normalize_experts_impl_on_the_fly kept grouped_mm and redirected every unsupported backend to it unconditionally. One capability constant now gates all three; without it, unsupported backends fall back to eager, as main forced. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/huggingface.py | 21 ++++++++++++++----- .../plugins/test_fused_experts.py | 11 ++++++++++ 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 45a32d7dc52..b25a762495e 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -58,6 +58,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 @@ -1026,7 +1032,7 @@ 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): @@ -1675,7 +1681,8 @@ def normalize_experts_impl_on_the_fly(model): 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). ``eager``, ``grouped_mm`` and ``torch_a2a`` are + 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 @@ -1693,8 +1700,12 @@ def normalize_experts_impl_on_the_fly(model): 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. - supported_impls = (None, "eager", "grouped_mm", "torch_a2a") + # 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 @@ -1731,7 +1742,7 @@ def _configs(cfg): if eager_only: cfg._experts_implementation = "eager" elif cfg._experts_implementation not in supported_impls: - cfg._experts_implementation = "grouped_mm" + cfg._experts_implementation = fallback_impl def assert_fused_experts_intercepted(model): diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 4842af200cf..3a9240d079a 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -1356,6 +1356,17 @@ def test_redirects_unsupported_backend_to_grouped_mm(self): normalize_experts_impl_on_the_fly(model) assert cfg._experts_implementation == "grouped_mm" + 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") + 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 From c95a5c397ceaa6ebca8b2f332c044b2db97b7368 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 22:56:01 +0000 Subject: [PATCH 05/42] Count an eager experts forward only when it routes a row to this rank's experts Under transformers' expert parallelism the router masks other ranks' experts to num_experts, which the eager forward skips. A rank whose experts no calibration row was routed to made no F.linear call, so the interception check raised on that rank alone. The grouped paths call _grouped_linear whatever rows they are handed, so only the eager path reads the routing. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/huggingface.py | 16 +++++++++++++--- .../quantization/plugins/test_fused_experts.py | 15 +++++++++------ 2 files changed, 22 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index b25a762495e..18de12ba7eb 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1094,9 +1094,19 @@ def _quantized_weight(self, name, quantizers, idx, dtype): def forward(self, *args, **kwargs): self._down_proj_linear = False hidden_states = args[0] if args else kwargs.get("hidden_states") - # A forward with no rows makes no projection call (the eager path loops over the experts - # some token reached), so only one with rows is work the interception check expects. - if hidden_states is None or hidden_states.numel(): + top_k_index = args[1] if len(args) > 1 else kwargs.get("top_k_index") + # 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. + 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: diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 3a9240d079a..dc90c75fb1a 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -164,7 +164,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] @@ -2944,15 +2944,18 @@ def test_experts_no_calibration_token_reached_are_not_an_error(self): 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 - def test_experts_called_with_no_rows_are_not_an_error(self): - """A forward handed no tokens makes no projection call on the eager path, which is not a - bypassed backend: a modality's experts called on its empty token subset, say.""" + @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): - no_rows = torch.zeros(0, HIDDEN_DIM) - m.moe.experts(no_rows, torch.zeros(0, TOP_K, dtype=torch.long), torch.zeros(0, TOP_K)) + 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) From a5987855762f4ad16c58153094524c6ba067a1ae Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 22:58:58 +0000 Subject: [PATCH 06/42] Require interception for expert weight quantizers only when their inputs are captured Under max (and MSE) every weight quantizer calibrates on its weight alone in weight_only_quantize, so weight-only experts whose forward bypasses both hooks are calibrated. The interception check now counts an enabled weight quantizer only when the calibration captured the experts' inputs to calibrate their weights (GPTQ, local Hessian, imatrix), and its message says what falls back: the input scales, and an input-reading weight calibration. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 2 + modelopt/torch/quantization/model_quant.py | 1 + .../torch/quantization/plugins/huggingface.py | 31 ++++++++------- .../plugins/test_fused_experts.py | 38 +++++++++++++++++++ 4 files changed, 59 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 33cff29d026..cf103015990 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1252,6 +1252,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 # grouped_mm and torch_a2a feed one expert-sorted batch per projection; eager # feeds one call per expert keyed by _current_expert_idx. make_hook = ( @@ -2680,6 +2681,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 834c2e65efc..23bf72475ed 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -160,6 +160,7 @@ def forward_loop(model): 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 diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 18de12ba7eb..3e521c8319f 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -915,6 +915,9 @@ def _setup(self): # 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: @@ -1773,10 +1776,13 @@ def assert_fused_experts_intercepted(model): ``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 at least one enabled quantizer are reported: - an experts module no calibration token reached (vision experts under text-only calibration) - falls back to weight-derived scales as before, and one the recipe leaves unquantized loses - nothing by not being intercepted. + 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. """ missed = [] for name, module in model.named_modules(): @@ -1789,13 +1795,11 @@ def assert_fused_experts_intercepted(model): continue first = module._first_proj_attr - quantizers = [] - for attr in (f"{first}_weight_quantizers", "down_proj_weight_quantizers"): - quantizers.extend(getattr(module, attr, None) or []) - for attr in (f"{first}_input_quantizer", "down_proj_input_quantizer"): - quantizer = getattr(module, attr, None) - if quantizer is not None: - quantizers.append(quantizer) + 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 []) if not any(getattr(q, "is_enabled", False) for q in quantizers): continue @@ -1807,8 +1811,9 @@ def assert_fused_experts_intercepted(model): raise RuntimeError( "Fused-experts calibration collected nothing: the forward of the module(s) below reached " - "neither `F.linear` nor `_grouped_linear`, so not one of their quantizers was invoked and " - "every expert would be exported with a weight-derived fallback scale.\n" + "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 " diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index dc90c75fb1a..48f0ead99f9 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -2987,6 +2987,44 @@ def forward_loop(m): finally: TestFusedExpertsCalibration._cleanup_registry(_BypassingExperts) + @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 From 6b69d6b9935c4986d97315b24f5a250ca3ac1301 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 00:13:48 +0000 Subject: [PATCH 07/42] Stop counting experts forwards once an interception hook has fired The interception check reads the forward count only for a module whose hooks never fired, so once one has, the count is not needed. Testing an eager forward's routing syncs the host, so it now runs only until the first hook call of each calibration rather than on every forward, inference and QAT included. The check's outcome is unchanged. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/huggingface.py | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 3e521c8319f..d9ea318d4b5 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1096,21 +1096,23 @@ def _quantized_weight(self, name, quantizers, idx, dtype): def forward(self, *args, **kwargs): self._down_proj_linear = False - 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") # 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. - 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 + # 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) From 947729228e06ad34363fbb5b98fe9a75beceafbc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 00:24:24 +0000 Subject: [PATCH 08/42] Wrap the grouped calibration test's quantize calls at 100 columns Signed-off-by: Shengliang Xu --- .../unit/torch/quantization/plugins/test_fused_experts.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 48f0ead99f9..4fca82094a4 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -2038,8 +2038,12 @@ def test_calibration_matches_eager(self): grouped_model.load_state_dict(eager_model.state_dict()) try: - 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) + 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 From d02b3e3887bd1d163f556e952c0f9a938a2b6d32 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 00:49:25 +0000 Subject: [PATCH 09/42] Capture each expert's inputs on the path that ran, not the one configured The weight-input captures (local Hessian, imatrix, GPTQ) chose their hook from config._experts_implementation when they were registered and raised when the other path ran. The interceptors already record the path that ran: _grouped_offs is set by the grouped one and cleared by the eager one. One hook now reads it at call time, so the captures follow whichever forward ran and the refusal and the second hook are gone. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 71 +++++-------------- .../plugins/test_fused_experts.py | 23 +++--- 2 files changed, 33 insertions(+), 61 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index cf103015990..bca7190f9f1 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1130,23 +1130,6 @@ def _refuse_packed_fused_expert_weights(name, module, method, updates): ) -def _check_experts_path(experts, grouped: bool) -> None: - """Raise if ``experts`` ran on the other path than the one its capture hook was set up for. - - A capture hook is chosen by the path ``config._experts_implementation`` names; the interceptors - record the path that ran (``_grouped_offs`` set on the grouped one, cleared on the eager one), - and a capture fed by the other path would attribute inputs to the wrong experts. - """ - ran_grouped = getattr(experts, "_grouped_offs", None) is not None - if ran_grouped != grouped: - impl = getattr(getattr(experts, "config", None), "_experts_implementation", None) - raise RuntimeError( - f"{type(experts).__name__} ran its experts on the {'grouped' if ran_grouped else 'eager'} " - f"path while its config selects {impl!r}; the weight-input capture would attribute " - "its inputs to the wrong experts." - ) - - def _expert_slices(offs: torch.Tensor): """``(expert, start, end)`` for the rows of each expert in a batch sorted by expert. @@ -1183,40 +1166,31 @@ def _register_weight_input_hooks(model, names, capture, on_weight=None): handles: list = [] def _make_expert_hook(expert_module, weight_name, quantizers, enabled): - def _expert_hook(_input_quantizer, args): - if not args: - return - _check_experts_path(expert_module, grouped=False) - 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 as the grouped hook below is (see ``_local_experts``). - weight = expert_module._local_experts(getattr(expert_module, weight_name)) - capture(quantizers[idx], weight[idx], args[0]) - - return _expert_hook - - def _make_grouped_expert_hook(expert_module, weight_name, quantizers, enabled): - """Per-expert capture for the grouped dispatchers (``grouped_mm``, ``torch_a2a``). + """Per-expert capture for whichever forward ran. - Both feed the input quantizer a single expert-sorted batch ``(S, cin)`` rather than - one call per expert, so we split it back into per-expert chunks using ``_grouped_offs`` - (cumulative tokens-per-expert, stashed by the _QuantFusedExperts grouped interceptor). - Rows beyond ``offs[-1]`` (EP sentinels) belong to no expert and are skipped. + 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 - _check_experts_path(expert_module, grouped=True) + # 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 - x = args[0] - # This rank's experts, in quantizer order: an expert-parallel weight is a DTensor whose - # index names a global expert (see ``_local_experts``). + 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): if idx in enabled: - capture(quantizers[idx], weight[idx], x[start:end]) + capture(quantizers[idx], weight[idx], args[0][start:end]) return _expert_hook @@ -1253,16 +1227,11 @@ def _dense_hook(linear, args): if not enabled: continue module._weight_input_capture = True - # grouped_mm and torch_a2a feed one expert-sorted batch per projection; eager - # feeds one call per expert keyed by _current_expert_idx. - make_hook = ( - _make_grouped_expert_hook - if getattr(module, "_running_grouped_mm", False) - else _make_expert_hook - ) handles.append( input_quantizer.register_forward_pre_hook( - _WeightInputHook(make_hook(module, weight_name, quantizers, enabled)) + _WeightInputHook( + _make_expert_hook(module, weight_name, quantizers, enabled) + ) ) ) return handles @@ -2648,9 +2617,7 @@ 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): - grouped = getattr(experts, "_running_grouped_mm", False) - _check_experts_path(experts, grouped) - if grouped: + 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): diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 4fca82094a4..d0b57bcb1b6 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -2268,25 +2268,30 @@ class _EagerOnlyExperts(_QuantFusedExperts): @pytest.mark.parametrize( ("selected", "ran"), [("grouped_mm", "eager"), ("eager", "grouped_mm")] ) - def test_a_capture_refuses_a_forward_on_the_other_path(self, selected, ran): - """The capture hooks are chosen by the path the config selects; fed by the other path they - would attribute inputs to the wrong experts, so they raise instead.""" + 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.""" 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), lambda *a: None) + handles = _register_weight_input_hooks(model, module_name_maps(model), capture) model.config._experts_implementation = ran - linear, grouped_linear = F.linear, _moe_integration._grouped_linear try: - with pytest.raises(RuntimeError, match="wrong experts"): - self._forward_loop_all_experts(model) - # The forward's interceptors are undone even though it raised. - assert (F.linear, _moe_integration._grouped_linear) == (linear, grouped_linear) + self._forward_loop_all_experts(model) 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("impl", ["grouped_mm", "torch_a2a"]) def test_local_hessian_captures_per_expert_in_grouped_mode(self, cpu_dist_workers, impl): From 717e9b7afabf17bc0b5563745dd3a0c72077c8c3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 00:49:42 +0000 Subject: [PATCH 10/42] Say that the interception check is rank-local Under data parallelism the forward count depends on each rank's data, so a bypassing backend raises only on the ranks whose experts saw rows. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index d9ea318d4b5..51c25805981 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1785,6 +1785,9 @@ def assert_fused_experts_intercepted(model): 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. + + The check is rank-local: under data parallelism it raises on the ranks whose experts saw rows, + and a rank whose experts saw none goes on to its next collective. """ missed = [] for name, module in model.named_modules(): From 7323814f6b648f7f69575e3f035427f2281aeea7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 07:19:02 +0000 Subject: [PATCH 11/42] Drop the note that grouped_mm's weight-input captures skip its batch Here the captures split the grouped batch per expert, so local Hessian, imatrix and GPTQ calibrate the experts' weights from their inputs on grouped_mm as on eager. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 51c25805981..7e8e35cbe9e 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -822,11 +822,6 @@ class _QuantFusedExperts(_QuantFunctionalMixin): kernel) backends are not intercepted; ``normalize_experts_impl_on_the_fly`` redirects them to ``grouped_mm`` before calibration. - On ``grouped_mm`` the weight-input captures keyed on one expert (``local_hessian``, - ``imatrix``, ``gptq``) skip the batch, which holds every expert's rows, so those - algorithms calibrate the experts' weights without their inputs; max calibration - is unaffected. - The non-gated variant (``up_proj`` instead of ``gate_up_proj``, used by NemotronH) is handled by :class:`_QuantNonGatedFusedExperts` via the ``_first_proj_attr`` / ``_is_gated`` hooks below; each layout names the From 272e67cd1ff6485d22bdbac532a76bba308d16ad Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 07:19:06 +0000 Subject: [PATCH 12/42] Pin that forwards after a hook fired are not counted The interception check passes once a hook fired, so the forward stops counting and skips its routing test, which syncs the host. The test fails with that gate removed. Signed-off-by: Shengliang Xu --- .../quantization/plugins/test_fused_experts.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index d0b57bcb1b6..8301ab880bf 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -2971,6 +2971,24 @@ def forward_loop(m): finally: TestFusedExpertsCalibration._cleanup_registry(_SyntheticFusedExperts) + 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] + @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 From 512252827371615540e5ba37a31edd74e5a31a12 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 07:46:13 +0000 Subject: [PATCH 13/42] Clear the grouped offsets once the experts forward returns The weight-input captures pick their branch by _grouped_offs, which the interceptors set before every input quantizer call. Left set after the forward, a capture reached any other way would split its input by offsets from a finished forward instead of skipping it, and the module would keep the offsets tensor alive between forwards. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 7e8e35cbe9e..5af2a4b6aab 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1114,6 +1114,8 @@ def forward(self, *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. From 4fc2ed0b7372731a0497e1ff1213caf97bf5e7f8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 07:46:21 +0000 Subject: [PATCH 14/42] Say where a rank-local interception error leaves the other ranks A rank whose experts saw no rows does not raise; it waits in its next collective until the launcher stops it or the collective times out, so the error is in another rank's log. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 5af2a4b6aab..5e6dea40540 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1784,7 +1784,8 @@ def assert_fused_experts_intercepted(model): back to weight-derived scales as before. The check is rank-local: under data parallelism it raises on the ranks whose experts saw rows, - and a rank whose experts saw none goes on to its next collective. + 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. """ missed = [] for name, module in model.named_modules(): From 86d774b51dc48e6c64447ceb0d49381ed0e7bcf8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 18:31:52 +0000 Subject: [PATCH 15/42] Name normalize_experts_impl_on_the_fly in the FP8 experts wrapper's comment The comment #2745 added names the function by the name it has there; this PR renames it, and it still keeps such a model on eager. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 5e6dea40540..7361a36c18b 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1196,7 +1196,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 From 9f2b8b95305cbefe5bda66060d85fbc30d2f5503 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 18:49:53 +0000 Subject: [PATCH 16/42] Check fused-experts interception through the loaded HF plugin only calibrate() imported the HF plugin for its interception check and caught ImportError only. A plugin whose body fails otherwise (a transformers version skew raising AttributeError, say) is one import_plugin only warns about, and the import reran it and failed every calibration with that error. The check now reads the plugin from sys.modules: one that did not load installed no fused-experts wrappers to check. A test breaks the plugin's import and calibrates; it fails without the change. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 14 +++++++------- .../quantization/plugins/test_fused_experts.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 23bf72475ed..9392f555927 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 @@ -192,13 +193,12 @@ def forward_loop(model): # 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. - try: - from .plugins.huggingface import assert_fused_experts_intercepted - except ImportError: - pass # no transformers -> no HF fused-experts modules to check - else: - assert_fused_experts_intercepted(model) + # 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. + 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/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 8301ab880bf..c9cbee05eac 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -16,6 +16,7 @@ """Tests for _QuantFusedExperts: generic fused MoE quantization and export.""" import copy +import sys import warnings import weakref from contextlib import nullcontext @@ -3014,6 +3015,22 @@ def forward_loop(m): 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 From 55dd0ed816b95d3751eb6985bc1945c0e2374ae1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 19:40:44 +0000 Subject: [PATCH 17/42] Refuse grouped offsets the per-expert split cannot read The weight-input captures split an expert-sorted batch by the grouped offsets, read as exclusive cumulative ends, one per expert, as transformers' grouped_mm forward and a2a_moe build them. Offsets of another form (a leading zero, too few entries, ends past the rows) would credit experts with other experts' rows without a word; the split now raises. A test feeds each form and fails without the check. The _grouped_offs comment names every capture that reads it, and its other job: telling a capture which forward ran. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 23 ++++++++++++++----- .../torch/quantization/plugins/huggingface.py | 6 +++-- .../plugins/test_fused_experts.py | 11 +++++++++ 3 files changed, 32 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index bca7190f9f1..daa6fca22a2 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1130,14 +1130,21 @@ def _refuse_packed_fused_expert_weights(name, module, method, updates): ) -def _expert_slices(offs: torch.Tensor): +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``): an expert without rows is - skipped, and rows beyond ``offs[-1]`` (expert-parallel sentinels) belong to no 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: a leading zero, say, would credit each expert with the previous one's rows. """ + ends = offs.tolist() + if len(ends) != num_experts or (ends and ends[-1] > rows): + raise RuntimeError( + f"Grouped experts offsets {ends} do not split {rows} rows over {num_experts} experts." + ) start = 0 - for idx, end in enumerate(offs.tolist()): + for idx, end in enumerate(ends): if end > start: yield idx, start, end start = end @@ -1188,7 +1195,9 @@ def _expert_hook(_input_quantizer, args): 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): + 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]) @@ -2620,7 +2629,9 @@ 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): + 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]) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 7361a36c18b..49ac321c1fa 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -899,8 +899,10 @@ 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 current grouped_mm forward, - # stashed so local-Hessian capture can split the expert-sorted batch. + # 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 diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index c9cbee05eac..3b043bc1618 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -56,6 +56,7 @@ 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, local_hessian_calibrate, ) @@ -2294,6 +2295,16 @@ def capture(quantizer, _weight, inputs): 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]], ids=["leading-zero", "short", "past-the-rows"] + ) + 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 From 329225637370f541e3a5b87461c8b9648084f39b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 20:04:20 +0000 Subject: [PATCH 18/42] Refuse grouped offsets whose ends go back The offsets check covered the count of ends and the last one, but ends out of order passed it and credited an expert with another expert's rows. It now also requires the ends never to decrease, from zero, and the test gains that case, which fails without it. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 9 +++++++-- .../torch/quantization/plugins/test_fused_experts.py | 4 +++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index daa6fca22a2..b9cc1478b55 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1136,10 +1136,15 @@ def _expert_slices(offs: torch.Tensor, num_experts: int, rows: int): ``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: a leading zero, say, would credit each expert with the previous one's rows. + another form raise (a leading zero, say, 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): + 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." ) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 3b043bc1618..255553d81f5 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -2296,7 +2296,9 @@ def capture(quantizer, _weight, inputs): assert rows.get(id(q)) == 16 * TOP_K // NUM_EXPERTS @pytest.mark.parametrize( - "offs", [[0, 2, 2, 5], [2, 5], [2, 2, 7]], ids=["leading-zero", "short", "past-the-rows"] + "offs", + [[0, 2, 2, 5], [2, 5], [2, 2, 7], [2, 1, 5]], + ids=["leading-zero", "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 From 4076d235c902f82bbd73a8a270618713e08de3e6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 20:17:14 +0000 Subject: [PATCH 19/42] Name the inclusive offsets form the split refuses, not a leading zero A first end of zero is fine (expert 0 got no rows); what the check refuses is the inclusive [0, c0, c0 + c1, ...] form, one end too many. The docstring and the test case's id now say so. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_calib.py | 4 ++-- tests/unit/torch/quantization/plugins/test_fused_experts.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index b9cc1478b55..756b8d96c7d 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1136,8 +1136,8 @@ def _expert_slices(offs: torch.Tensor, num_experts: int, rows: int): ``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 (a leading zero, say, or ends out of order), since they would credit experts - with other experts' rows. + 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 ( diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 255553d81f5..4cbb5d078b1 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -2298,7 +2298,7 @@ def capture(quantizer, _weight, inputs): @pytest.mark.parametrize( "offs", [[0, 2, 2, 5], [2, 5], [2, 2, 7], [2, 1, 5]], - ids=["leading-zero", "short", "past-the-rows", "out-of-order"], + 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 From a3b1816c0350f178c5f9cc017b1e4392146b536e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 20:34:51 +0000 Subject: [PATCH 20/42] Add --experts_implementation to hf_ptq.py Calibration now runs fused experts on the backend transformers picks, grouped_mm by default, and the way back to earlier releases' scales is loading with experts_implementation="eager", which hf_ptq.py had no way to pass. The flag reaches every load path: get_model (the compressed- tensors branch included), the low-memory load, and, through the config the meta model is built from, the FSDP2 one. A test loads a tiny Qwen3-MoE with eager and fails without the plumbing; the changelog entry names the flag. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- examples/hf_ptq/example_utils.py | 4 ++++ examples/hf_ptq/hf_ptq.py | 20 +++++++++++++++++--- tests/examples/hf_ptq/test_example_utils.py | 10 +++++++++- 4 files changed, 31 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8b0c2ddef37..90daae84ecc 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -55,7 +55,7 @@ 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. -- Hugging Face MoE calibration now runs fused experts on the ``grouped_mm`` backend Transformers selects by default instead of forcing ``eager``, so calibrated activation scales can differ from earlier releases: some of Qwen3-30B-A3B's moved by up to about 10%. Load the model with ``experts_implementation="eager"`` to calibrate exactly as earlier releases did. +- Hugging Face MoE calibration now runs fused experts on the ``grouped_mm`` backend Transformers selects by default instead of forcing ``eager``, so calibrated activation scales can differ from earlier releases: some of Qwen3-30B-A3B's moved by up to about 10%. Load the model with ``experts_implementation="eager"`` (``--experts_implementation eager`` in ``hf_ptq.py``) to calibrate exactly as earlier releases did. - 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..fa5f91f0fad 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,8 @@ 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") + if experts_implementation is not None: + model_kwargs["experts_implementation"] = experts_implementation # 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 +723,7 @@ def has_pack_quantized_config(config): device_map="auto", trust_remote_code=trust_remote_code, dtype="auto", + experts_implementation=experts_implementation, ) 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..07adc6eacf7 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -320,9 +320,10 @@ def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False if args.use_fsdp2: - hf_config = AutoConfig.from_pretrained( - args.local_checkpoint_path, trust_remote_code=args.trust_remote_code - ) + config_kwargs = {"trust_remote_code": args.trust_remote_code} + if args.experts_implementation is not None: + config_kwargs["experts_implementation"] = args.experts_implementation + hf_config = AutoConfig.from_pretrained(args.local_checkpoint_path, **config_kwargs) validate_fsdp2_supported(args, hf_config) full_model = parallel_load_and_prepare_fsdp2( args.local_checkpoint_path, @@ -342,6 +343,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 +369,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 +1350,16 @@ 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." + ), + default=None, + type=str, + ) parser.add_argument( "--auto_quantize_checkpoint", type=str, 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"]) From 538d8dd7c23e0a87e5e6b6edce57bc6509a9cb9f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 20:36:40 +0000 Subject: [PATCH 21/42] Name calibrate's new fused-experts error in the changelog calibrate() now raises when a fused-experts forward ran but reached neither quantizer hook, where main exported weight-derived input scales with a warning: a script that used to produce a checkpoint now stops. The Backward Breaking Changes section says so. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 90daae84ecc..7cdf4c44f68 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -56,6 +56,7 @@ Changelog - 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. - Hugging Face MoE calibration now runs fused experts on the ``grouped_mm`` backend Transformers selects by default instead of forcing ``eager``, so calibrated activation scales can differ from earlier releases: some of Qwen3-30B-A3B's moved by up to about 10%. Load the model with ``experts_implementation="eager"`` (``--experts_implementation eager`` in ``hf_ptq.py``) to calibrate exactly 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. From 19fe775d38b11cf583cc5a17f71d61161c9eae70 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 20:54:43 +0000 Subject: [PATCH 22/42] Set --experts_implementation on the FSDP2 config directly; pass it only when given The FSDP2 load builds its meta model from a pre-fetched config, so the flag is set on that config's _experts_implementation, as build_meta_causal_lm sets a pre-fetched config's attn_implementation, rather than through an AutoConfig kwarg. A test captures the config the FSDP2 loader receives and fails without it. get_model passes the flag on its compressed-tensors branch only when given, as on every other load, instead of an explicit None. Signed-off-by: Shengliang Xu --- examples/hf_ptq/example_utils.py | 9 +++++--- examples/hf_ptq/hf_ptq.py | 9 +++++--- tests/examples/hf_ptq/test_hf_ptq_args.py | 25 +++++++++++++++++++++++ 3 files changed, 37 insertions(+), 6 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index fa5f91f0fad..0ef0c4fa389 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -642,8 +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") - if experts_implementation is not None: - model_kwargs["experts_implementation"] = experts_implementation + # 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. @@ -723,7 +726,7 @@ def has_pack_quantized_config(config): device_map="auto", trust_remote_code=trust_remote_code, dtype="auto", - experts_implementation=experts_implementation, + **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 07adc6eacf7..5a9f4e3d3a1 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -320,10 +320,13 @@ def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False if args.use_fsdp2: - config_kwargs = {"trust_remote_code": args.trust_remote_code} + hf_config = AutoConfig.from_pretrained( + args.local_checkpoint_path, trust_remote_code=args.trust_remote_code + ) if args.experts_implementation is not None: - config_kwargs["experts_implementation"] = args.experts_implementation - hf_config = AutoConfig.from_pretrained(args.local_checkpoint_path, **config_kwargs) + # 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. + hf_config._experts_implementation = args.experts_implementation validate_fsdp2_supported(args, hf_config) full_model = parallel_load_and_prepare_fsdp2( args.local_checkpoint_path, diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 75b611b68ff..ce584384d6d 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 Qwen3MoeConfig from modelopt.recipe import load_recipe from modelopt.recipe.config import ( @@ -414,6 +415,30 @@ def distributed_load(checkpoint, *_args, **_kwargs): hf_ptq.load_model(args) +def test_fsdp2_load_builds_with_the_requested_experts_implementation(monkeypatch, tmp_path): + """--experts_implementation reaches the config the FSDP2 load builds its meta model from.""" + Qwen3MoeConfig(architectures=["Qwen3MoeForCausalLM"]).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" + 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): From 39ffe7db6b447f7dda751ab53b1c189daba4807a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 21:18:05 +0000 Subject: [PATCH 23/42] Warn when quantization moves the fused experts to another backend normalize_experts_impl_on_the_fly rewrote the experts backend without a word. On main it always forced eager, but a caller can now ask for one (a load's experts_implementation, hf_ptq's --experts_implementation), and an unintercepted backend redirected to grouped_mm, or experts that quantize on the eager forward only, override that choice. Each change now warns once, naming both backends and why; a supported value, the common case, stays quiet. The redirect test expects the warning and fails without it. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/huggingface.py | 17 ++++++++++++++++- .../quantization/plugins/test_fused_experts.py | 10 +++++++--- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 49ac321c1fa..053d42f32a7 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1752,11 +1752,26 @@ def _configs(cfg): "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 cfg._experts_implementation not in supported_impls: + elif was not in supported_impls: cfg._experts_implementation = fallback_impl + if cfg._experts_implementation != was: + 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). + warnings.warn( + f"Quantizing the fused experts on {now!r} instead of experts_implementation={was!r}: " + f"{reason}." + ) def assert_fused_experts_intercepted(model): diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 4cbb5d078b1..3d01b756c8e 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -1355,7 +1355,10 @@ 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) - normalize_experts_impl_on_the_fly(model) + 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_falls_back_to_eager_without_the_grouped_forward(self, monkeypatch): @@ -1462,14 +1465,15 @@ class _ScaledOnGroupedMM(_ScaledOnEager): _eager_only = False @pytest.mark.parametrize("impl", ["eager", "grouped_mm", None]) - def test_leaves_a_supported_value_unchanged(self, impl): + 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.""" + supported, so they are left as they are, without a word.""" cfg = _StubConfig() cfg._experts_implementation = impl model = _TinyMoEModelWithConfig(cfg) 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)] # --------------------------------------------------------------------------- From 730f1b7ac3e29d69b441713e105afa92a8171f58 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 21:50:19 +0000 Subject: [PATCH 24/42] Do not warn about pinning a None experts backend to eager None already dispatches to the eager F.linear forward, so moving it to eager for experts that quantize on the eager forward only changes nothing; it warned anyway, for every such model, through the experts' own config too. The FP8 test now covers a grouped_mm and a None backend, and counts the warnings: one, then none. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 3 ++- .../torch/quantization/plugins/test_fused_experts.py | 11 ++++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 053d42f32a7..a044b0931b0 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1759,7 +1759,8 @@ def _configs(cfg): cfg._experts_implementation = "eager" elif was not in supported_impls: cfg._experts_implementation = fallback_impl - if cfg._experts_implementation != was: + # 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" diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 3d01b756c8e..1d62fa62b68 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -1403,11 +1403,14 @@ def test_no_crash_when_impl_attr_missing(self): normalize_experts_impl_on_the_fly(model) assert not hasattr(cfg, "_experts_implementation") - def test_forces_eager_for_experts_quantized_on_the_eager_forward_only(self): + @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.""" - model = _NonMoEModelWithConfig(_StubConfig(impl="grouped_mm")) + 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" @@ -1416,6 +1419,8 @@ def test_forces_eager_for_experts_quantized_on_the_eager_forward_only(self): ) 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 From afa4ad7ade8f4ec79d2f6ce43f95d1488d9613a2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 22:08:24 +0000 Subject: [PATCH 25/42] Say a fused-experts backend redirect on rank 0 only Every rank redirects its experts alike, so an expert-parallel run printed one copy of the notice per rank. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index a044b0931b0..0e1c0395077 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -39,6 +39,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 @@ -1768,8 +1769,9 @@ def _configs(cfg): 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). - warnings.warn( + # 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}." ) From 40c2c856b42562555a8c5e88cab1ec92ea4215fc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 22:25:14 +0000 Subject: [PATCH 26/42] Pin that --experts_implementation reaches a VLM's text config on FSDP2 Transformers' _experts_implementation setter carries the value to the sub-configs, so the one assignment on the pre-fetched config also reaches the text config Qwen3.5-MoE's experts dispatch on. The FSDP2 test now runs a VLM config too and checks its text config; it fails if the value reaches the root config only. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 3 ++- tests/examples/hf_ptq/test_hf_ptq_args.py | 18 ++++++++++++++---- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 5a9f4e3d3a1..c55a39b19c4 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -325,7 +325,8 @@ def load_model(args: argparse.Namespace): ) 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. + # 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( diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index ce584384d6d..2f0b425de40 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -26,7 +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 Qwen3MoeConfig +from transformers import Qwen3_5MoeConfig, Qwen3MoeConfig from modelopt.recipe import load_recipe from modelopt.recipe.config import ( @@ -415,9 +415,18 @@ def distributed_load(checkpoint, *_args, **_kwargs): hf_ptq.load_model(args) -def test_fsdp2_load_builds_with_the_requested_experts_implementation(monkeypatch, tmp_path): - """--experts_implementation reaches the config the FSDP2 load builds its meta model from.""" - Qwen3MoeConfig(architectures=["Qwen3MoeForCausalLM"]).save_pretrained(tmp_path) +@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, @@ -432,6 +441,7 @@ def test_fsdp2_load_builds_with_the_requested_experts_implementation(monkeypatch 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) From 5ab9038f71f295ccb6358fbbaf8d22888686e45e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 22:37:50 +0000 Subject: [PATCH 27/42] Name the ways forward in the fused-experts interception error A class no backend change redirects needs its own quantized wrapper; meanwhile the rest of the model quantizes with these modules' quantizers disabled, which the check honours. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 0e1c0395077..53fe4be4165 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1841,7 +1841,10 @@ def assert_fused_experts_intercepted(model): + "\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`." + "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." ) From 4386dab61224c3dad500a81060d64bf7beff4be7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:23:00 +0000 Subject: [PATCH 28/42] Give the vision-calibration test's arguments the new experts_implementation load_model now reads args.experts_implementation; the test's stand-in for hf_ptq's parsed arguments lacked it, so both of its cases failed with an AttributeError in the trtllm hf_ptq lane. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py | 1 + 1 file changed, 1 insertion(+) 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, From a69068ef17589cbda8e1e192c55f74c572dfba60 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 17:49:33 +0000 Subject: [PATCH 29/42] Offer --experts_implementation's two useful values as argparse choices A typo in the backend name surfaced only inside transformers' model load. Quantization redirects every backend but eager and grouped_mm to one of them, so those two are the choices. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index c55a39b19c4..47d6a8b2c19 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -1359,10 +1359,11 @@ def parse_args() -> argparse.Namespace: 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." + "supports it; eager calibrates exactly as earlier releases did. Quantization redirects " + "every other backend to one of these two." ), default=None, - type=str, + choices=["eager", "grouped_mm"], ) parser.add_argument( "--auto_quantize_checkpoint", From 8108562dae93b5233e6f8673487c618746372344 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 17:50:59 +0000 Subject: [PATCH 30/42] Pin that transformers' experts take the routing indices second The fused-experts forward reads the routing indices by position to tell the forwards that route rows to this rank's experts. A test now reads every experts class transformers decorates with use_experts_implementation and checks that its forward takes them right after the hidden states. Signed-off-by: Shengliang Xu --- .../plugins/test_fused_experts.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 1d62fa62b68..45d7dd72d40 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -15,11 +15,13 @@ """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 @@ -34,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 @@ -3012,6 +3015,30 @@ def forward_loop(m): 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() + 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 From da24de9dfee0587415523176e71a8e8fcc044c7a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 17:53:40 +0000 Subject: [PATCH 31/42] Reduce the interception check over the experts' expert-parallel group The check was rank-local, so under expert parallelism a rank whose experts saw no rows did not raise and waited in its next collective while the others raised. Each experts module's verdict is now reduced over the group its experts are sharded over, which every member already joins in the experts' forward: one all-reduce per group, issued whatever the rank found, and only when experts are sharded that way. A rank that raises for another rank's experts says so. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/huggingface.py | 79 +++++++++++++------ .../plugins/test_fused_experts.py | 24 ++++++ 2 files changed, 80 insertions(+), 23 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 53fe4be4165..891adf502f9 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 @@ -1777,6 +1778,32 @@ def _configs(cfg): ) +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``.""" + 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. @@ -1803,31 +1830,37 @@ def assert_fused_experts_intercepted(model): experts module no calibration token reached (vision experts under text-only calibration) falls back to weight-derived scales as before. - The check is rank-local: under data parallelism 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. + Under expert parallelism, 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: under data parallelism 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. """ - missed = [] - for name, module in model.named_modules(): - if not isinstance(module, _QuantFusedExperts): - continue - # 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 - ): - continue - - 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 []) - if not any(getattr(q, "is_enabled", False) for q in quantizers): - continue + 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) - impl = getattr(getattr(module, "config", None), "_experts_implementation", None) - missed.append(f" {name} ({type(module).__name__}, experts_implementation={impl!r})") + 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 diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 45d7dd72d40..fdcea2f39d2 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -70,6 +70,7 @@ _is_sparse_sequaential_moe_block, _QuantFusedExperts, _QuantNonGatedFusedExperts, + assert_fused_experts_intercepted, normalize_experts_impl_on_the_fly, register_fused_experts_on_the_fly, register_sparse_moe_on_the_fly, @@ -3353,3 +3354,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) From 6ed89f8dcaf25a4208ad26cab4df5a51122c5234 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 18:13:53 +0000 Subject: [PATCH 32/42] Say the interception check's reduction needs transformers 5.16's experts The expert-parallel group is read from the experts' weight, a DTensor sharded over the experts under transformers 5.16 and later. Earlier versions' expert parallelism keeps each rank's experts as a plain tensor, which names no group, so the check stays rank-local there, as it was. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/huggingface.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 891adf502f9..670574bcd3b 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1796,7 +1796,12 @@ def _bypassed_both_hooks(module) -> bool: def _expert_parallel_group(module): - """The process group ``module``'s experts are sharded over (expert parallelism), else ``None``.""" + """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 @@ -1830,12 +1835,13 @@ def assert_fused_experts_intercepted(model): experts module no calibration token reached (vision experts under text-only calibration) falls back to weight-derived scales as before. - Under expert parallelism, 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: under data parallelism 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. + 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)) From 9f0d37f28ebc2d114306f45e040d13c0f2b94f63 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 18:34:53 +0000 Subject: [PATCH 33/42] Say the interception check's gates must agree across ranks Under expert parallelism the check all-reduces over the experts' group, so whether calibrate() runs it (the forward loop ran, the HF plugin is loaded) must come out the same on every rank, or the ranks that run it wait for the ones that did not. A comment at the gate says so. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/model_quant.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 9392f555927..87c8677985c 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -195,7 +195,10 @@ def forward_loop(model): # 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. + # 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) From 9e02ecc14d43302fa0db86cbc43487fd193324a2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:12:21 +0000 Subject: [PATCH 34/42] Give the FSDP2 KV AutoQuantize test's arguments experts_implementation load_model's FSDP2 branch now reads args.experts_implementation; this 2-GPU test's stand-in for hf_ptq's parsed arguments lacked it, so it would fail with an AttributeError once run (the nightly runs it; the 1-GPU PR runners skip it). Signed-off-by: Shengliang Xu --- tests/gpu/torch/quantization/test_fsdp2_autoquant.py | 1 + 1 file changed, 1 insertion(+) 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, From baf18d1e16f023c55319dc42158e87e3ae25073d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 04:24:26 +0000 Subject: [PATCH 35/42] Return the experts' weight itself when its quantizers only collect statistics During calibration the weight quantizers collect statistics without quantizing, so each hands its slice back unchanged. The grouped forward still built its result slice by slice into a new weight-sized buffer: a copy of every expert's weight per projection on each calibration forward, now that calibration runs on grouped_mm by default. Such quantizers now collect per slice and the weight itself is returned, as for disabled ones. A test checks the result is the weight and each expert's statistics are collected; it fails without the change. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/huggingface.py | 14 +++++++--- .../plugins/test_fused_experts.py | 26 +++++++++++++++++++ 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 670574bcd3b..8072e8001fc 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -955,17 +955,25 @@ 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 this rank's experts 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. """ weight = self._local_experts(weight) - # 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 ): + if any(q.is_enabled for q in quantizers): + for q, w in zip(quantizers, weight, strict=True): + if q.is_enabled: + 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 diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index fdcea2f39d2..e039a8daa25 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -61,6 +61,8 @@ 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 @@ -2116,6 +2118,30 @@ def one_alive(_quantizer, _args, output): expected = torch.stack([q(weight[idx]) for idx, q in enumerate(quantizers)]) torch.testing.assert_close(out, expected, rtol=0, atol=0) + def test_calibrating_a_projection_returns_its_weight_without_a_copy(self): + """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.""" + 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 + for q in quantizers: + q.reset_amax() + enable_stats_collection(experts) + with torch.no_grad(): + out = experts._quantize_grouped_weight(weight, quantizers) + finish_stats_collection(experts) + + assert out is weight + 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("transform", ["pre_quant_scale", "rotate"]) def test_disabled_quantizers_keep_their_pre_quant_scale_and_rotation( self, monkeypatch, transform From 4d130ad7552e29b70bd0dc4dd6be615112d4f068 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 04:24:41 +0000 Subject: [PATCH 36/42] Say the experts stay on grouped_mm for every forward after quantize or restore The backend normalization runs at convert and at restore, so after mtq.quantize or mto.restore the fused experts keep grouped_mm for every forward, not only calibration: QAT/QAD training, AutoQuantize scoring and fake-quant evaluation too. The changelog entry now says so, and names experts_implementation="eager" as the way back for all of them. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7cdf4c44f68..89fb2a8b8d7 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -55,7 +55,7 @@ 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. -- Hugging Face MoE calibration now runs fused experts on the ``grouped_mm`` backend Transformers selects by default instead of forcing ``eager``, so calibrated activation scales can differ from earlier releases: some of Qwen3-30B-A3B's moved by up to about 10%. Load the model with ``experts_implementation="eager"`` (``--experts_implementation eager`` in ``hf_ptq.py``) to calibrate exactly as earlier releases did. +- ``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. From ef69ee33d1cdbc9a2931b1de7f5cd845a5a06863 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 06:04:02 +0000 Subject: [PATCH 37/42] Drop the local-expert indexing the grouped weight quantization never needs _quantize_grouped_weight is reached only from the grouped interceptor, inside the experts forward. Under expert parallelism transformers swaps each DTensor expert parameter for its local shard around that forward (the torch_a2a experts style asks for local tensors), and before 5.16 the shards were plain tensors, so the weight it receives is already this rank's experts and _local_experts returned it unchanged on every path. The docstring now says why no indexing is needed. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 8072e8001fc..ec2b4db4b37 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -954,11 +954,13 @@ 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 this rank's experts untouched when every + eager per-expert ``F.linear`` semantics. Returns the weight untouched when every quantizer would pass its slice through unchanged, so passthrough stays exact (and cheap): a disabled quantizer, 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. """ - weight = self._local_experts(weight) # 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 From f7c22a1c9b40918282528804339277c874e11cb0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 06:04:53 +0000 Subject: [PATCH 38/42] Call only calibrating quantizers on the grouped calibration pass-through The pass-through called every enabled weight quantizer on its slice, but only one that calibrates does anything with it: enable_stats_collection leaves a quantizer with a constant amax neither quantizing nor calibrating, so its calls were no-ops. The loop now calls only the calibrating ones. The calibration test gains a constant-amax case that calls no quantizer and fails without the change. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/huggingface.py | 9 +++++--- .../plugins/test_fused_experts.py | 21 ++++++++++++------- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index ec2b4db4b37..06dd770df6c 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -972,9 +972,12 @@ def _quantize_grouped_weight(self, weight: torch.Tensor, quantizers: nn.ModuleLi and not q.rotate_is_enabled for q in quantizers ): - if any(q.is_enabled for q in quantizers): - for q, w in zip(quantizers, weight, strict=True): - if q.is_enabled: + # 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 diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index e039a8daa25..59951fae8b5 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -2118,10 +2118,12 @@ def one_alive(_quantizer, _args, output): expected = torch.stack([q(weight[idx]) for idx, q in enumerate(quantizers)]) torch.testing.assert_close(out, expected, rtol=0, atol=0) - def test_calibrating_a_projection_returns_its_weight_without_a_copy(self): + @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.""" + 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) @@ -2129,18 +2131,23 @@ def test_calibrating_a_projection_returns_its_weight_without_a_copy(self): 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) - finish_stats_collection(experts) assert out is weight - for idx, q in enumerate(quantizers): - torch.testing.assert_close( - q.amax.flatten(), weight[idx].detach().abs().amax(dim=1), rtol=0, atol=0 - ) + 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("transform", ["pre_quant_scale", "rotate"]) def test_disabled_quantizers_keep_their_pre_quant_scale_and_rotation( From 5b82e723ab71c22e66a088724dc4d43255e0cf58 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 06:05:48 +0000 Subject: [PATCH 39/42] Pin that a grouped experts forward counts whatever its routing says The interception check counts an eager forward only when a row routes to this module's experts, but a grouped one (grouped_mm, torch_a2a) whatever its routing indices hold: those forwards call _grouped_linear whatever rows they are handed, and torch_a2a's indices are global expert ids. A test runs a grouped-configured forward that reaches neither hook with every index past the module's experts, and checks calibration still raises; without the grouped term the forward goes uncounted. Signed-off-by: Shengliang Xu --- .../plugins/test_fused_experts.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 59951fae8b5..b8ac1ac8367 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -3031,6 +3031,25 @@ def forward_loop(m): 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.""" From 5ef19b97fda517fe01017dce9823f2508d91aede Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 06:06:29 +0000 Subject: [PATCH 40/42] Pin that a capture after a grouped forward does not use its offsets The experts forward clears the grouped offsets when it returns, so a weight-input capture reached outside a forward does not split its input by a finished forward's offsets. The capture test now calls an input quantizer directly after a grouped forward and checks that nothing is captured; without the clear, it splits by the stale offsets and raises. Signed-off-by: Shengliang Xu --- .../unit/torch/quantization/plugins/test_fused_experts.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index b8ac1ac8367..177e3e217d3 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -2318,7 +2318,9 @@ class _EagerOnlyExperts(_QuantFusedExperts): ) 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.""" + 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 = {} @@ -2331,6 +2333,10 @@ def capture(quantizer, _weight, inputs): 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() From f1e5e1b6194ed082229b41dca2a834d974331ebb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 18:04:15 +0000 Subject: [PATCH 41/42] Read transformers' modeling sources as UTF-8 in the routing-order test The test read each transformers modeling file with the platform's default encoding, which is cp1252 on Windows; 16 of them, such as modeling_altclip.py, hold characters cp1252 cannot decode, so the test failed in the windows unit job. It now reads them as UTF-8. Signed-off-by: Shengliang Xu --- tests/unit/torch/quantization/plugins/test_fused_experts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 177e3e217d3..0d2b44d7e96 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -3081,7 +3081,7 @@ def test_transformers_experts_take_the_routing_indices_second(self): 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() + source = path.read_text(encoding="utf-8") if "use_experts_implementation" not in source: continue for node in ast.parse(source).body: From a6758d77eb36a801d11b794a6c79efef637131fa Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:53:58 +0000 Subject: [PATCH 42/42] Calibrate grouped inputs on the rows the grouped GEMM computes Under transformers' own expert parallelism the grouped forward sorts the rows routed to other ranks' experts past offs[-1], where the grouped GEMM neither reads nor writes. Before transformers 5.17 nothing zeroes the up projection's output there, so the down projection's input holds those rows uninitialized, and the shared input quantizer calibrated on them: an input scale from garbage, silently when the garbage is finite. While an input quantizer collects statistics, it now sees only the rows before offs[-1]; the tail passes to the GEMM unread. The host sync this needs is paid during calibration only. A test hands the grouped projections a NaN or a large finite tail, as pre-5.17 expert parallelism leaves it, and checks the input amaxes match a run without one; both cases fail without the change. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/huggingface.py | 26 ++++++++-- .../plugins/test_fused_experts.py | 50 +++++++++++++++++++ 2 files changed, 72 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 06dd770df6c..dd086013118 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1068,10 +1068,10 @@ def _quantized_grouped_linear(input, weight, offs, bias=None, is_transposed=Fals 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( @@ -1082,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. diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 0d2b44d7e96..d7f03854837 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -399,6 +399,29 @@ def forward(self, hidden_states, top_k_index, top_k_weights): 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__() @@ -2149,6 +2172,33 @@ def test_calibrating_a_projection_returns_its_weight_without_a_copy(self, consta 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