diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 48aed5eb8b9..f2c9cffe5a8 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,7 @@ Changelog **New Features** +- Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager. - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 `_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change). - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred. - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``. diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 0ca30d18448..76a9d673748 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -890,8 +890,8 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): description=( "If True, max-calibration synchronizes the weight quantizer amax across local " "experts within each SequentialMLP layer, so all experts in that layer share " - "one effective weight amax. TEGroupedMLP already fuses experts into a single " - "GEMM with one weight quantizer, so this flag is irrelevant there." + "one effective weight amax. TEGroupedMLP keeps a per-expert weight quantizer " + "(GroupedQuantizer) whose amax follows the same expert-parallel sync rule." ), ) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 7e5bb85c09b..b7f3c9bf0fb 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -42,7 +42,13 @@ from .calib import MseCalibrator, NVFP4MSECalibrator, _Calibrator from .conversion import create_and_replace_svdquant_linear_on_the_fly, set_quantizer_by_cfg_context -from .nn import NVFP4StaticQuantizer, QuantModule, SequentialQuantizer, TensorQuantizer +from .nn import ( + GroupedQuantizer, + NVFP4StaticQuantizer, + QuantModule, + SequentialQuantizer, + TensorQuantizer, +) from .utils import ( SHARED_PATTERNS, SharedWeightGlobalAmaxState, @@ -205,7 +211,7 @@ def _has_expert_parallelism(module: nn.Module) -> bool: def _iter_leaf_quantizers(quantizer): - if isinstance(quantizer, SequentialQuantizer): + if isinstance(quantizer, (SequentialQuantizer, GroupedQuantizer)): for _q in quantizer: yield from _iter_leaf_quantizers(_q) return @@ -314,12 +320,12 @@ def max_calibrate( for name, module in model.named_modules(): if isinstance(module, QuantModule) and _has_expert_parallelism(module): for child in module.children(): - if isinstance(child, (TensorQuantizer, SequentialQuantizer)): + if isinstance(child, (TensorQuantizer, SequentialQuantizer, GroupedQuantizer)): _check_moe_calibration_complete(child, module.parallel_state) def sync_quantizer_amax_across_dp_ep(quantizer, parallel_state, parent_name, child_name): """Sync amax across DP (always) and EP (filtered — see _should_sync_amax_across_ep).""" - if isinstance(quantizer, SequentialQuantizer): + if isinstance(quantizer, (SequentialQuantizer, GroupedQuantizer)): for _q in quantizer: sync_quantizer_amax_across_dp_ep(_q, parallel_state, parent_name, child_name) return @@ -333,7 +339,7 @@ def sync_quantizer_amax_across_dp_ep(quantizer, parallel_state, parent_name, chi for name, module in model.named_modules(): if isinstance(module, QuantModule): for child_name, child in module.named_children(): - if isinstance(child, (TensorQuantizer, SequentialQuantizer)): + if isinstance(child, (TensorQuantizer, SequentialQuantizer, GroupedQuantizer)): sync_quantizer_amax_across_dp_ep(child, module.parallel_state, name, child_name) # Step 3: TP sync # Objective: the quantization parameters when TP = 8 then changed to TP=4 then back to TP=8 should be the same @@ -355,7 +361,7 @@ def sync_quantizer_amax_across_tp( parallel_state: ParallelState, ): # Syncing amax across TP for sequential quantizer - if isinstance(quantizer, SequentialQuantizer): + if isinstance(quantizer, (SequentialQuantizer, GroupedQuantizer)): for _q in quantizer: sync_quantizer_amax_across_tp( _q, linear_name, quantizer_type, axes_for_sync, parallel_state diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 98d9e0dcb1e..6a634cd806d 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -65,6 +65,7 @@ from ..functional import normalized_hadamard_transform __all__ = [ + "GroupedQuantizer", "HardDisabledTensorQuantizer", "NVFP4StaticQuantizer", "SequentialQuantizer", @@ -1641,3 +1642,71 @@ def convert_to_single_quantizer(model, indx: int = 0): ) in original_sequential_quantizers.items(): for name, sequential_quantizer in sequential_quantizers_list: setattr(parent_module, name, sequential_quantizer) + + +class GroupedQuantizer(nn.ModuleList): + """A container for per-group :class:`TensorQuantizer` modules. + + Used when a single linear holds several independently-quantized weights — e.g. the + fused experts of a TEGroupedLinear, where each of the ``num_gemms`` weights needs its + own ``amax``. Unlike :class:`SequentialQuantizer` (an ``nn.Sequential`` that *chains* + quantizers over one tensor), the contained quantizers act on *different* tensors, so + there is no inherent forward path: index in with ``grouped[i](weight_i)``. + + Property reads (``amax``, ``is_enabled``) delegate to the first quantizer — all members + share one config, so the first is representative for "is this calibrated/enabled" + checks; the real per-group values live on the members and are used via indexing. + Lifecycle/config methods broadcast to every member. + """ + + _delegated_properties = ["fake_quant", "is_enabled", "amax"] + _delegated_methods = [ + "reset_amax", + "disable", + "enable", + "load_calib_amax", + "load_calib_bias", + ] + + def __init__(self, *quantizers: "TensorQuantizer | SequentialQuantizer"): + """Initialize GroupedQuantizer module.""" + super().__init__(quantizers) + assert all(isinstance(q, (TensorQuantizer, SequentialQuantizer)) for q in self), ( + "All quantizers must be a TensorQuantizer or SequentialQuantizer." + ) + + def forward(self, inputs): + """Apply the representative quantizer for single-weight compatibility paths.""" + return self[0](inputs) + + def __getattr__(self, name): + """Delegate property reads to the first member and method calls to all members.""" + if name in self._delegated_properties: + return getattr(self[0], name) + + if name in self._delegated_methods: + + def method_wrapper(*args, **kwargs): + return [getattr(quantizer, name)(*args, **kwargs) for quantizer in self] + + return method_wrapper + + return super().__getattr__(name) + + def __setattr__(self, name, value): + if name in self._delegated_properties: + for quantizer in self: + setattr(quantizer, name, value) + else: + super().__setattr__(name, value) + + def set_from_attribute_config(self, attributes): + """Set the attributes of contained quantizers; a single config broadcasts to all.""" + if not isinstance(attributes, (list, tuple)): + attributes = [attributes] * len(self) + for attribute, quantizer in zip(attributes, self): + quantizer.set_from_attribute_config(attribute) + + def get_modelopt_state(self) -> dict[str, Any]: + """Get meta state to be saved in checkpoint.""" + return {"num_quantizers": len(self), "is_grouped_quantizer": True} diff --git a/modelopt/torch/quantization/plugins/custom.py b/modelopt/torch/quantization/plugins/custom.py index f480d245daa..d16596908c9 100644 --- a/modelopt/torch/quantization/plugins/custom.py +++ b/modelopt/torch/quantization/plugins/custom.py @@ -24,7 +24,13 @@ from modelopt.torch.utils.distributed import ParallelState -from ..nn import NVFP4StaticQuantizer, QuantModule, SequentialQuantizer, TensorQuantizer +from ..nn import ( + GroupedQuantizer, + NVFP4StaticQuantizer, + QuantModule, + SequentialQuantizer, + TensorQuantizer, +) from ..nn.modules.quant_linear import _QuantLinear from ..utils import multi_context, replace_function @@ -134,11 +140,19 @@ def _check_unsupported_states(quantizer: TensorQuantizer): def _has_state(quantizer, name): # Handling for SequentialQuantizer - quantizer = quantizer[0] if isinstance(quantizer, SequentialQuantizer) else quantizer + quantizer = ( + quantizer[0] + if isinstance(quantizer, (SequentialQuantizer, GroupedQuantizer)) + else quantizer + ) return hasattr(quantizer, name) def _has_complete_static_nvfp4_weight_state(quantizer, weight): - quantizer = quantizer[0] if isinstance(quantizer, SequentialQuantizer) else quantizer + quantizer = ( + quantizer[0] + if isinstance(quantizer, (SequentialQuantizer, GroupedQuantizer)) + else quantizer + ) if not isinstance(quantizer, NVFP4StaticQuantizer): return False amax = getattr(quantizer, "_amax", None) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 752dd801a6e..66f43801eca 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -15,6 +15,7 @@ """Support quantization for megatron linear layers.""" +import re import types from contextlib import contextmanager from typing import Any @@ -42,7 +43,13 @@ from ..algorithms import AutoQuantizeGradientSearcher from ..conversion import maybe_promote_nvfp4_static_quantizer -from ..nn import QuantModule, QuantModuleRegistry, SequentialQuantizer, TensorQuantizer +from ..nn import ( + GroupedQuantizer, + QuantModule, + QuantModuleRegistry, + SequentialQuantizer, + TensorQuantizer, +) from ..nn.modules.quant_linear import RealQuantLinear from ..qtensor import QTensorWrapper from ..utils import sync_moe_expert_amax @@ -90,7 +97,7 @@ def _check_nvfp4_static_tp_supported(model: torch.nn.Module) -> None: continue leaves = ( list(weight_quantizer) - if isinstance(weight_quantizer, SequentialQuantizer) + if isinstance(weight_quantizer, (SequentialQuantizer, GroupedQuantizer)) else [weight_quantizer] ) if any(leaf.is_nvfp4_static for leaf in leaves): @@ -697,8 +704,13 @@ def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): return super()._load_from_state_dict(filtered_state_dict, prefix, *args, **kwargs) 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 quantizers have independent checkpoint keys. Preserve their native + # scalar, channel, or block shape instead of flattening them through the legacy + # single-quantizer path. + if re.match(r"weight_quantizer\.\d+\..+_amax$", k): + quantizer_state_dict[k] = v + else: + quantizer_state_dict[k] = v.view(-1) if v.numel() == 1 else v @QuantModuleRegistry.register( {TEColumnParallelGroupedLinear: "megatron_TEColumnParallelGroupedLinear"} diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index d0efcc52db1..38004f905d5 100644 --- a/modelopt/torch/quantization/plugins/transformer_engine.py +++ b/modelopt/torch/quantization/plugins/transformer_engine.py @@ -15,7 +15,9 @@ """Support quantization for Transformer Engine layers.""" +import copy import inspect +import os import warnings import torch @@ -27,11 +29,13 @@ from modelopt.torch.quantization.utils import replace_function -from ..nn import QuantModuleRegistry +from ..nn import GroupedQuantizer, QuantModuleRegistry, SequentialQuantizer from .custom import _ParallelLinear _TE_VERSION = Version(te.__version__) +_COMPILE_TEGROUPED_WEIGHT_LOOP_ENV = "MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP" + def _assert_te_fp8_enabled(): """Check if Transformer Engine FP8 autocast is enabled and raise error if so.""" @@ -48,6 +52,13 @@ def _assert_te_fp8_enabled(): pass # Older TE versions may not have this API +def _is_calibrating(quantizer): + """Return whether a tensor or sequential quantizer is collecting calibration stats.""" + if isinstance(quantizer, SequentialQuantizer): + return any(getattr(q, "_if_calib", False) for q in quantizer) + return getattr(quantizer, "_if_calib", False) + + @QuantModuleRegistry.register({te.pytorch.Linear: "te_Linear"}) class _QuantTELinear(_ParallelLinear): @property @@ -137,8 +148,27 @@ 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. + # Each fused expert gets its own weight quantizer (independent amax). Storing them in a + # GroupedQuantizer (an nn.ModuleList) surfaces them as ``weight_quantizer.{i}``, which the + # fused-experts name normalizer maps to ``*weight_quantizer`` so the stock configs apply. + self.weight_quantizer = GroupedQuantizer( + *(copy.deepcopy(self.weight_quantizer) for _ in range(self.num_gemms)) + ) + + # Compile only the per-expert quantizer loop. The surrounding TE grouped GEMM remains + # eager, and the opt-in flag leaves the default execution path unchanged. + if os.getenv(_COMPILE_TEGROUPED_WEIGHT_LOOP_ENV, "0") == "1": + quantizers = tuple(self.weight_quantizer) + + def quantize_weights(*weights): + return tuple(quantizer(weight) for quantizer, weight in zip(quantizers, weights)) + + self._compiled_weight_quantizer_loop = torch.compile( + quantize_weights, + backend="inductor", + fullgraph=False, + mode="reduce-overhead", + ) def modelopt_post_restore(self, prefix: str = ""): # GroupedMLP stores the weights as weight0, weight1, etc. To run post_restore in order to @@ -150,12 +180,29 @@ def modelopt_post_restore(self, prefix: str = ""): # Remove self.weight after post_restore. delattr(self, "weight") + # Base post_restore only re-calibrates the first member; the other per-expert + # quantizers also need re-calibration so a TP/EP change between save and restore + # produces correctly shaped per-channel _amax. Mirror base behavior: only + # re-calibrate quantizers whose loaded state had _amax (skip unused ones). + from modelopt.torch.quantization.model_calib import max_calibrate + + for i in range(self.num_gemms): + weight_i = getattr(self, f"weight{i}", None) + if weight_i is None: + continue + wq_i = self.weight_quantizer[i] + q = wq_i[0] if isinstance(wq_i, SequentialQuantizer) else wq_i + if not hasattr(q, "_amax"): + continue + wq_i.reset_amax() + max_calibrate(wq_i, lambda wq, w=weight_i: wq(w), distributed_sync=False) + def iter_weights_for_calibration(self): """Yield ``(weight_i, weight_quantizer)`` for each of the ``num_gemms`` grouped weights.""" for i in range(self.num_gemms): weight_i = getattr(self, f"weight{i}", None) if weight_i is not None: - yield weight_i, self.weight_quantizer + yield weight_i, self.weight_quantizer[i] @staticmethod def te_grouped_quantized_linear_fn(package, func_name, self, *args): @@ -184,8 +231,19 @@ 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]) + weights = tuple(args[weights_start : weights_start + num_gemms]) + # Calibration mutates collector state and must stay outside Inductor/CUDAGraph capture. + use_compiled_loop = hasattr(self, "_compiled_weight_quantizer_loop") and not any( + _is_calibrating(quantizer) for quantizer in self.weight_quantizer + ) + if use_compiled_loop: + quantized_weights = self._compiled_weight_quantizer_loop(*weights) + else: + quantized_weights = tuple( + self.weight_quantizer[gemm_idx](weight) for gemm_idx, weight in enumerate(weights) + ) + for gemm_idx, quantized_weight in enumerate(quantized_weights): + new_args[weights_start + gemm_idx] = quantized_weight 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. diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 3fded1c2864..29092c12fd5 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -14,6 +14,7 @@ # limitations under the License. import copy +import math from contextlib import nullcontext from functools import partial from pathlib import Path @@ -58,9 +59,13 @@ from modelopt.torch.quantization.algorithms import QuantRecipe, _AutoQuantizeBaseSearcher from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins.megatron import ( + _QuantMegatronTEGroupedLinear, _QuantTEMCoreRowParallelLinear, get_mcore_layerwise_calibration_layers, ) +from modelopt.torch.quantization.plugins.transformer_engine import ( + _COMPILE_TEGROUPED_WEIGHT_LOOP_ENV, +) from modelopt.torch.quantization.utils import is_quantized_linear from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -734,10 +739,10 @@ def _test_te_grouped_vs_sequential_quantize_helper(tp_size, ep_size, quant_cfg, # Quantize grouped model mtq.quantize(te_grouped_moe_model, quant_cfg, forward) - # Quantize non-grouped model with synced weight amax to match TEGroupedMLP behavior - seq_quant_cfg = copy.deepcopy(quant_cfg) - seq_quant_cfg["algorithm"] = {"method": "max", "sync_expert_weight_amax": True} - mtq.quantize(sequential_moe_model, seq_quant_cfg, forward) + # TEGroupedMLP now quantizes per-expert by default (GroupedQuantizer), matching + # SequentialMLP's per-expert quantizers, so no amax sync override is needed for the + # two models to produce identical quantized outputs. + mtq.quantize(sequential_moe_model, copy.deepcopy(quant_cfg), forward) # Compare model outputs after quantization te_grouped_moe_quant_output = forward(te_grouped_moe_model) @@ -757,6 +762,234 @@ def test_te_grouped_vs_sequential_quantize(dist_workers_size_4, quant_cfg): ) +def test_te_grouped_process_quantizer_amax_preserves_per_expert_shape(): + """Per-expert amax buffers retain their native checkpoint shape.""" + value = torch.randn(3, 2) + state_dict = {} + + _QuantMegatronTEGroupedLinear._process_quantizer_amax( + None, "weight_quantizer.2._amax", value, state_dict + ) + + assert state_dict["weight_quantizer.2._amax"] is value + assert state_dict["weight_quantizer.2._amax"].shape == (3, 2) + + +@pytest.mark.parametrize("compile_enabled", [False, True]) +def test_te_grouped_compiled_weight_quantizer_loop( + distributed_setup_size_1, monkeypatch, compile_enabled +): + """The opt-in flag controls compilation and preserves per-expert backward.""" + compile_kwargs = [] + compiled_calls = [] + + def fake_compile(fn, **kwargs): + compile_kwargs.append(kwargs) + + def compiled(*args): + compiled_calls.append(len(args)) + return fn(*args) + + return compiled + + if compile_enabled: + monkeypatch.setenv(_COMPILE_TEGROUPED_WEIGHT_LOOP_ENV, "1") + else: + monkeypatch.delenv(_COMPILE_TEGROUPED_WEIGHT_LOOP_ENV, raising=False) + monkeypatch.setattr(torch, "compile", fake_compile) + initialize_for_megatron(seed=SEED) + model = _gpt_model_provider( + tp_size=1, + hidden_size=32, + moe_grouped_gemm=True, + transformer_impl="transformer_engine", + num_moe_experts=4, + ) + forward = get_forward(model) + for module in model.modules(): + if isinstance(module, TopKRouter): + module.topk = module.num_experts + + mtq.quantize(model, copy.deepcopy(mtq.INT8_DEFAULT_CFG), forward) + grouped_modules = [ + module + for module in model.modules() + if isinstance(getattr(module, "weight_quantizer", None), mtq.nn.GroupedQuantizer) + ] + compiled_modules = [ + module for module in model.modules() if hasattr(module, "_compiled_weight_quantizer_loop") + ] + assert grouped_modules + assert len(compiled_modules) == (len(grouped_modules) if compile_enabled else 0) + assert len(compile_kwargs) == (len(grouped_modules) if compile_enabled else 0) + assert all( + kwargs == {"backend": "inductor", "fullgraph": False, "mode": "reduce-overhead"} + for kwargs in compile_kwargs + ) + # Calibration mutates collector state and must stay eager even when the flag is enabled. + assert not compiled_calls + + loss = forward(model).sum() + loss.backward() + if compile_enabled: + assert compiled_calls + assert set(compiled_calls) == {4} + else: + assert not compiled_calls + assert all( + torch.isfinite(getattr(module, f"weight{i}").grad).all() + for module in grouped_modules + for i in range(module.num_gemms) + ) + destroy_model_parallel() + + +def _test_te_grouped_vs_sequential_default_amax_helper(tp_size, ep_size, quant_cfg, rank, size): + """TEGrouped keeps a per-expert weight quantizer (GroupedQuantizer) by default; each + expert's amax should match the corresponding SequentialMLP expert (no cross-expert sharing).""" + initialize_for_megatron( + tensor_model_parallel_size=tp_size, + expert_model_parallel_size=ep_size, + seed=SEED, + ) + + te_grouped = _gpt_model_provider( + tp_size=tp_size, + ep_size=ep_size, + hidden_size=32, + moe_grouped_gemm=True, + transformer_impl="transformer_engine", + num_moe_experts=4, + ) + forward = get_forward(te_grouped, batch_size=8) + + sequential = _gpt_model_provider( + tp_size=tp_size, + ep_size=ep_size, + hidden_size=32, + moe_grouped_gemm=False, + num_moe_experts=4, + transformer_impl="modelopt", + ) + copy_weights_from_grouped_to_non_grouped(te_grouped, sequential) + + for module in te_grouped.modules(): + if isinstance(module, TopKRouter): + module.topk = module.num_experts + for module in sequential.modules(): + if isinstance(module, TopKRouter): + module.topk = module.num_experts + + mtq.quantize(te_grouped, quant_cfg, forward) + mtq.quantize(sequential, quant_cfg, forward) + + te_modules = [m for m in te_grouped.modules() if isinstance(m, TEGroupedMLP)] + seq_modules = [m for m in sequential.modules() if isinstance(m, SequentialMLP)] + assert len(te_modules) == len(seq_modules) + + saw_per_expert_divergence = False + for te_mlp, seq_mlp in zip(te_modules, seq_modules): + for linear_name in ("linear_fc1", "linear_fc2"): + te_wq = getattr(te_mlp, linear_name).weight_quantizer + # One weight quantizer per local expert, not a single shared one. + assert len(te_wq) == len(seq_mlp.local_experts), ( + f"{linear_name}: expected {len(seq_mlp.local_experts)} per-expert quantizers, " + f"got {len(te_wq)}" + ) + + expert_amaxes = [] + for i, expert in enumerate(seq_mlp.local_experts): + te_amax = te_wq[i].amax + seq_amax = getattr(expert, linear_name).weight_quantizer.amax + assert te_amax is not None + assert torch.allclose(te_amax, seq_amax, atol=1e-5, rtol=1e-5), ( + f"TEGrouped expert {i} amax != Sequential expert {i} amax for {linear_name}" + ) + expert_amaxes.append(te_amax.reshape(-1)[0]) + + stacked = torch.stack(expert_amaxes) + if (stacked.max() - stacked.min()).item() > 1e-5: + saw_per_expert_divergence = True + + assert saw_per_expert_divergence, ( + "Expected per-expert weight amax to diverge across experts (proves no cross-expert sharing)." + ) + + +@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG]) +def test_te_grouped_vs_sequential_default_amax(dist_workers_size_4, quant_cfg): + dist_workers_size_4.run( + partial(_test_te_grouped_vs_sequential_default_amax_helper, 1, 2, quant_cfg) + ) + + +def _test_te_grouped_vs_sequential_default_loss_helper(tp_size, ep_size, quant_cfg, rank, size): + """TEGrouped quantized output should diverge from BF16 more than SequentialMLP under default sync=False.""" + initialize_for_megatron( + tensor_model_parallel_size=tp_size, + expert_model_parallel_size=ep_size, + seed=SEED, + ) + + te_grouped = _gpt_model_provider( + tp_size=tp_size, + ep_size=ep_size, + hidden_size=32, + moe_grouped_gemm=True, + transformer_impl="transformer_engine", + num_moe_experts=4, + ) + forward = get_forward(te_grouped, batch_size=8) + + sequential = _gpt_model_provider( + tp_size=tp_size, + ep_size=ep_size, + hidden_size=32, + moe_grouped_gemm=False, + num_moe_experts=4, + transformer_impl="modelopt", + ) + copy_weights_from_grouped_to_non_grouped(te_grouped, sequential) + + for module in te_grouped.modules(): + if isinstance(module, TopKRouter): + module.topk = module.num_experts + for module in sequential.modules(): + if isinstance(module, TopKRouter): + module.topk = module.num_experts + + ref_te = forward(te_grouped) + ref_seq = forward(sequential) + + mtq.quantize(te_grouped, quant_cfg, forward) + mtq.quantize(sequential, quant_cfg, forward) + + out_te = forward(te_grouped) + out_seq = forward(sequential) + + err_te = (out_te - ref_te).abs().mean().item() + err_seq = (out_seq - ref_seq).abs().mean().item() + + if rank == 0: + print( + f"\n[default-amax] TEGrouped quant-err={err_te:.6f}, " + f"Sequential quant-err={err_seq:.6f}, ratio TE/Seq={err_te / max(err_seq, 1e-12):.3f}" + ) + + # At toy scale (4 small experts) the per-tensor amax difference is dominated + # by other numerical noise (~few %); the effect amplifies at production scale + # (e.g. 128 experts in Nemotron Nano). Just sanity-check both errors are finite. + assert err_te > 0 and err_seq > 0 + assert math.isfinite(err_te) and math.isfinite(err_seq) + + +@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG]) +def test_te_grouped_vs_sequential_default_loss(dist_workers_size_4, quant_cfg): + dist_workers_size_4.run( + partial(_test_te_grouped_vs_sequential_default_loss_helper, 1, 2, quant_cfg) + ) + + def _test_auto_quantize_moe_ep_helper(rank, size): initialize_for_megatron( tensor_model_parallel_size=1, diff --git a/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py b/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py index 56019d8cf75..46ebc2cae50 100644 --- a/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quantizer_cpu.py @@ -15,12 +15,15 @@ """Tests of tensor quantizer.""" +import torch from _test_utils.torch.quantization.tensor_quantizer_common import ( BlockQuantTester, SequentialQuantizerTester, TensorQuantizerTester, ) +from modelopt.torch.quantization.nn import GroupedQuantizer, TensorQuantizer + class TestTensorQuantizerCPU(TensorQuantizerTester): device = "cpu" @@ -32,3 +35,15 @@ class TestBlockQuantCPU(BlockQuantTester): class TestSequentialQuantizerCPU(SequentialQuantizerTester): device = "cpu" + + +def test_grouped_quantizer_forward_uses_representative_quantizer(): + """Single-weight compatibility paths should dispatch to the first group.""" + representative = TensorQuantizer() + other = TensorQuantizer() + other.disable() + grouped = GroupedQuantizer(representative, other) + inputs = torch.tensor([0.1234, -0.5678]) + + assert torch.equal(grouped(inputs), representative(inputs)) + assert not torch.equal(grouped(inputs), other(inputs))