[OMNIML-4998] Per-expert weight quantization on TEGroupedLinear

modelopt/torch/quantization/model_calib.py

@@ -357,6 +357,21 @@ def sync_quantizer_amax_across_tp(
             quantizer.sync_amax_across_distributed_group(parallel_state.tensor_parallel_group)
 
     # Step 2: Sync amax across relevant parallelism (such as TP / EP)
+    def _weight_axes_for_sync(module, base_axes):
+        # For column-parallel TEGroupedLinear in per-expert mode (axis=0), the
+        # expert weights are NOT TP-sharded (etp=1 case) — axis=0 indexes experts,
+        # not output channels. Per-rank reductions still produce slightly
+        # divergent amax across TP (BF16/FP16 sums are not bit-identical across
+        # ranks even with the same input), and dist-checkpoint save treats
+        # _amax as replicated and captures only one rank's view. Without this
+        # sync, model_ref's per-rank-different amax mismatches the loaded
+        # model_test on every TP rank except the one whose value was saved.
+        if hasattr(module, "weight0"):
+            quantizer = getattr(module, "weight_quantizer", None)
+            if isinstance(quantizer, TensorQuantizer) and quantizer.axis in (0, (0,)):
+                return list(base_axes) + [0]
+        return base_axes
+
     for name, module in model.named_modules():
         if getattr(module, "_parallel_state", None) is None:
             continue
@@ -373,7 +388,7 @@ def sync_quantizer_amax_across_tp(
                 module.weight_quantizer,
                 name,
                 "weight_quantizer",
-                axes_for_sync=[None, -1],
+                axes_for_sync=_weight_axes_for_sync(module, [None, -1]),
                 parallel_state=module.parallel_state,
             )
 

modelopt/torch/quantization/plugins/megatron.py

@@ -680,6 +680,64 @@ class _QuantTELayerNormColumnParallelLinear(
 
     # Quantized subclasses to support TEGroupedMLP quantization
     class _QuantMegatronTEGroupedLinear(_QuantTEGroupedLinear, _MegatronParallelLinear):
+        def _ep_group(self):
+            # Return the expert_model_parallel_group iff it's initialized AND has >1 rank.
+            # _MegatronTEGroupedMLP._setup populates parallel_state with the EP group;
+            # outside that wrapping it may be unset (e.g. ad-hoc unit tests).
+            ps = getattr(self, "parallel_state", None)
+            if ps is None:
+                return None
+            ep = ps.expert_model_parallel_group
+            if not ep.is_initialized() or ep.world_size() <= 1:
+                return None
+            return ep
+
+        def _gather_global_per_expert_amax(self):
+            # Return the global [num_gemms_global] per-expert amax, gathered across
+            # the EP group, without mutating the live weight_quantizer._amax buffer
+            # (the forward path indexes _amax by local expert position and would
+            # break if we replaced it with the global shape).
+            #
+            # Returns None when the layer isn't per-expert (per-tensor stays the
+            # legacy scalar path) or when EP == 1 (local == global, caller can
+            # reshape the live buffer directly).
+            if not self._is_per_expert_weight_quant():
+                return None
+            quantizer = self.weight_quantizer
+            amax = getattr(quantizer, "_amax", None)
+            if amax is None:
+                return None
+            ep = self._ep_group()
+            if ep is None:
+                return amax.view(self.num_gemms)  # EP=1: local IS global.
+
+            global_size = self.num_gemms * ep.world_size()
+            if amax.numel() != self.num_gemms:
+                raise AssertionError(
+                    f"TEGroupedLinear weight_quantizer._amax numel {amax.numel()}; "
+                    f"expected {self.num_gemms} (local per-expert) or "
+                    f"{global_size} (global per-expert)."
+                )
+            v_local = amax.view(self.num_gemms)
+            gathered = [torch.empty_like(v_local) for _ in range(ep.world_size())]
+            torch.distributed.all_gather(gathered, v_local, group=ep.group)
+            return torch.cat(gathered, dim=0)
+
+        def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
+            # Gather the global per-expert amax ONCE before Megatron's save traversal.
+            # Stash on a temporary attribute that _process_quantizer_amax consumes; the
+            # live weight_quantizer._amax buffer stays local so forward keeps working.
+            #
+            # Done at the top of sharded_state_dict (not inside _process_quantizer_amax)
+            # so the EP collective completes BEFORE Megatron's dist-checkpoint save
+            # kicks off its own default-PG ALLGATHER metadata exchanges. Interleaving
+            # EP gathers with default-PG collectives deadlocks NCCL.
+            self._cached_global_per_expert_amax = self._gather_global_per_expert_amax()
+            try:
+                return super().sharded_state_dict(prefix, sharded_offsets, metadata)
+            finally:
+                self._cached_global_per_expert_amax = None
+
         def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
             # _sharded_state_dict_grouped adds _extra_state{gemm_idx} for gemm_idx:[1, num_gemms] in
             # sharded_state_dict which is same as _extra_state. The _extra_state{gemm_idx} is used for
@@ -690,11 +748,52 @@ def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
                 for k, v in state_dict.items()
                 if not any(k.endswith(f"_extra_state{num}") for num in range(1, self.num_gemms))
             }
+            # Per-expert amax was saved as the gathered global [num_gemms_global]
+            # tensor. On restore each EP rank narrows it back to its local
+            # [num_gemms] slice before the base class's view_as() reshape.
+            ep = self._ep_group()
+            if ep is not None:
+                global_size = self.num_gemms * ep.world_size()
+                offset = ep.rank() * self.num_gemms
+                for k in list(filtered_state_dict.keys()):
+                    if (
+                        "weight_quantizer" in k
+                        and k.endswith("_amax")
+                        and filtered_state_dict[k].numel() == global_size
+                    ):
+                        filtered_state_dict[k] = (
+                            filtered_state_dict[k]
+                            .view(global_size)
+                            .narrow(0, offset, self.num_gemms)
+                        )
             return super()._load_from_state_dict(filtered_state_dict, prefix, *args, **kwargs)
 
+        def _get_shard_axis_dict(self, state_dict):
+            # The column-parallel parent marks weight_quantizer._amax as TP-sharded
+            # along axis 0 whenever weight_quantizer.axis is not None — correct for
+            # per-channel-along-output, but wrong for our per-expert path: the
+            # gathered amax is [num_gemms_global], not output-channel-sharded.
+            # Strip the marker so the dist-checkpoint framework treats it as
+            # replicated across TP. EP-aware handling lives in
+            # _gather_global_per_expert_amax (called from sharded_state_dict on
+            # save) and _load_from_state_dict (narrow on load).
+            shard_axis_dict = super()._get_shard_axis_dict(state_dict)
+            if self._is_per_expert_weight_quant():
+                for k in list(shard_axis_dict):
+                    if "weight_quantizer" in k and k.endswith("_amax"):
+                        del shard_axis_dict[k]
+            return shard_axis_dict
+
         def _process_quantizer_amax(self, k, v, quantizer_state_dict):
-            assert v.numel() == 1, "TEGroupedLinear only supports per-tensor quantization"
-            quantizer_state_dict[k] = v.view(-1)
+            # Per-expert weight amax: emit the gathered global tensor cached by
+            # sharded_state_dict's pre-pass. No collective fires here, so this
+            # call is safe to run inside Megatron's save traversal.
+            # Per-tensor (or non-weight) amax: just reshape the local value.
+            cached = getattr(self, "_cached_global_per_expert_amax", None)
+            if cached is not None and "weight_quantizer" in k and k.endswith("_amax"):
+                quantizer_state_dict[k] = cached.view(cached.numel())
+                return
+            quantizer_state_dict[k] = v.view(v.numel())
 
     @QuantModuleRegistry.register(
         {TEColumnParallelGroupedLinear: "megatron_TEColumnParallelGroupedLinear"}

modelopt/torch/quantization/plugins/transformer_engine.py

@@ -27,7 +27,7 @@
 
 from modelopt.torch.quantization.utils import replace_function
 
-from ..nn import QuantModuleRegistry
+from ..nn import QuantModuleRegistry, TensorQuantizer
 from .custom import _ParallelLinear
 
 _TE_VERSION = Version(te.__version__)
@@ -103,6 +103,43 @@ def te_quantized_linear_fn(package, func_name, self, *args, **kwargs):
     _quantized_linear_fn = te_quantized_linear_fn
 
 
+def _reshape_loaded_amax_to_buffer_shape(
+    module,
+    state_dict,
+    prefix,
+    local_metadata,
+    strict,
+    missing_keys,
+    unexpected_keys,
+    error_msgs,
+):
+    """Reshape a loaded `_amax` tensor to match the live buffer shape on this rank.
+
+    PyTorch's `load_state_dict` enforces an exact-shape check at the leaf module
+    level. For per-expert TEGroupedLinear quantization, the dist-checkpoint tensor
+    is saved flat as `[num_gemms]` (the legacy `_process_quantizer_amax` flatten)
+    while the live buffer is registered as `[num_gemms, 1, 1]` from the metadata
+    path (axis=0 calibration with keepdims). Reshape the loaded tensor in the
+    state_dict the strict check is about to read.
+
+    The hook is leaf-level and rank-local — it only touches the local
+    `state_dict[prefix + "_amax"]` and the local buffer's shape, and only acts
+    when shapes differ but numel matches. EP>1 narrowing happens in the outer
+    `_QuantMegatronTEGroupedLinear._load_from_state_dict` before this hook runs.
+    """
+    key = prefix + "_amax"
+    if key not in state_dict:
+        return
+    buf = getattr(module, "_amax", None)
+    if buf is None:
+        return
+    loaded = state_dict[key]
+    if loaded.shape == buf.shape:
+        return
+    if loaded.numel() == buf.numel():
+        state_dict[key] = loaded.reshape(buf.shape).contiguous()
+
+
 # Register the public te.pytorch.GroupedLinear class
 @QuantModuleRegistry.register({te_grouped_linear.GroupedLinear: "te_GroupedLinear"})
 class _QuantTEGroupedLinear(_ParallelLinear):
@@ -137,21 +174,72 @@ def _setup(self):
         # Remove self.weight after setup.
         delattr(self, "weight")
 
-        # TODO: GroupedLinear supports weights split by `num_gemms`, to support quantization
-        # with static parameters beyond per-tensor, we need to support a unique quantizer for each gemm.
+        # Reconcile the per-expert _amax storage shape on restore. The save path keeps
+        # the dist-checkpoint tensor as flat [num_gemms_global] (legacy convention for
+        # quantizer amax) while _pytorch_state_metadata records the live multi-dim
+        # buffer shape [num_gemms, 1, 1] (axis=0 calibration + keepdims). On restore,
+        # the metadata path registers an [num_gemms, 1, 1] buffer first; load_state_dict
+        # then tries to copy the flat tensor into it and raises a size mismatch.
+        #
+        # Reshape inside a pre-hook at the leaf quantizer level so it fires AFTER
+        # PyTorch has narrowed state_dict to this submodule's prefix but BEFORE the
+        # strict size check. The hook is pure-CPU and rank-local — no per-rank state
+        # queries beyond the leaf's own buffer shape, so it doesn't introduce the
+        # rank-asymmetric collective ordering that the parent-level state_dict
+        # mutation approach did. EP>1 still relies on the outer narrow path.
+        if hasattr(self, "weight_quantizer") and isinstance(self.weight_quantizer, TensorQuantizer):
+            self.weight_quantizer._register_load_state_dict_pre_hook(
+                _reshape_loaded_amax_to_buffer_shape, with_module=True
+            )
+
+    def _is_per_expert_weight_quant(self) -> bool:
+        # Per-expert mode: axis=0 on the weight quantizer means the stacked
+        # [num_gemms, out, in] weight is reduced over (1, 2), producing one amax
+        # per expert. Anything else (axis=None, axis=(0, 1), ...) keeps the
+        # legacy per-tensor path.
+        quantizer = getattr(self, "weight_quantizer", None)
+        if not isinstance(quantizer, TensorQuantizer):
+            return False
+        axis = quantizer.axis
+        return axis in (0, (0,))
 
     def modelopt_post_restore(self, prefix: str = ""):
         # GroupedMLP stores the weights as weight0, weight1, etc. To run post_restore in order to
-        # initialize the quantizer states, self.weight is used to extract shape, dtype etc. Assigning
-        # self.weight0 to self.weight to run the quantizer states initialization.
+        # initialize the quantizer states, self.weight is used to extract shape, dtype etc.
+        #
+        # Per-expert mode (axis=0) stores _amax as [num_gemms, 1, 1] — the same shape
+        # iter_weights_for_calibration and te_grouped_quantized_linear_fn produce by stacking
+        # the per-expert weights first. Mirror that here so the restore path initializes the
+        # placeholder _amax to the same shape; otherwise dist-checkpoint load builds _amax from
+        # the un-stacked [out, in] view (shape [out, 1]) and _gather_global_per_expert_amax
+        # later fails on `amax.view(num_gemms)`.
         assert not hasattr(self, "weight"), "self.weight should not exist for TEGroupedLinear"
-        self.weight = self.weight0
+        if self._is_per_expert_weight_quant():
+            weights = [getattr(self, f"weight{i}") for i in range(self.num_gemms)]
+            self.weight = torch.stack(weights, dim=0)
+        else:
+            self.weight = self.weight0
         super().modelopt_post_restore(prefix=prefix)
         # Remove self.weight after post_restore.
         delattr(self, "weight")
 
     def iter_weights_for_calibration(self):
-        """Yield ``(weight_i, weight_quantizer)`` for each of the ``num_gemms`` grouped weights."""
+        """Yield grouped weights for calibration.
+
+        Per-tensor (``axis=None``): yields ``(weight_i, weight_quantizer)`` for each
+        expert separately; the calibrator accumulates a single amax across all experts.
+
+        Per-expert (``axis=0``): yields ``(stacked, weight_quantizer)`` once, where
+        ``stacked`` has shape ``[num_gemms, out, in]``. The quantizer's axis-0 reduction
+        then produces one amax per expert (shape ``[num_gemms, 1, 1]``).
+        """
+        if self._is_per_expert_weight_quant():
+            weights = [getattr(self, f"weight{i}", None) for i in range(self.num_gemms)]
+            if any(w is None for w in weights):
+                return
+            yield torch.stack(weights, dim=0), self.weight_quantizer
+            return
+
         for i in range(self.num_gemms):
             weight_i = getattr(self, f"weight{i}", None)
             if weight_i is not None:
@@ -182,8 +270,17 @@ def te_grouped_quantized_linear_fn(package, func_name, self, *args):
 
         new_args = list(args)
         new_args[inp_pos] = self.input_quantizer(args[inp_pos])
-        for i in range(weights_start, weights_start + num_gemms):
-            new_args[i] = self.weight_quantizer(args[i])
+        if self._is_per_expert_weight_quant():
+            # Stack the N expert weights into one [N, out, in] tensor so the
+            # quantizer's axis-0 reduction sees them together; the resulting
+            # _amax of shape [N, 1, 1] broadcasts to give per-expert fake-quant.
+            stacked = torch.stack(list(args[weights_start : weights_start + num_gemms]), dim=0)
+            q_stacked = self.weight_quantizer(stacked)
+            for i in range(num_gemms):
+                new_args[weights_start + i] = q_stacked[i]
+        else:
+            for i in range(weights_start, weights_start + num_gemms):
+                new_args[i] = self.weight_quantizer(args[i])
         output = getattr(package, func_name)(*new_args)
         # TE 2.15+ returns `(out, new_workspaces)`; TE <= 2.14 returns just `out`.
         # Only the activation tensor participates in output quantization.

tests/_test_utils/torch/megatron/models.py

@@ -140,6 +140,7 @@ def get_mcore_gpt_model(
     use_cpu_initialization: bool = False,
     bf16: bool = True,
     use_te: bool = False,
+    sequence_parallel: bool = False,
     # MoE-specific parameters
     moe_grouped_gemm: bool = False,
     moe_ffn_hidden_size: int | None = None,
@@ -168,7 +169,7 @@ def squared_relu(x):
         pipeline_model_parallel_size=pipeline_model_parallel_size,
         expert_model_parallel_size=expert_model_parallel_size,
         expert_tensor_parallel_size=expert_tensor_parallel_size,
-        sequence_parallel=False,
+        sequence_parallel=sequence_parallel,
         num_layers=num_layers,
         num_layers_in_first_pipeline_stage=num_layers_in_first_pipeline_stage,
         num_layers_in_last_pipeline_stage=num_layers_in_last_pipeline_stage,

tests/gpu_megatron/torch/quantization/plugins/test_megatron.py

@@ -289,6 +289,7 @@ def _gpt_model_provider(
                 use_cpu_initialization=meta_device,
                 num_moe_experts=num_moe_experts,
                 moe_grouped_gemm=moe_grouped_gemm,
+                sequence_parallel=(tp_size > 1),  # Required for MoE + TP (mirrors hybrid path)
             )
 
     if not meta_device:
@@ -693,6 +694,55 @@ def test_te_grouped_vs_sequential_quantize(dist_workers_size_4, quant_cfg):
     )
 
 
+# OMNIML-4998: per-expert weight amax (axis=0) on TEGroupedLinear should round-trip
+# through sharded_state_dict / dist-checkpoint with bit-for-bit equality across EP.
+TE_GROUPED_PER_EXPERT_CFG = {
+    "algorithm": "max",
+    "quant_cfg": [
+        # Disable everything (attention, embeddings, non-MoE MLP) so the test
+        # isolates the per-expert path on TEGroupedMLP experts.
+        {"quantizer_name": "*", "enable": False},
+        # Re-enable axis=0 on TEGrouped MoE experts only -- this triggers the
+        # per-expert path inside _QuantTEGroupedLinear.
+        {
+            "quantizer_name": "*experts.linear_fc*.weight_quantizer",
+            "cfg": {"num_bits": 8, "axis": 0, "enable": True},
+        },
+    ],
+}
+
+
+def test_te_grouped_per_expert_sharded_state_dict(dist_workers, need_4_gpus, tmp_path):
+    """Per-expert (axis=0) weight amax round-trips through dist-checkpoint on TEGroupedMLP.
+
+    Mirrors ``test_moe_sharded_state_dict``'s setup exactly (dist_workers fixture,
+    hidden_size=256, tp=2 ep=2 etp=1 num_moe_experts=4 moe_grouped_gemm=True), with
+    only the quant_cfg varying. The existing test is known to pass at this layout
+    for FP8_DEFAULT_CFG / NVFP4_DEFAULT_CFG; this test exercises the OMNIML-4998
+    per-expert (axis=0) path on top of the same infrastructure.
+    """
+    moe_config = {
+        "tp_size": 2,
+        "ep_size": 2,
+        "etp_size": 1,
+        "num_moe_experts": 4,
+        "moe_grouped_gemm": True,
+        "transformer_impl": "transformer_engine",
+    }
+    dist_workers.run(
+        partial(
+            _test_sharded_state_dict,
+            tmp_path,
+            copy.deepcopy(TE_GROUPED_PER_EXPERT_CFG),
+            256,
+            None,
+            False,
+            False,
+            moe_config,
+        ),
+    )
+
+
 @pytest.mark.parametrize("ep_size", [1, 2])
 @pytest.mark.parametrize("sync_weight_amax", [True, False])
 def test_layer_sync_moe_local_experts_amax(dist_workers, ep_size, sync_weight_amax):