From 77fb51fee70188d0b5836eebf6137139d1677f2b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 17:42:37 +0000 Subject: [PATCH 1/3] Index expert-parallel fused expert weights by this rank's experts Under transformers' expert parallelism a fused expert weight is a DTensor sharded over the experts, and num_experts is the per-rank count. Weight calibration, folding, the weight-input captures and GPTQ read those weights outside the forward by local expert index, but an index into the DTensor names a global expert: every rank but the first calibrated and folded another rank's experts, and every read all-gathered all of them. _local_experts hands those readers the rank's own shard, and refuses a DTensor whose local shard does not hold the module's experts. The native FP8 wrapper's unpack_weight refuses weights sharded over the experts rather than fill a buffer for every expert from this rank's experts. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + modelopt/torch/quantization/model_calib.py | 20 +- .../torch/quantization/plugins/huggingface.py | 40 ++- .../plugins/test_fused_experts.py | 243 +++++++++++++++++- 4 files changed, 293 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index e30bb9e245e..9e588b1c027 100755 --- a/CHANGELOG.rst +++ b/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 Transformers' expert parallelism, which used another rank's expert weights on every rank but the first. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 7e432cce9fb..ed61d0496fa 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/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 diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 0e6b1ad7271..756266ebeef 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/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,26 @@ 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 keeps a fused expert weight as a DTensor sharded over + the experts and sets ``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 +1066,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 +1161,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) @@ -1166,6 +1190,12 @@ def unpack_weight(self): for name in (self._first_proj_attr, "down_proj"): weight = getattr(self, name) if weight.element_size() == 1: + if self._local_experts(weight) is not weight: + # Sharded over the experts: one unpacked parameter would have to stay sharded too. + raise NotImplementedError( + f"{type(self).__name__} cannot unpack native FP8 expert weights sharded " + "over the experts. Quantize the model without expert parallelism." + ) 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)) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 7c8dac3d500..fbe2d0c8fbd 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -19,17 +19,23 @@ import warnings import weakref from contextlib import nullcontext +from unittest import mock import pytest import torch import torch.nn as nn import torch.nn.functional as F +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor +from torch.distributed.tensor.debug import CommDebugMode pytest.importorskip("transformers") from _test_utils.torch.quantization.tied_modules import tie_fused_experts_3d_params -from transformers import PretrainedConfig +from transformers import PretrainedConfig, Qwen3MoeConfig from transformers.integrations.finegrained_fp8 import FP8Experts +from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeExperts import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq @@ -62,6 +68,7 @@ register_fused_experts_on_the_fly, register_sparse_moe_on_the_fly, ) +from modelopt.torch.quantization.utils.calib_utils import GPTQHelper from modelopt.torch.quantization.utils.core_utils import module_name_maps, weight_attr_names @@ -2681,3 +2688,237 @@ def test_native_fp8_experts_reject_packed_weight_updates(cfg, algorithm): mtq.calibrate(experts, algorithm, forward_loop=lambda m: m(inputs, indices, routing)) for name, weight in unpacked.items(): assert not torch.equal(getattr(experts, name), weight) + + +# --------------------------------------------------------------------------- +# Fused experts under expert parallelism +# --------------------------------------------------------------------------- +_EP_EXPERTS, _EP_HIDDEN = 4, 8 + + +class _ExpertsHolder(nn.Module): + def __init__(self, experts): + super().__init__() + self.experts = experts + + def forward(self, *args): + return self.experts(*args) + + +_WEIGHT_ONLY_INT8_EXPERTS_CFG = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*proj_weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": "max", +} + + +def _ep_experts_config(): + return Qwen3MoeConfig( + hidden_size=_EP_HIDDEN, + moe_intermediate_size=2 * _EP_HIDDEN, + intermediate_size=4 * _EP_HIDDEN, + num_experts=_EP_EXPERTS, + num_experts_per_tok=TOP_K, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + vocab_size=32, + ) + + +def _experts_with_distinct_scales(config): + """Whole fused experts with a distinct weight scale per expert, so a mix-up cannot pass.""" + experts = Qwen3MoeExperts(config) + with torch.no_grad(): + for e in range(_EP_EXPERTS): + experts.gate_up_proj[e].normal_(std=0.1 * (e + 1)) + experts.down_proj[e].normal_(std=0.1 * (e + 1)) + return experts + + +def _sharded_like_transformers(whole, config, mesh, size): + """``whole``'s experts as transformers' expert parallelism shards them (``MoEParamShard``): a + DTensor ``Shard(0)`` over the experts, ``num_experts`` set to the local count.""" + experts = Qwen3MoeExperts(config) + experts.num_experts = _EP_EXPERTS // size + for name in ("gate_up_proj", "down_proj"): + full = getattr(whole, name).detach().clone() + experts._parameters[name] = nn.Parameter(distribute_tensor(full, mesh, [Shard(0)])) + return experts + + +def _ep_sharded_experts_use_their_own_weights_job(rank, size, layout): + """Each rank's per-expert quantizers must calibrate on, and fold into, the rank's own experts, as + in the whole module -- not on the global experts their local index names. ``fsdp2-dense`` is the + dense-sharded expert-parallel layout: the decoder layer FSDP2-sharded over ``dp`` with its + experts left out (``ignored_params``), so weight calibration opens FSDP2's weight-access window + on the layer, which must leave the expert DTensors alone -- gathering them would hand every rank + every expert, indexed by local position.""" + config = _ep_experts_config() + whole = _ExpertsHolder(_experts_with_distinct_scales(config)) + n_local = _EP_EXPERTS // size + owned = slice(rank * n_local, (rank + 1) * n_local) + # Unnamed meshes give the window nothing but FSDP2's own bookkeeping to tell its shards by. + named = layout != "fsdp2-dense-unnamed-meshes" + ep_mesh = init_device_mesh("cpu", (size,), mesh_dim_names=("ep",) if named else None) + experts = _sharded_like_transformers(whole.experts, config, ep_mesh, size) + if layout != "plain": + layer = nn.Module() + layer.dense = nn.Linear(_EP_HIDDEN, _EP_HIDDEN) + layer.experts = experts + dp_mesh = init_device_mesh("cpu", (size,), mesh_dim_names=("dp",) if named else None) + fully_shard(layer, mesh=dp_mesh, ignored_params=set(experts.parameters())) + model = nn.Module() + model.layer = layer + else: + model = _ExpertsHolder(experts) + + mtq.quantize(whole, _WEIGHT_ONLY_INT8_EXPERTS_CFG) + mtq.quantize(model, _WEIGHT_ONLY_INT8_EXPERTS_CFG) + sharded = model.experts if layout == "plain" else model.layer.experts + for name in ("gate_up_proj_weight_quantizers", "down_proj_weight_quantizers"): + local_quantizers = getattr(sharded, name) + assert len(local_quantizers) == n_local + for i, quantizer in enumerate(local_quantizers): + reference = getattr(whole.experts, name)[owned.start + i] + torch.testing.assert_close( + quantizer.amax, reference.amax, atol=0, rtol=0, msg=f"{name}[{i}]" + ) + # The expert weights stay sharded: nothing gathered them for good. + assert sharded.gate_up_proj.to_local().shape[0] == n_local + + if layout == "plain": + whole.experts.fold_weight() + sharded.fold_weight() + for name in ("gate_up_proj", "down_proj"): + torch.testing.assert_close( + getattr(sharded, name).to_local(), + getattr(whole.experts, name)[owned], + atol=0, + rtol=0, + ) + + +@pytest.mark.parametrize("layout", ["plain", "fsdp2-dense", "fsdp2-dense-unnamed-meshes"]) +def test_ep_sharded_experts_calibrate_and_fold_their_own_weights(cpu_dist_workers, layout): + """Weight calibration and folding run outside the forward, where expert-parallel weights are still + DTensors sharded over the experts: each rank must use its own experts, not the global ones. With + the dense part FSDP2-sharded, the calibration runs inside FSDP2's weight-access window, which + must not gather the expert-parallel weights FSDP2 does not own, its meshes named or not.""" + cpu_dist_workers.run(_ep_sharded_experts_use_their_own_weights_job, layout=layout) + + +def _expert_layouts_the_wrappers_cannot_index_job(rank, size): + mesh = init_device_mesh("cpu", (size,)) + n_local = _EP_EXPERTS // size + # Every expert on every rank, with num_experts at the local count: its local shard is not this + # rank's experts, so indexing it by local expert would read the global ones. + experts = _experts_with_distinct_scales(_ep_experts_config()) + experts.num_experts = n_local + for name in ("gate_up_proj", "down_proj"): + full = getattr(experts, name).detach().clone() + experts._parameters[name] = nn.Parameter(distribute_tensor(full, mesh, [Replicate()])) + with pytest.raises(RuntimeError, match=f"holds {_EP_EXPERTS} on this rank"): + mtq.quantize(_ExpertsHolder(experts), _WEIGHT_ONLY_INT8_EXPERTS_CFG) + + # Native FP8 experts sharded over the experts would unpack into one buffer for every expert. + config = PretrainedConfig( + hidden_size=32, intermediate_size=32, num_experts=_EP_EXPERTS, hidden_act="silu" + ) + fp8 = FP8Experts(config, block_size=None, has_gate=True) + for name in ("gate_up_proj", "down_proj"): + with torch.no_grad(): + getattr(fp8, name).copy_(torch.randn(getattr(fp8, name).shape)) + getattr(fp8, f"{name}_scale_inv").fill_(0.01) + mtq.quantize(fp8, _WEIGHT_ONLY_INT8_EXPERTS_CFG) + owned = fp8.gate_up_proj.detach()[rank * n_local : (rank + 1) * n_local] + fp8.gate_up_proj = nn.Parameter( + DTensor.from_local(owned, mesh, [Shard(0)], run_check=False), requires_grad=False + ) + fp8.num_experts = n_local + with pytest.raises(NotImplementedError, match="sharded over the experts"): + fp8.unpack_weight() + + +def test_expert_weights_the_quantizers_cannot_index_are_refused(cpu_dist_workers): + """A DTensor expert weight whose local shard does not hold the module's experts would be indexed + as global experts; native FP8 experts sharded over the experts would unpack into a buffer for + every expert. Both raise instead of calibrating or exporting the wrong experts.""" + cpu_dist_workers.run(_expert_layouts_the_wrappers_cannot_index_job) + + +def _ep_weight_input_hooks_read_local_experts_job(rank, size): + """Registering the weight-input hooks (local Hessian, imatrix) hands ``on_weight`` each + projection's first LOCAL expert, the one its first quantizer quantizes, and runs no collective: + indexing the expert-parallel DTensor would name a global expert and all-gather every expert.""" + config = _ep_experts_config() + torch.manual_seed(0) + whole = _experts_with_distinct_scales(config) + mesh = init_device_mesh("cpu", (size,)) + model = _ExpertsHolder(_sharded_like_transformers(whole, config, mesh, size)) + n_local = _EP_EXPERTS // size + mtq.quantize(model, _WEIGHT_ONLY_INT8_EXPERTS_CFG) + + seen, captured = {}, [] + with CommDebugMode() as comms: + handles = _register_weight_input_hooks( + model, + module_name_maps(model), + lambda quantizer, weight, inputs: captured.append(weight), + on_weight=lambda name, weight, quantizer: seen.__setitem__(name, (weight, quantizer)), + ) + # An eager forward's per-expert call: the capture reads this rank's first expert too. + model.experts._current_expert_idx = 0 + model.experts.gate_up_proj_input_quantizer(torch.randn(2, _EP_HIDDEN)) + for handle in handles: + handle.remove() + assert comms.get_total_counts() == 0, dict(comms.get_comm_counts()) + for name in ("gate_up_proj", "down_proj"): + weight, quantizer = seen[f"experts.{name}"] + assert quantizer is getattr(model.experts, f"{name}_weight_quantizers")[0] + local_first = getattr(whole, name)[rank * n_local].detach() + torch.testing.assert_close(weight, local_first, atol=0, rtol=0, msg=name) + torch.testing.assert_close( + captured[0], whole.gate_up_proj[rank * n_local].detach(), atol=0, rtol=0 + ) + + +def test_ep_weight_input_hooks_read_this_ranks_experts_without_a_collective(cpu_dist_workers): + """Hook registration only needs one expert's shape, so it must not gather other ranks' experts.""" + cpu_dist_workers.run(_ep_weight_input_hooks_read_local_experts_job) + + +def _ep_gptq_helpers_take_local_experts_job(rank, size): + """GPTQ builds each expert's helper from this rank's expert, the one its quantizer quantizes: + indexing the expert-parallel DTensor would name a global expert and all-gather every expert.""" + config = _ep_experts_config() + torch.manual_seed(0) + whole = _experts_with_distinct_scales(config) + mesh = init_device_mesh("cpu", (size,)) + model = _ExpertsHolder(_sharded_like_transformers(whole, config, mesh, size)) + n_local = _EP_EXPERTS // size + built = {} + init = GPTQHelper.__init__ + + def recording(helper, module, name, *args, weight=None, **kwargs): + if weight is not None: + built[name] = weight if isinstance(weight, DTensor) else weight.detach().clone() + init(helper, module, name, *args, weight=weight, **kwargs) + + cfg = {**_WEIGHT_ONLY_INT8_EXPERTS_CFG, "algorithm": {"method": "gptq", "block_size": 8}} + with mock.patch.object(GPTQHelper, "__init__", recording): + mtq.quantize(model, cfg, forward_loop=lambda m: None) + + assert len(built) == 2 * n_local, sorted(built) + for name, weight in built.items(): + _, proj, idx = name.rsplit(".", 2) + assert not isinstance(weight, DTensor), name + expected = getattr(whole, proj)[rank * n_local + int(idx)].detach() + torch.testing.assert_close(weight, expected, atol=0, rtol=0, msg=name) + + +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) From c56d68304264ee1b4aae787962e8f08df42ba2f9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 18:10:14 +0000 Subject: [PATCH 2/3] Refuse expert-sharded FP8 experts before unpacking either projection unpack_weight checked each projection as it unpacked it, so with only down_proj sharded over the experts it unpacked gate_up_proj and then raised, leaving the module half-unpacked: a state fold_weight refuses and unpack_weight cannot finish. Both projections are now checked first. The test shards down_proj alone and checks gate_up_proj is still FP8. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/huggingface.py | 27 ++++++++++--------- .../plugins/test_fused_experts.py | 6 +++-- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 756266ebeef..5eb1a3ba723 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1187,19 +1187,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: - if self._local_experts(weight) is not weight: - # Sharded over the experts: one unpacked parameter would have to stay sharded too. - raise NotImplementedError( - f"{type(self).__name__} cannot unpack native FP8 expert weights sharded " - "over the experts. Quantize the model without expert parallelism." - ) - 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): diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index fbe2d0c8fbd..22c91454332 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -2833,13 +2833,15 @@ def _expert_layouts_the_wrappers_cannot_index_job(rank, size): getattr(fp8, name).copy_(torch.randn(getattr(fp8, name).shape)) getattr(fp8, f"{name}_scale_inv").fill_(0.01) mtq.quantize(fp8, _WEIGHT_ONLY_INT8_EXPERTS_CFG) - owned = fp8.gate_up_proj.detach()[rank * n_local : (rank + 1) * n_local] - fp8.gate_up_proj = nn.Parameter( + # Only the second projection sharded: the refusal must come before the first is unpacked. + owned = fp8.down_proj.detach()[rank * n_local : (rank + 1) * n_local] + fp8.down_proj = nn.Parameter( DTensor.from_local(owned, mesh, [Shard(0)], run_check=False), requires_grad=False ) fp8.num_experts = n_local with pytest.raises(NotImplementedError, match="sharded over the experts"): fp8.unpack_weight() + assert fp8.gate_up_proj.dtype == torch.float8_e4m3fn def test_expert_weights_the_quantizers_cannot_index_are_refused(cpu_dist_workers): From 1e74c31b49ad04890fe55671fbfb3f8e216a6df1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 18:13:15 +0000 Subject: [PATCH 3/3] Say the expert-parallel DTensor experts are transformers 5.16's transformers keeps expert-parallel expert weights as DTensors sharded over the experts from 5.16 on (MoEParamShard); before, GroupedGemmParallel kept each rank's slice as a plain tensor, which ModelOpt indexed correctly. The changelog entry and _local_experts' docstring now say so. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- modelopt/torch/quantization/plugins/huggingface.py | 11 ++++++----- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9e588b1c027..6b3dc978b29 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -96,7 +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 Transformers' expert parallelism, which used another rank's expert weights on every rank but the first. +- 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) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 5eb1a3ba723..8a2d0001086 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -901,11 +901,12 @@ def _setup(self): 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 keeps a fused expert weight as a DTensor sharded over - the experts and sets ``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. + 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()