[OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts

CHANGELOG.rst

@@ -96,6 +96,7 @@ Changelog
 - Fix the column-major ``DequantizeLinear`` surgery leaving a negative or default ``axis`` attribute unchanged while transposing the weight it refers to, which changed the model's output. Re-run the surgery on affected models.
 - Fix QAT of Hugging Face fused MoE experts with non-reentrant activation checkpointing, as in the ``examples/llm_qat`` FSDP2 config: the experts' ``F.linear`` interception stayed installed after the recompute, and training crashed with ``aten.mm.default got mixed torch.Tensor and DTensor``.
 - Fix quantization-aware training of Diffusers models failing in the backward pass with ``NotImplementedError`` once their attention modules are quantized, with or without MHA quantization.
+- Fix weight calibration, GPTQ and weight folding of Hugging Face fused MoE experts sharded by the expert parallelism of Transformers 5.16 and later, which used another rank's expert weights on every rank but the first.
 
 0.48.0 (2026-10-xx)
 ^^^^^^^^^^^^^^^^^^^

modelopt/torch/quantization/model_calib.py

@@ -1158,8 +1158,11 @@ def _expert_hook(_input_quantizer, args):
                 return
             idx = expert_module._current_expert_idx
             if idx in enabled:
-                # Read the weight fresh (valid under accelerate/FSDP re-materialization).
-                capture(quantizers[idx], getattr(expert_module, weight_name)[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``).
+                weight = expert_module._local_experts(getattr(expert_module, weight_name))
+                capture(quantizers[idx], weight[idx], args[0])
 
         return _expert_hook
 
@@ -1185,7 +1188,11 @@ def _dense_hook(linear, args):
                     module
                 ):
                     if on_weight is not None:
-                        on_weight(f"{name}.{weight_name}", weight[0], quantizers[0])
+                        # This rank's first expert, the one quantizers[0] quantizes: indexing an
+                        # expert-parallel DTensor names a global expert and all-gathers every expert.
+                        on_weight(
+                            f"{name}.{weight_name}", module._local_experts(weight)[0], quantizers[0]
+                        )
                     # Snapshot which experts are enabled now, before the caching forward silences
                     # all weight quantizers — so we don't capture (and discard) disabled experts.
                     enabled = {i for i, q in enumerate(quantizers) if q.is_enabled}
@@ -2592,10 +2599,13 @@ def _expert_hook(_input_quantizer, _args, output):
             for weight_name, weight, quantizers, input_quantizer in _fused_expert_projections(
                 module
             ):
-                # An expert's logged error would cost two full matmuls; gptq() logs a summary.
+                # This rank's experts, in quantizer order: an expert-parallel weight is a DTensor
+                # whose index names a global expert and all-gathers every expert. An expert's logged
+                # error would cost two full matmuls; gptq() logs a summary.
+                local = module._local_experts(weight)
                 per_expert = {
                     idx: _make_gptq_handle(
-                        f"{name}.{weight_name}.{idx}", module, q, weight[idx], log_mse=False
+                        f"{name}.{weight_name}.{idx}", module, q, local[idx], log_mse=False
                     )
                     for idx, q in enumerate(quantizers)
                     if q.is_enabled

modelopt/torch/quantization/plugins/huggingface.py

@@ -27,7 +27,7 @@
 import torch.nn as nn
 import transformers
 from torch import Tensor
-from torch.distributed.tensor import Shard
+from torch.distributed.tensor import DTensor, Shard
 from torch.nn.functional import linear
 from transformers.integrations.finegrained_fp8 import FP8Experts, FP8Linear
 from transformers.models.t5.modeling_t5 import T5Attention
@@ -898,6 +898,27 @@ def _setup(self):
         # Per forward: each projection's experts, fake-quantized together (see _quantized_weight).
         self._register_temp_attribute("_stacked_weights", {})
 
+    def _local_experts(self, weight: torch.Tensor) -> torch.Tensor:
+        """``weight`` indexed the way the per-expert quantizers are: by this module's own experts.
+
+        Under expert parallelism transformers 5.16 and later keep a fused expert weight as a DTensor
+        sharded over the experts (earlier ones keep a plain tensor of this rank's experts) and set
+        ``num_experts`` to the local count, so indexing the DTensor would read the GLOBAL expert: from
+        rank 1 on, another rank's. Its local shard holds this rank's experts in quantizer order. A
+        DTensor holding every expert (FSDP2) is indexed as it is, and one whose local shard holds
+        another number of experts than the quantizers is refused.
+        """
+        if isinstance(weight, DTensor) and weight.shape[0] != self.num_experts:
+            local = weight.to_local()
+            if local.shape[0] != self.num_experts:
+                raise RuntimeError(
+                    f"{type(self).__name__} has {self.num_experts} experts, but its expert weight, a "
+                    f"DTensor placed {weight.placements}, holds {local.shape[0]} on this rank: only a "
+                    "weight sharded over the experts can be indexed by local expert."
+                )
+            return local
+        return weight
+
     def _quantize_grouped_weight(self, weight: torch.Tensor, quantizers: nn.ModuleList):
         """Fake-quantize a fused 3-D weight per expert.
 
@@ -1046,8 +1067,12 @@ def forward(self, *args, **kwargs):
             self._stacked_weights = {}
 
     def get_expert_weight(self, name, idx, dtype=None):
-        """Return one expert's projection in the calibration/export weight domain."""
-        return getattr(self, name)[idx]
+        """Return one expert's projection in the calibration/export weight domain.
+
+        Indexed by this module's own experts, as the per-expert quantizers are (see
+        :meth:`_local_experts`).
+        """
+        return self._local_experts(getattr(self, name))[idx]
 
     def iter_weights_for_calibration(self):
         """Yield ``(weight_slice, quantizer)`` per-expert pairs.
@@ -1137,11 +1162,11 @@ def _setup(self):
         super()._setup()
 
     def get_expert_weight(self, name, idx, dtype=None):
-        weight = getattr(self, name)[idx]
+        weight = self._local_experts(getattr(self, name))[idx]
         if weight.element_size() > 1:
             return weight if dtype is None else weight.to(dtype)
         dtype = dtype or torch.get_default_dtype()
-        scales = getattr(self, f"{name}_scale_inv")[idx].float()
+        scales = self._local_experts(getattr(self, f"{name}_scale_inv"))[idx].float()
         rows, cols = self.block_size or (weight.shape[0] // scales.shape[0], weight.shape[1])
         scales = scales.repeat_interleave(rows, 0).repeat_interleave(cols, 1)
         return (weight.float() * scales[: weight.shape[0], : weight.shape[1]]).to(dtype)
@@ -1163,13 +1188,22 @@ def fold_weight(self, keep_attrs: bool = False):
 
     def unpack_weight(self):
         """Materialize floating weights before folding; retain scales for native forward."""
-        for name in (self._first_proj_attr, "down_proj"):
+        names = [
+            n for n in (self._first_proj_attr, "down_proj") if getattr(self, n).element_size() == 1
+        ]
+        # Sharded over the experts, one unpacked parameter would have to stay sharded too. Checked for
+        # both projections before either is unpacked, so a refusal leaves the module as it was.
+        if any(self._local_experts(getattr(self, n)) is not getattr(self, n) for n in names):
+            raise NotImplementedError(
+                f"{type(self).__name__} cannot unpack native FP8 expert weights sharded over the "
+                "experts. Quantize the model without expert parallelism."
+            )
+        for name in names:
             weight = getattr(self, name)
-            if weight.element_size() == 1:
-                unpacked = torch.empty_like(weight, dtype=torch.get_default_dtype())
-                for idx in range(self.num_experts):
-                    unpacked[idx].copy_(self.get_expert_weight(name, idx))
-                setattr(self, name, nn.Parameter(unpacked, requires_grad=False))
+            unpacked = torch.empty_like(weight, dtype=torch.get_default_dtype())
+            for idx in range(self.num_experts):
+                unpacked[idx].copy_(self.get_expert_weight(name, idx))
+            setattr(self, name, nn.Parameter(unpacked, requires_grad=False))
 
 
 def _get_fused_experts_quantizer_attr_names(module):