diff --git a/CHANGELOG.rst b/CHANGELOG.rst index e30bb9e245e..6b3dc978b29 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 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/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..8a2d0001086 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,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): diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 7c8dac3d500..22c91454332 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,239 @@ 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) + # 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): + """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)