[OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager

CHANGELOG.rst

@@ -55,6 +55,8 @@ Changelog
 - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm.
 - The Windows FVD example rejects I3D checkpoints requiring unrestricted pickle loading. Custom TorchScript archives now require ``--weights-sha256`` with a trusted checksum; tensor state dictionaries continue to load directly.
 - A ``modelopt_state`` saved by an earlier release for a Hugging Face model whose fused MoE experts are stored transposed without a quantized wrapper of their own (for example Aria's) no longer restores, since those experts are now left unquantized. The checkpoints such a state exported had wrong expert weights, so there is nothing to migrate.
+- ``mtq.quantize`` and ``mto.restore`` no longer force a Hugging Face model's fused MoE experts onto the ``eager`` backend: calibration and every later forward (QAT/QAD training, AutoQuantize scoring, fake-quant evaluation) run on the ``grouped_mm`` backend Transformers selects by default, so calibrated activation scales (some of Qwen3-30B-A3B's moved by up to about 10%), step time and memory can differ from earlier releases. Load the model with ``experts_implementation="eager"`` (``--experts_implementation eager`` in ``hf_ptq.py``) to run as earlier releases did.
+- ``mtq.calibrate`` now raises an error, naming the modules, when a Hugging Face fused-experts forward reaches neither quantizer hook (a hand-written ``bmm`` forward that transformers cannot redirect, say), instead of exporting weight-derived input scales with a warning.
 
 - Layerwise-calibration resume (``checkpoint_dir``) now loads checkpoint files with ``weights_only=True``. A checkpoint directory written by an earlier release whose layer inputs hold objects other than tensors, plain containers or Hugging Face dynamic caches can no longer be resumed; recalibrate in a fresh directory, or disable checkpointing.
 

examples/hf_ptq/example_utils.py

@@ -600,6 +600,7 @@ def get_model(
     trust_remote_code=False,
     use_seq_device_map=False,
     attn_implementation=None,
+    experts_implementation=None,
     offload_folder=None,
     max_cpu_memory_gb=None,
     max_gpu_memory_gb=None,
@@ -641,6 +642,11 @@ def get_model(
     # Note: Forcibly converting the model precision between bf16 and fp16 may introduce accuracy drop
     model_kwargs = config_kwargs.copy()
     model_kwargs.setdefault("dtype", "auto")
+    # Only when asked, on every load below: an explicit None is a value transformers stores.
+    experts_kwargs = (
+        {} if experts_implementation is None else {"experts_implementation": experts_implementation}
+    )
+    model_kwargs.update(experts_kwargs)
 
     # DiffusionGemma ties encoder/decoder weights. device_map "auto" (balanced) can split
     # a tied pair across GPUs, leaving one side on the meta device and breaking generation.
@@ -720,6 +726,7 @@ def has_pack_quantized_config(config):
                 device_map="auto",
                 trust_remote_code=trust_remote_code,
                 dtype="auto",
+                **experts_kwargs,
             )
     elif get_original_hf_quant_method(hf_config) == "mxfp4":
         # Native MXFP4 checkpoints (e.g. openai/gpt-oss-*) must be dequantized to

examples/hf_ptq/hf_ptq.py

@@ -323,6 +323,11 @@ def load_model(args: argparse.Namespace):
         hf_config = AutoConfig.from_pretrained(
             args.local_checkpoint_path, trust_remote_code=args.trust_remote_code
         )
+        if args.experts_implementation is not None:
+            # Set on the config, as build_meta_causal_lm sets a pre-fetched config's
+            # attn_implementation: the meta model is built from it, with no kwarg of its own. The
+            # setter carries it to the sub-configs, such as the text config a VLM's experts read.
+            hf_config._experts_implementation = args.experts_implementation
         validate_fsdp2_supported(args, hf_config)
         full_model = parallel_load_and_prepare_fsdp2(
             args.local_checkpoint_path,
@@ -342,6 +347,7 @@ def load_model(args: argparse.Namespace):
             trust_remote_code=args.trust_remote_code,
             use_seq_device_map=args.use_seq_device_map,
             attn_implementation=args.attn_implementation,
+            experts_implementation=args.experts_implementation,
             offload_folder=args.offload_folder,
             max_cpu_memory_gb=args.max_cpu_memory_gb,
             max_gpu_memory_gb=args.max_gpu_memory_gb,
@@ -367,6 +373,8 @@ def load_model(args: argparse.Namespace):
             model_kwargs = {"trust_remote_code": args.trust_remote_code}
             if args.attn_implementation is not None:
                 model_kwargs["attn_implementation"] = args.attn_implementation
+            if args.experts_implementation is not None:
+                model_kwargs["experts_implementation"] = args.experts_implementation
             full_model = load_pretrained(
                 args.local_checkpoint_path,
                 AutoModelForCausalLM,
@@ -1346,6 +1354,17 @@ def parse_args() -> argparse.Namespace:
         default=None,
         type=str,
     )
+    parser.add_argument(
+        "--experts_implementation",
+        help=(
+            "Fused MoE experts backend to load the model with, passed to the HF model loading as "
+            "experts_implementation. By default transformers picks grouped_mm where the model "
+            "supports it; eager calibrates exactly as earlier releases did. Quantization redirects "
+            "every other backend to one of these two."
+        ),
+        default=None,
+        choices=["eager", "grouped_mm"],
+    )
     parser.add_argument(
         "--auto_quantize_checkpoint",
         type=str,

modelopt/torch/quantization/model_calib.py

@@ -1075,9 +1075,8 @@ def _fused_expert_projections(module: nn.Module):
 
     ``weight`` is the 3-D ``(num_experts, out, in)`` parameter, ``weight_quantizers`` holds one
     quantizer per expert, and ``input_quantizer`` is shared by the experts. Under eager it runs once
-    per expert forward with ``module._current_expert_idx`` set to that expert; under grouped_mm once
-    per projection, on every expert's rows, with the index cleared (``None``), so a capture keyed on
-    it skips that batch.
+    per expert forward with ``module._current_expert_idx`` set to that expert; under grouped_mm
+    and torch_a2a once per projection, on a batch sorted by expert (``module._grouped_offs``).
     """
     first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
     for weight_name, quantizers_name, input_q_name in (
@@ -1131,6 +1130,31 @@ def _refuse_packed_fused_expert_weights(name, module, method, updates):
             )
 
 
+def _expert_slices(offs: torch.Tensor, num_experts: int, rows: int):
+    """``(expert, start, end)`` for the rows of each expert in a batch sorted by expert.
+
+    ``offs`` is the cumulative token count per expert (``_grouped_offs``): exclusive ends, one per
+    expert, as transformers' grouped_mm forward and ``a2a_moe`` build them. An expert without rows is
+    skipped, and rows beyond ``offs[-1]`` (expert-parallel sentinels) belong to no expert. Offsets of
+    another form raise (the inclusive ``[0, c0, c0 + c1, ...]`` form, which carries one end too many,
+    or ends out of order), since they would credit experts with other experts' rows.
+    """
+    ends = offs.tolist()
+    if (
+        len(ends) != num_experts
+        or (ends and ends[-1] > rows)
+        or any(end < prev for prev, end in zip([0, *ends], ends))
+    ):
+        raise RuntimeError(
+            f"Grouped experts offsets {ends} do not split {rows} rows over {num_experts} experts."
+        )
+    start = 0
+    for idx, end in enumerate(ends):
+        if end > start:
+            yield idx, start, end
+        start = end
+
+
 class _WeightInputHook:
     """Identify active input collectors so auxiliary modules can request a calibration forward."""
 
@@ -1146,23 +1170,41 @@ def _register_weight_input_hooks(model, names, capture, on_weight=None):
 
     ``capture(weight_quantizer, weight, inputs)`` runs on every forward of a weight whose quantizer
     is enabled: dense quantized linears hook the layer input; HF fused-MoE experts hook the shared
-    input quantizers, keyed by the active expert (``_current_expert_idx``). Weights without a hook
-    (conv, SequentialQuantizer, non-eager experts) are never captured. ``on_weight(name, weight,
-    weight_quantizer)``, if given, is called once per hooked weight. Returns removable handles.
+    input quantizers, keyed by the active expert (``_current_expert_idx``) when eager, or split per
+    expert by ``_grouped_offs`` under ``grouped_mm``. Weights without a hook (conv,
+    SequentialQuantizer) are never captured. ``on_weight(name, weight, weight_quantizer)``, if
+    given, is called once per hooked weight. Returns removable handles.
     """
     handles: list = []
 
     def _make_expert_hook(expert_module, weight_name, quantizers, enabled):
+        """Per-expert capture for whichever forward ran.
+
+        The eager forward feeds the input quantizer one call per expert (``_current_expert_idx``);
+        the grouped ones (``grouped_mm``, ``torch_a2a``) feed it one expert-sorted batch per
+        projection, which ``_grouped_offs`` (set by the grouped interceptor, cleared by the eager one)
+        splits back per expert. Rows beyond ``offs[-1]`` (EP sentinels) belong to no expert.
+        """
+
         def _expert_hook(_input_quantizer, args):
             if not args:
                 return
-            idx = expert_module._current_expert_idx
-            if idx in enabled:
-                # Read the weight fresh (valid under accelerate/FSDP re-materialization), indexed by
-                # this rank's own experts: an expert-parallel weight is a DTensor whose index names a
-                # global expert (see ``_local_experts``).
-                weight = expert_module._local_experts(getattr(expert_module, weight_name))
-                capture(quantizers[idx], weight[idx], args[0])
+            # Read the weight fresh (valid under accelerate/FSDP re-materialization), indexed by this
+            # rank's own experts: an expert-parallel weight is a DTensor whose index names a global
+            # expert (see ``_local_experts``).
+            offs = expert_module._grouped_offs
+            if offs is None:
+                idx = expert_module._current_expert_idx
+                if idx in enabled:
+                    weight = expert_module._local_experts(getattr(expert_module, weight_name))
+                    capture(quantizers[idx], weight[idx], args[0])
+                return
+            weight = expert_module._local_experts(getattr(expert_module, weight_name))
+            for idx, start, end in _expert_slices(
+                offs, expert_module.num_experts, args[0].shape[0]
+            ):
+                if idx in enabled:
+                    capture(quantizers[idx], weight[idx], args[0][start:end])
 
         return _expert_hook
 
@@ -1198,6 +1240,7 @@ def _dense_hook(linear, args):
                     enabled = {i for i, q in enumerate(quantizers) if q.is_enabled}
                     if not enabled:
                         continue
+                    module._weight_input_capture = True
                     handles.append(
                         input_quantizer.register_forward_pre_hook(
                             _WeightInputHook(
@@ -2588,6 +2631,16 @@ def _make_gptq_handle(name, m, weight_quantizer=None, weight=None, log_mse=True)
 
     def _make_expert_hook(experts, per_expert):
         def _expert_hook(_input_quantizer, _args, output):
+            if experts._grouped_offs is not None:
+                # grouped_mm and torch_a2a quantize one expert-sorted batch per projection: split
+                # it back per expert by the cumulative token counts, as the eager loop feeds it.
+                for idx, start, end in _expert_slices(
+                    experts._grouped_offs, experts.num_experts, output.shape[0]
+                ):
+                    handle = per_expert.get(idx)
+                    if handle is not None:
+                        handle.accumulate(output[start:end])
+                return
             handle = per_expert.get(experts._current_expert_idx)
             if handle is not None:
                 handle.accumulate(output)
@@ -2611,6 +2664,7 @@ def _expert_hook(_input_quantizer, _args, output):
                     if q.is_enabled
                 }
                 if per_expert:
+                    module._weight_input_capture = True
                     expert_handles.update((id(quantizers[i]), h) for i, h in per_expert.items())
                     hooks.append(
                         input_quantizer.register_forward_hook(_make_expert_hook(module, per_expert))

modelopt/torch/quantization/model_quant.py

@@ -18,6 +18,7 @@
 import fnmatch
 import inspect
 import os
+import sys
 import warnings
 from collections.abc import Callable, Iterable, Sequence
 from contextlib import contextmanager
@@ -143,6 +144,25 @@ def calibrate(
             def forward_loop(model):
                 return original_forward_loop()  # type: ignore[call-arg]
 
+    # Whether the algorithm actually invoked the loop. `max_calibrate` may skip it
+    # (`skip_forward_without_activation_calib`) when no enabled quantizer needs activation
+    # statistics, and a post-calibration check that assumed it ran would fire on every
+    # weight-only calibration.
+    forward_loop_ran = False
+    if forward_loop is not None:
+        _requested_forward_loop = forward_loop
+
+        def forward_loop(model):
+            nonlocal forward_loop_ran
+            forward_loop_ran = True
+            return _requested_forward_loop(model)
+
+        # The fused-experts interception check below judges this calibration's forwards alone.
+        for module in model.modules():
+            if hasattr(module, "_intercept_count"):
+                module._intercept_count = module._forward_count = 0
+                module._weight_input_capture = False
+
     # move the model to eval mode, remembering each module's mode; quantizers, and plain containers of them,
     # follow the module they quantize instead
     training_modes = {
@@ -169,6 +189,20 @@ def forward_loop(model):
 
     _restore_training_modes(model, training_modes)
 
+    if forward_loop_ran:
+        # A fused-experts module whose forward ran past both interception hooks collected
+        # nothing. Surface it here rather than let the export emit weight-derived fallback scales
+        # under a warning that blames the size of the calibration set. Raised once the model is
+        # back in its mode, for a caller that catches it to calibrate again. Only a loaded HF plugin
+        # installs fused-experts wrappers; importing it here would rerun one that failed to load,
+        # which import_plugin only warns about, and fail every calibration with its error. Both
+        # conditions must come out the same on every rank: under expert parallelism the check
+        # all-reduces over the experts' group, and a rank that skipped it would leave its peers
+        # waiting there.
+        huggingface = sys.modules.get("modelopt.torch.quantization.plugins.huggingface")
+        if huggingface is not None:
+            huggingface.assert_fused_experts_intercepted(model)
+
     return model