From 87bab303d8ff6405f4fbfae6426e37f030f7926e Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Wed, 22 Jul 2026 09:26:53 -0700 Subject: [PATCH 01/16] feat(quant): per-expert weight quantizer for TEGroupedMLP Give each fused expert in a TEGroupedLinear its own weight quantizer via a GroupedQuantizer (an nn.ModuleList surfaced as weight_quantizer.{i}), so per-expert amax is independent instead of expert-0's shared across all. Includes the amax-preserving restore fix (Issue 1): modelopt_post_restore keeps the loaded MSE/static-calibrated (and QAD-frozen) amax and only re-max_calibrates a quantizer whose loaded amax is shape-incompatible with its weight (a genuine TP/EP change), detected via a fake-quant dry-run. Adds an opt-in torch.compile path for the per-expert quantize loop (MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1; default stays eager). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jennifer Chen --- modelopt/torch/quantization/config.py | 4 +- modelopt/torch/quantization/model_calib.py | 19 +++-- .../nn/modules/tensor_quantizer.py | 69 ++++++++++++++++ modelopt/torch/quantization/plugins/custom.py | 20 ++++- .../plugins/transformer_engine.py | 82 +++++++++++++++++-- .../quantization/test_tensor_quantizer_cpu.py | 15 ++++ 6 files changed, 192 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 6257623d108..1e70ec62aab 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -963,8 +963,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 3fe38610a74..d2ab0323a25 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -43,7 +43,14 @@ 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 QuantModule, SequentialQuantizer, StaticBlockScaleQuantizer, TensorQuantizer +from .nn import ( + GroupedQuantizer, + NVFP4StaticQuantizer, + QuantModule, + SequentialQuantizer, + StaticBlockScaleQuantizer, + TensorQuantizer, +) from .utils import ( SHARED_PATTERNS, SharedWeightGlobalAmaxState, @@ -208,7 +215,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 @@ -376,12 +383,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 @@ -395,7 +402,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 @@ -417,7 +424,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 0e313dd97e1..2db9c0d98e7 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -71,6 +71,7 @@ _FP8_E4M3_MIN_POSITIVE = torch.finfo(torch.float8_e4m3fn).smallest_normal / (2**3) __all__ = [ + "GroupedQuantizer", "HardDisabledTensorQuantizer", "NVFP4StaticQuantizer", "SequentialQuantizer", @@ -1840,3 +1841,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/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index d0efcc52db1..57affbe3ac1 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,41 @@ def modelopt_post_restore(self, prefix: str = ""): # Remove self.weight after post_restore. delattr(self, "weight") + # Preserve the loaded (calibrated) per-expert amax. Recomputing via max_calibrate + # replaces MSE/static-calibrated (and QAD-frozen) amax with max|W|, which corrupts + # static recipes on export and overwrites frozen amax on every QAD resume. Only + # re-calibrate a quantizer whose loaded amax is shape-INCOMPATIBLE with its weight + # (a genuine TP/EP change between save and restore); otherwise keep it as-is. + 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") or q._amax is None: + continue + prev_fake = getattr(q, "_fake_quant", True) + q._fake_quant = True + try: + wq_i(weight_i) # dry-run: succeeds iff the loaded amax fits this weight + shape_ok = True + except Exception: + shape_ok = False + finally: + q._fake_quant = prev_fake + if shape_ok: + continue # loaded amax is valid -> keep it, do NOT recompute + 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 +243,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/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)) From b6659bcc4a6ba917f608fc1e71267fde20569dbf Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Wed, 22 Jul 2026 09:26:53 -0700 Subject: [PATCH 02/16] fix(quant): shard TE grouped per-expert amax with global expert identity _QuantMegatronTEGroupedLinear inherited the base sharded_state_dict, which emitted per-expert amax under local keys weight_quantizer.{0..num_local-1} with no expert offset, so every EP rank wrote identical keys and torch_dist dedup kept only one rank's experts (EP16: 128 -> 8). Override it to emit each per-expert amax with global_expert_idx = ep_rank*num_gemms + gemm_idx (mirroring MCore _sharded_state_dict_grouped), so all num_global_experts persist and reshard at any EP. Fixes both the collapsed save and the EP>1 first-load mis-map that corrupted static-NVFP4 QAD. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jennifer Chen --- .../torch/quantization/plugins/megatron.py | 149 +++++++++++++++++- 1 file changed, 141 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 752dd801a6e..ea4ddcd3572 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 @@ -24,12 +25,13 @@ import megatron.core.transformer.mlp as megatron_mlp import megatron.core.transformer.moe.experts as megatron_moe import torch +from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding from megatron.core.parallel_state import get_data_parallel_group from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region from megatron.core.transformer import MegatronModule from megatron.core.transformer.attention import Attention from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint -from megatron.core.utils import get_tensor_model_parallel_group_if_none +from megatron.core.utils import get_pg_rank, get_pg_size, get_tensor_model_parallel_group_if_none from modelopt.torch.opt.dynamic import DynamicModule from modelopt.torch.opt.plugins.megatron import ( @@ -42,7 +44,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 +98,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): @@ -224,7 +232,17 @@ def quant_module_set_extra_state(self, state: Any): if quantizer_state is not None: for name, module in self.named_modules(): if isinstance(module, TensorQuantizer): - quantizer_substate = quantizer_state[name] + quantizer_substate = quantizer_state.get(name) + if quantizer_substate is None: + # Per-expert quantizers ("weight_quantizer.") are saved per EP rank, so a + # module loaded at smaller EP (e.g. EP1 export from an EP16 ckpt) has more + # experts than the saved state. Per-expert properties are uniform across + # experts (amax rides separately as globally-indexed sharded tensors), so + # fall back to expert 0's state. + fallback = re.sub(r"\.\d+$", ".0", name) + quantizer_substate = quantizer_state.get(fallback) + if quantizer_substate is None: + continue maybe_promote_nvfp4_static_quantizer(module, quantizer_substate) module.set_from_modelopt_state(quantizer_substate, properties_only=False) self.modelopt_post_restore() @@ -438,7 +456,8 @@ def _get_shard_axis_dict(self, state_dict): """ shard_axis_dict = {} for k in state_dict: - # Static NVFP4 _global_amax is a replicated scalar; only per-block _amax shards. + # _global_amax needs no channel shard axis (replicated scalar; for grouped experts it + # rides with the global expert identity assigned in the grouped sharded_state_dict). if k.endswith("_global_amax"): continue if "weight_quantizer." in k: @@ -469,7 +488,8 @@ def _get_shard_axis_dict(self, state_dict): """ shard_axis_dict = {} for k in state_dict: - # Static NVFP4 _global_amax is a replicated scalar; only per-block _amax shards. + # _global_amax needs no channel shard axis (replicated scalar; for grouped experts it + # rides with the global expert identity assigned in the grouped sharded_state_dict). if k.endswith("_global_amax"): continue if "weight_quantizer." in k: @@ -697,8 +717,121 @@ 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 + + def _expert_parallel_groups(self): + """Return the (ep, expt_dp) process groups used to place fused experts globally.""" + pg_collection = getattr(self, "_pg_collection", None) + if pg_collection is not None: + return pg_collection.ep, pg_collection.expt_dp + return ( + mcore_parallel.get_expert_model_parallel_group(), + mcore_parallel.get_expert_data_parallel_group(), + ) + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + """Emit per-expert quantizer amax with the same global expert identity as the weights. + + The base linear emits ``weight_quantizer.{local_i}._amax`` with the local index and no + expert offset, so every EP rank writes identical keys and ``torch_dist`` dedup keeps only + one rank's experts. Here we mirror Megatron ``TEGroupedLinear._sharded_state_dict_grouped``: + each fused expert comes with its ``global_expert_idx`` (baked into the key prefix under + ``singleton_local_shards``, otherwise an EP sharded-offset) so all ``num_global_experts`` + persist and reshard to any EP. Shared, whole-linear quantizer buffers (e.g. + ``input_quantizer``) keep the plain replicated path. + """ + metadata = ensure_metadata_has_dp_cp_group(metadata) + singleton_local_shards = bool((metadata or {}).get("singleton_local_shards", False)) + + # Weights/bias/_extra_state come from the wrapped TE grouped linear, which already + # assigns each expert its global identity. Skip _MegatronParallelLinear's local-index + # amax emission by starting from the base MCore module's sharded_state_dict. + sharded_state_dict = super(_MegatronParallelLinear, self).sharded_state_dict( + prefix, sharded_offsets, metadata + ) + + # Collect the quantizer buffers exactly like _MegatronParallelLinear.sharded_state_dict. + quantizer_state_dict = {} + for k, v in self.state_dict(prefix="", keep_vars=True).items(): + if "_quantizer" in k and "_amax" in k: + self._process_quantizer_amax(k, v, quantizer_state_dict) + elif k == "input_quantizer._pre_quant_scale": + self._process_activation_quantizer_pre_quant_scale(k, v, quantizer_state_dict) + elif self._parameter_to_keep_in_quantizer_state_dict(k): + quantizer_state_dict[k] = v + elif "quantizer" in k: + warn_rank_0( + f"Quantizer state {k} is not supported for sharded_state_dict. " + "Please use regular state_dict." + ) + + # Channel shard axes (per real key); _global_amax stays un-sharded along channels but + # still rides with the expert identity below. + shard_axis_dict = self._get_shard_axis_dict(quantizer_state_dict) + + # Split per-expert weight_quantizer.{i}.* from shared (input/output) quantizer buffers. + expert_re = re.compile(r"^weight_quantizer\.(\d+)\.(.+)$") + per_expert_subs = [[] for _ in range(self.num_gemms)] + shared_state = {} + for k, v in quantizer_state_dict.items(): + m = expert_re.match(k) + if m: + per_expert_subs[int(m.group(1))].append((m.group(2), v, shard_axis_dict.get(k))) + else: + shared_state[k] = v + + # Shared quantizer buffers: replicated across experts, plain base offsets. + shared_axis_dict = {k: shard_axis_dict[k] for k in shared_state if k in shard_axis_dict} + sharded_state_dict.update( + make_sharded_tensors_for_checkpoint( + shared_state, prefix, shared_axis_dict, sharded_offsets + ) + ) + + # Per-expert amax: assign the same global expert identity the weights use. + ep_group, expt_dp_group = self._expert_parallel_groups() + num_global_experts = get_pg_size(ep_group) * self.num_gemms + local_expert_indices_offset = get_pg_rank(ep_group) * self.num_gemms + edp_replica_id = get_pg_rank(expt_dp_group) + ep_axis = len(sharded_offsets) + for gemm_idx, subs in enumerate(per_expert_subs): + if not subs: + continue + global_expert_idx = local_expert_indices_offset + gemm_idx + if singleton_local_shards: + expert_prefix = f"{global_expert_idx}.{prefix}" + new_sharded_offsets = sharded_offsets + else: + expert_prefix = prefix + new_sharded_offsets = ( + *sharded_offsets, + (ep_axis, global_expert_idx, num_global_experts), + ) + expert_state = {f"{gemm_idx}.weight_quantizer.{sub}": v for sub, v, _ in subs} + expert_axis = { + f"{gemm_idx}.weight_quantizer.{sub}": axis + for sub, _, axis in subs + if axis is not None + } + sub_sd = make_sharded_tensors_for_checkpoint( + expert_state, "", expert_axis, new_sharded_offsets + ) + # Rewrite each ShardedTensor.key to carry the global expert identity (dict keys, + # which map to the local buffers on restore, are left untouched). + replace_prefix_for_sharding(sub_sd, f"{gemm_idx}.", expert_prefix) + for sub, _, _ in subs: + sh_ten = sub_sd[f"{gemm_idx}.weight_quantizer.{sub}"] + replica_id = sh_ten.replica_id + if len(replica_id) == 3: + sh_ten.replica_id = (*replica_id[:2], edp_replica_id) + sharded_state_dict[f"{prefix}weight_quantizer.{gemm_idx}.{sub}"] = sh_ten + return sharded_state_dict @QuantModuleRegistry.register( {TEColumnParallelGroupedLinear: "megatron_TEColumnParallelGroupedLinear"} From d2b46a4e9a69798482611c265b3e7f4b0d18ae95 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Wed, 22 Jul 2026 09:26:53 -0700 Subject: [PATCH 03/16] feat(export): per-expert HF export for TEGroupedMLP Export each fused expert with its own qformat/scales by swapping in that expert's weight{i} and TensorQuantizer, instead of applying weight0's scales to every expert. Combined with the EP-gather path (local_expert_indices + all_gather_object across the EP group, collective-safe missing-key check) so EP>1 exports gather all global experts; EP=1 reduces to the plain per-expert loop. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jennifer Chen --- .../torch/export/unified_export_megatron.py | 131 ++++++++++-------- 1 file changed, 74 insertions(+), 57 deletions(-) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 0e443391f39..a6ab24447ed 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -70,6 +70,7 @@ process_layer_quant_config, to_quantized_weight, ) +from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer with import_plugin("transformers", verbose=False): import transformers @@ -1053,24 +1054,11 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): Reverse of _grouped_mlp_merging in the importer. """ num_experts = module.num_gemms + state_dict = module.state_dict() - # TEGroupedLinear doesn't have module.weight (it has weight0, weight1, ...). - # Temporarily assign weight = weight0 so _get_quantized_state can extract - # qformat, scales, and input_scale from the module's quantizers. has_weight = hasattr(module, "weight") - if not has_weight: - module.weight = module.weight0 - try: - name_to_value, qformat, block_size = self._get_quantized_state( - module, self.dtype, prefix=prefix - ) - weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) - name_to_value.pop("weight", None) - finally: - if not has_weight and hasattr(module, "weight"): - delattr(module, "weight") - - state_dict = module.state_dict() + grouped_wq = getattr(module, "weight_quantizer", None) + per_expert_wq = isinstance(grouped_wq, GroupedQuantizer) ep_size = ( get_expert_model_parallel_world_size() if torch.distributed.is_initialized() else 1 @@ -1119,48 +1107,78 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): elif local_missing: raise ValueError(f"TEGroupedMLP missing expert weights: {local_missing}") - # Move shared scales/aux to CPU once so the gather payload avoids GPU clones. - weight_scale_cpu = weight_scale.detach().cpu().clone() if weight_scale is not None else None - weight_scale_2_cpu = ( - weight_scale_2.detach().cpu().clone() if weight_scale_2 is not None else None - ) - name_to_value_cpu = { - k: v.detach().cpu().clone() for k, v in name_to_value.items() if k != "output_scale" - } - - # Record quant config for ALL global experts on every rank; otherwise the writer's - # hf_quant_config.json would miss (EP-1)/EP of the routed experts. All experts in - # a TEGroupedMLP layer share qformat/block_size, so local values apply globally. - num_total_experts = num_experts * ep_size - for global_id in range(num_total_experts): - self._record_layer_quant_config(prefix.format(global_id) + ".", qformat, block_size) - + # Per expert, temporarily assign weight = weight{i} and, for the per-expert + # quantizer layout (GroupedQuantizer), swap in that expert's own TensorQuantizer, + # so _get_quantized_state extracts each expert's own qformat/scales instead of + # applying weight0's scales to every expert. local_expert_state: dict[str, torch.Tensor] = {} + seen_qformat = None + seen_block_size = None + try: + for local_id in range(num_experts): + global_id = local_expert_indices[local_id] + expert_prefix = prefix.format(global_id) + "." + weight_key = f"weight{local_id}" + + module.weight = getattr(module, weight_key) + if per_expert_wq: + module.weight_quantizer = grouped_wq[min(local_id, len(grouped_wq) - 1)] + # Dynamic-NVFP4 per-expert quantizers carry no stored amax, but + # weight_scale_2 derivation asserts one. Max-calibration weight amax + # is exactly max(|W|), so compute it from this expert's weight. + _wq = module.weight_quantizer + if getattr(_wq, "_amax", None) is None and getattr(_wq, "is_enabled", False): + _wq.amax = module.weight.detach().abs().max().float() + + name_to_value, qformat, block_size = self._get_quantized_state( + module, self.dtype, prefix=prefix + ) + weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) + name_to_value.pop("weight", None) + seen_qformat, seen_block_size = qformat, block_size - for local_id in range(num_experts): - global_id = local_expert_indices[local_id] - expert_prefix = prefix.format(global_id) + "." - weight_key = f"weight{local_id}" + weight = state_dict[weight_key].to(self.dtype).cpu() + weight_scale_cpu = ( + weight_scale.detach().cpu().clone() if weight_scale is not None else None + ) + weight_scale_2_cpu = ( + weight_scale_2.detach().cpu().clone() if weight_scale_2 is not None else None + ) - weight = state_dict[weight_key].to(self.dtype).cpu() + if weight_scale_cpu is None: + local_expert_state[expert_prefix + "weight"] = weight + else: + local_expert_state[expert_prefix + "weight"] = to_quantized_weight( + weight, + weight_scale_cpu, + qformat, + weight_scale_2_cpu, + block_size, + ) + local_expert_state[expert_prefix + "weight_scale"] = weight_scale_cpu.clone() - if weight_scale_cpu is None: - local_expert_state[expert_prefix + "weight"] = weight - else: - local_expert_state[expert_prefix + "weight"] = to_quantized_weight( - weight, - weight_scale_cpu, - qformat, - weight_scale_2_cpu, - block_size, - ) - local_expert_state[expert_prefix + "weight_scale"] = weight_scale_cpu.clone() + if weight_scale_2_cpu is not None: + local_expert_state[expert_prefix + "weight_scale_2"] = weight_scale_2_cpu.clone() - if weight_scale_2_cpu is not None: - local_expert_state[expert_prefix + "weight_scale_2"] = weight_scale_2_cpu.clone() + for key, val in name_to_value.items(): + if key == "output_scale": + continue + local_expert_state[expert_prefix + key] = val.detach().cpu().clone() + finally: + if per_expert_wq: + module.weight_quantizer = grouped_wq + if not has_weight and hasattr(module, "weight"): + delattr(module, "weight") - for key, val in name_to_value_cpu.items(): - local_expert_state[expert_prefix + key] = val.clone() + # Record quant config for ALL global experts on every rank; otherwise the writer's + # hf_quant_config.json would miss (EP-1)/EP of the routed experts. All experts in + # a TEGroupedMLP layer share qformat/block_size, so local values apply globally. + if seen_qformat is not None: + num_total_experts = num_experts * ep_size + for global_id in range(num_total_experts): + self._record_layer_quant_config( + prefix.format(global_id) + ".", seen_qformat, seen_block_size + ) if ep_size > 1: # all_gather_object pickles trip on quantized uint8 tensors whose @@ -1175,11 +1193,10 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): ) del local_bytes for b in gathered_bytes: - # weights_only=False: bytes are our own torch.save output from a sibling - # EP rank in this job's collective, not user-supplied. weights_only=True - # rejects quantized uint8 tensors (custom storage outside the allowlist). - s = torch.load(io.BytesIO(b), map_location="cpu", weights_only=False) - self._state_dict.update(s) + # weights_only=False: our own torch.save output from a sibling EP rank + # in this job's collective, not user-supplied. + s_loaded = torch.load(io.BytesIO(b), map_location="cpu", weights_only=False) + self._state_dict.update(s_loaded) del gathered_bytes else: self._state_dict.update(local_expert_state) From e2742fc163e0b81d052698a30560c581808cb595 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Wed, 22 Jul 2026 09:34:57 -0700 Subject: [PATCH 04/16] test(quant): TEGrouped per-expert quantization tests Add GPU tests asserting per-expert amax independence, sharded_state_dict global expert identity, and the opt-in compile path; plus the changelog entry. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jennifer Chen --- .../quantization/plugins/test_megatron.py | 341 +++++++++++++++++- 1 file changed, 336 insertions(+), 5 deletions(-) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 36f80787931..ea3698ee23a 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -14,6 +14,8 @@ # limitations under the License. import copy +import math +import re from contextlib import nullcontext from functools import partial from pathlib import Path @@ -56,11 +58,15 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.quantization.algorithms import QuantRecipe, _AutoQuantizeBaseSearcher -from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer 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 @@ -737,10 +743,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) @@ -760,6 +766,331 @@ 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 _te_grouped_expert_identity_from_sharded_state(module): + """Return {local_key: (global_expert_idx, num_global_experts)} for per-expert amax shards. + + The grouped linear must give each fused expert the same global identity the weights use: + the dict key keeps the local expert index (maps to the local buffer on restore) while the + ShardedTensor carries the global expert offset. Called with sharded_offsets=() so the expert + axis is the (only) prepended axis at index 0. + """ + sharded_sd = module.sharded_state_dict(prefix="", sharded_offsets=(), metadata=None) + identity = {} + for key, sh_ten in sharded_sd.items(): + if re.match(r"weight_quantizer\.\d+\..*_amax$", key): + assert sh_ten.prepend_axis_num >= 1, f"{key}: expected a prepended expert axis" + identity[key] = (int(sh_ten.global_offset[0]), int(sh_ten.global_shape[0])) + return identity + + +def _test_te_grouped_sharded_state_dict_global_expert_identity_helper( + tp_size, ep_size, quant_cfg, rank, size +): + """Per-expert quantizer amax must persist all num_global_experts across EP. + + With EP>1 the base linear emitted ``weight_quantizer.{local_i}._amax`` at the local index with + no expert offset, so every rank wrote identical keys and torch_dist dedup collapsed them to a + single rank's experts. Assert each rank's fused experts now carry distinct global identities so + the union across ranks covers every global expert. + """ + initialize_for_megatron( + tensor_model_parallel_size=tp_size, + expert_model_parallel_size=ep_size, + seed=SEED, + ) + num_experts = 4 + num_local = num_experts // ep_size + + 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=num_experts, + ) + forward = get_forward(te_grouped, batch_size=8) + for module in te_grouped.modules(): + if isinstance(module, TopKRouter): + module.topk = module.num_experts + mtq.quantize(te_grouped, quant_cfg, forward) + + grouped_linears = [ + m for m in te_grouped.modules() if isinstance(m, _QuantMegatronTEGroupedLinear) + ] + assert grouped_linears, "No grouped quant linears found" + + expected_global = {rank * num_local + i for i in range(num_local)} + for linear in grouped_linears: + # Give each expert a distinct amax so a value mix-up would also be observable. + for i in range(linear.num_gemms): + wq = linear.weight_quantizer[i] + leaves = list(wq) if isinstance(wq, SequentialQuantizer) else [wq] + for leaf in leaves: + if hasattr(leaf, "_amax") and leaf._amax is not None: + leaf._amax.fill_(1.0 + rank * num_local + i) + + identity = _te_grouped_expert_identity_from_sharded_state(linear) + # One entry per local expert per amax buffer; dict keys keep the LOCAL index. + local_keys = {int(re.search(r"weight_quantizer\.(\d+)\.", k).group(1)) for k in identity} + assert local_keys == set(range(num_local)), ( + f"Expected local expert keys {set(range(num_local))}, got {local_keys}" + ) + # ShardedTensor global identity: this rank owns experts {rank*num_local + i}. + local_global = {gidx for gidx, _ in identity.values()} + assert local_global == expected_global, ( + f"rank {rank}: expected global experts {expected_global}, got {local_global}" + ) + assert all(total == num_experts for _, total in identity.values()), ( + f"num_global_experts should be {num_experts}, got {identity}" + ) + + # Gather the global expert indices across all EP ranks: the union must cover every expert. + gathered = [None] * size + torch.distributed.all_gather_object(gathered, sorted(expected_global)) + union = set() + for part in gathered: + union.update(part) + assert union == set(range(num_experts)), ( + f"Union of global experts across EP ranks should be {set(range(num_experts))}, got {union}" + ) + + +@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG]) +def test_te_grouped_sharded_state_dict_global_expert_identity(dist_workers_size_2, quant_cfg): + dist_workers_size_2.run( + partial(_test_te_grouped_sharded_state_dict_global_expert_identity_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, From 2883da00750ab5fce3d18ea05d039170bd968422 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Wed, 22 Jul 2026 11:52:44 -0700 Subject: [PATCH 05/16] refactor(quant): alias the any-quantizer isinstance tuple in model_calib Introduce module-level _ANY_QUANTIZER = (TensorQuantizer, SequentialQuantizer, GroupedQuantizer) and use it at the two "is this child any quantizer" sites, instead of repeating the three-way isinstance union. No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jennifer Chen --- modelopt/torch/quantization/model_calib.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index d2ab0323a25..fe06fe0e426 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -77,6 +77,9 @@ "svdquant", ] +# Any quantizer instance: the leaf TensorQuantizer or a quantizer container. +_ANY_QUANTIZER = (TensorQuantizer, SequentialQuantizer, GroupedQuantizer) + def _collect_weight_stats(quantizer: nn.Module, weight: torch.Tensor) -> None: quantizer(weight) @@ -383,7 +386,7 @@ 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 | GroupedQuantizer): + if isinstance(child, _ANY_QUANTIZER): _check_moe_calibration_complete(child, module.parallel_state) def sync_quantizer_amax_across_dp_ep(quantizer, parallel_state, parent_name, child_name): @@ -402,7 +405,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 | GroupedQuantizer): + if isinstance(child, _ANY_QUANTIZER): 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 From 5b8f9d7d18fe7799db36d54adf4f9adaee9afa8d Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Fri, 24 Jul 2026 11:43:39 -0700 Subject: [PATCH 06/16] refactor(quant): centralize quantizer type checks on AnyQuantizer Introduce AnyQuantizer = (TensorQuantizer, SequentialQuantizer, GroupedQuantizer) in quantization.nn and route the leaf/container isinstance checks through it, replacing the local _ANY_QUANTIZER tuple and the (TensorQuantizer, SequentialQuantizer)-only checks that skipped GroupedQuantizer: - model_calib: MoE amax completeness + DP/EP amax sync - vllm_fakequant_hf: weight-quantizer-disabled export guard Also: - representative_weight_quantizer: handle a singular GroupedQuantizer (TEGroupedLinear fused experts) by returning its first expert, so the exported hf_quant_config qformat/scales are correct instead of missing. - Drop GroupedQuantizer.get_modelopt_state; the container needs no serialized meta state. - Fix mypy in the per-expert export path (narrowing asserts) and drop the now-unused NVFP4StaticQuantizer import. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jennifer Chen --- modelopt/torch/export/plugins/vllm_fakequant_hf.py | 11 +++++++---- modelopt/torch/export/unified_export_megatron.py | 8 ++++++-- modelopt/torch/quantization/model_calib.py | 9 +++------ modelopt/torch/quantization/nn/__init__.py | 3 +++ .../torch/quantization/nn/modules/tensor_quantizer.py | 4 ---- modelopt/torch/quantization/utils/core_utils.py | 8 ++++++-- 6 files changed, 25 insertions(+), 18 deletions(-) diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py index acb1968e070..9168f988146 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py @@ -29,7 +29,12 @@ import modelopt.torch.opt as mto from modelopt.torch.quantization.conversion import quantizer_state from modelopt.torch.quantization.model_calib import enable_stats_collection, finish_stats_collection -from modelopt.torch.quantization.nn import QuantModule, SequentialQuantizer, TensorQuantizer +from modelopt.torch.quantization.nn import ( + AnyQuantizer, + QuantModule, + SequentialQuantizer, + TensorQuantizer, +) from modelopt.torch.quantization.utils import get_quantizer_state_dict from modelopt.torch.quantization.utils.core_utils import enable_weight_access_and_writeback from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -125,9 +130,7 @@ def _check_all_weight_quantizers_disabled(model: nn.Module) -> None: if not isinstance(module, QuantModule): continue for attr_name, quantizer in module.named_children(): - if attr_name.endswith("weight_quantizer") and isinstance( - quantizer, (TensorQuantizer, SequentialQuantizer) - ): + if attr_name.endswith("weight_quantizer") and isinstance(quantizer, AnyQuantizer): if quantizer.is_enabled: raise RuntimeError( f"vLLM fakequant export: {attr_name!r} must be disabled before saving " diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index a6ab24447ed..813638390d0 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -34,6 +34,7 @@ from safetensors.torch import save_file from modelopt import __version__ +from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer from modelopt.torch.utils import import_plugin from .convert_hf_config import convert_hf_quant_config_format @@ -70,7 +71,6 @@ process_layer_quant_config, to_quantized_weight, ) -from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer with import_plugin("transformers", verbose=False): import transformers @@ -1122,6 +1122,7 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): module.weight = getattr(module, weight_key) if per_expert_wq: + assert isinstance(grouped_wq, GroupedQuantizer) module.weight_quantizer = grouped_wq[min(local_id, len(grouped_wq) - 1)] # Dynamic-NVFP4 per-expert quantizers carry no stored amax, but # weight_scale_2 derivation asserts one. Max-calibration weight amax @@ -1158,7 +1159,9 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): local_expert_state[expert_prefix + "weight_scale"] = weight_scale_cpu.clone() if weight_scale_2_cpu is not None: - local_expert_state[expert_prefix + "weight_scale_2"] = weight_scale_2_cpu.clone() + local_expert_state[expert_prefix + "weight_scale_2"] = ( + weight_scale_2_cpu.clone() + ) for key, val in name_to_value.items(): if key == "output_scale": @@ -1174,6 +1177,7 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): # hf_quant_config.json would miss (EP-1)/EP of the routed experts. All experts in # a TEGroupedMLP layer share qformat/block_size, so local values apply globally. if seen_qformat is not None: + assert seen_block_size is not None num_total_experts = num_experts * ep_size for global_id in range(num_total_experts): self._record_layer_quant_config( diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index fe06fe0e426..dfaf1ac59fa 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -44,8 +44,8 @@ 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 ( + AnyQuantizer, GroupedQuantizer, - NVFP4StaticQuantizer, QuantModule, SequentialQuantizer, StaticBlockScaleQuantizer, @@ -77,9 +77,6 @@ "svdquant", ] -# Any quantizer instance: the leaf TensorQuantizer or a quantizer container. -_ANY_QUANTIZER = (TensorQuantizer, SequentialQuantizer, GroupedQuantizer) - def _collect_weight_stats(quantizer: nn.Module, weight: torch.Tensor) -> None: quantizer(weight) @@ -386,7 +383,7 @@ 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, _ANY_QUANTIZER): + if isinstance(child, AnyQuantizer): _check_moe_calibration_complete(child, module.parallel_state) def sync_quantizer_amax_across_dp_ep(quantizer, parallel_state, parent_name, child_name): @@ -405,7 +402,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, _ANY_QUANTIZER): + if isinstance(child, AnyQuantizer): 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 diff --git a/modelopt/torch/quantization/nn/__init__.py b/modelopt/torch/quantization/nn/__init__.py index 2e6bc64054e..794f3b70cc2 100644 --- a/modelopt/torch/quantization/nn/__init__.py +++ b/modelopt/torch/quantization/nn/__init__.py @@ -26,3 +26,6 @@ from .modules.quant_pooling import * from .modules.quant_rnn import * from .modules.tensor_quantizer import * + +# Every quantizer instance type: the leaf TensorQuantizer or a quantizer container. +AnyQuantizer = (TensorQuantizer, SequentialQuantizer, GroupedQuantizer) # noqa: F405 diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 2db9c0d98e7..cf8537882e8 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -1905,7 +1905,3 @@ def set_from_attribute_config(self, attributes): 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/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 1bdf23da64a..8898bb0d42d 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -214,21 +214,25 @@ def reduce_sum(input, axis=None, keepdims=True): def representative_weight_quantizer(module: nn.Module, weight_name: str = "weight"): """Return the representative weight quantizer for ``weight_name`` on ``module``. - Handles two layouts: + Handles three layouts: - singular ``_weight_quantizer`` — standard ``nn.Linear`` / ``_QuantLinear``. + - singular ``_weight_quantizer`` that is a ``GroupedQuantizer`` — TEGroupedLinear + fused experts (one quantizer per expert); the first is representative. - plural ``_weight_quantizers`` (``nn.ModuleList``) — fused-experts modules (``_QuantFusedExperts``) hold one ``TensorQuantizer`` per expert. Per-expert formats are identical, so the first element is representative. Returns ``None`` if no matching quantizer is found. """ - from ..nn import SequentialQuantizer, TensorQuantizer + from ..nn import GroupedQuantizer, SequentialQuantizer, TensorQuantizer singular = quantizer_attr_names(weight_name).weight_quantizer q = getattr(module, singular, None) if isinstance(q, (TensorQuantizer, SequentialQuantizer)): return q + if isinstance(q, GroupedQuantizer) and len(q) > 0: + return q[0] plural = getattr(module, singular + "s", None) if isinstance(plural, nn.ModuleList) and len(plural) > 0: From 4708d8534da2d1d92c0d93378686cb5b0f214156 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Tue, 28 Jul 2026 12:09:15 -0700 Subject: [PATCH 07/16] make tests more robust Signed-off-by: Jennifer Chen --- .../quantization/plugins/test_megatron.py | 139 ++++++++++++++++-- 1 file changed, 126 insertions(+), 13 deletions(-) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index ea3698ee23a..bfaa9fd7974 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -848,9 +848,96 @@ def compiled(*args): destroy_model_parallel() +def test_te_grouped_real_compile_weight_quantizer_loop(distributed_setup_size_1, monkeypatch): + """Real (unpatched) torch.compile parity for the per-expert weight-quantizer loop. + + Complements test_te_grouped_compiled_weight_quantizer_loop, which fakes torch.compile to + assert wiring only. Here torch.compile is left intact so the opt-in loop is actually + compiled, executed, and back-propagated, and its numerics are checked against the eager + path built from identical weights and calibrated amax. + """ + initialize_for_megatron(seed=SEED) + + def build(): + model = _gpt_model_provider( + tp_size=1, + hidden_size=32, + moe_grouped_gemm=True, + transformer_impl="transformer_engine", + num_moe_experts=4, + ) + for module in model.modules(): + if isinstance(module, TopKRouter): + module.topk = module.num_experts + return model + + # Two identical models (same raw weights); one stays eager, one is real-compiled. + model_eager = build() + model_compiled = build() + model_compiled.load_state_dict(model_eager.state_dict()) + + # One cached input batch, shared across both models for an apples-to-apples compare. + forward = get_forward(model_eager) + + monkeypatch.delenv(_COMPILE_TEGROUPED_WEIGHT_LOOP_ENV, raising=False) + mtq.quantize(model_eager, copy.deepcopy(mtq.INT8_DEFAULT_CFG), forward) + + monkeypatch.setenv(_COMPILE_TEGROUPED_WEIGHT_LOOP_ENV, "1") + mtq.quantize(model_compiled, copy.deepcopy(mtq.INT8_DEFAULT_CFG), forward) + + grouped_eager = [ + m + for m in model_eager.modules() + if isinstance(getattr(m, "weight_quantizer", None), mtq.nn.GroupedQuantizer) + ] + grouped_compiled = [ + m + for m in model_compiled.modules() + if isinstance(getattr(m, "weight_quantizer", None), mtq.nn.GroupedQuantizer) + ] + assert grouped_compiled and len(grouped_eager) == len(grouped_compiled) + # The opt-in path attached the real compiled loop (torch.compile left unpatched); the + # eager control model did not. + assert all(hasattr(m, "_compiled_weight_quantizer_loop") for m in grouped_compiled) + assert all(not hasattr(m, "_compiled_weight_quantizer_loop") for m in grouped_eager) + + # Forward parity: the first call on model_compiled triggers real compilation. + out_eager = forward(model_eager) + out_compiled = forward(model_compiled) + torch.testing.assert_close(out_compiled, out_eager, rtol=1e-3, atol=1e-3) + + # Backward parity: per-expert weight grads must be finite and match the eager path. + out_eager.sum().backward() + out_compiled.sum().backward() + for m_e, m_c in zip(grouped_eager, grouped_compiled): + for i in range(m_c.num_gemms): + g_e = getattr(m_e, f"weight{i}").grad + g_c = getattr(m_c, f"weight{i}").grad + assert g_c is not None and torch.isfinite(g_c).all() + torch.testing.assert_close(g_c, g_e, rtol=1e-2, atol=1e-2) + + destroy_model_parallel() + + +def _te_grouped_expert_magnitude(linear_name, local_idx): + """Distinct, known weight magnitude for each (linear, local-expert) pair. + + Chosen so every per-expert weight quantizer sees a different amax (and fc1 vs fc2 differ + too), making divergence guaranteed by construction rather than by random initialization. + """ + return {"linear_fc1": 0.25, "linear_fc2": 1.25}[linear_name] + 0.5 * local_idx + + 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).""" + """TEGrouped keeps a per-expert weight quantizer (GroupedQuantizer) by default; each expert's + amax must equal the corresponding SequentialMLP expert's (no cross-expert sharing). + + Divergence is made causal: each local expert's weights are filled with a distinct known + magnitude, so its weight amax is that magnitude by construction. The test then asserts + (a) grouped == sequential per expert, (b) each amax equals ITS OWN set magnitude, and + (c) the per-expert quantizer objects are distinct instances. A cross-expert-sharing + regression therefore fails deterministically, not by luck of the random init. + """ initialize_for_megatron( tensor_model_parallel_size=tp_size, expert_model_parallel_size=ep_size, @@ -875,6 +962,19 @@ def _test_te_grouped_vs_sequential_default_amax_helper(tp_size, ep_size, quant_c num_moe_experts=4, transformer_impl="modelopt", ) + + # Fill each local expert's grouped weights with a distinct, known magnitude so the per-expert + # weight amax is deterministic (== that magnitude) and diverges across experts by construction. + for te_mlp in (m for m in te_grouped.modules() if isinstance(m, TEGroupedMLP)): + for linear_name in ("linear_fc1", "linear_fc2"): + grouped_linear = getattr(te_mlp, linear_name) + for i in range(grouped_linear.num_gemms): + with torch.no_grad(): + getattr(grouped_linear, f"weight{i}").fill_( + _te_grouped_expert_magnitude(linear_name, i) + ) + + # Propagate the identical per-expert weights to the sequential model. copy_weights_from_grouped_to_non_grouped(te_grouped, sequential) for module in te_grouped.modules(): @@ -891,7 +991,6 @@ def _test_te_grouped_vs_sequential_default_amax_helper(tp_size, ep_size, quant_c 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 @@ -901,23 +1000,37 @@ def _test_te_grouped_vs_sequential_default_amax_helper(tp_size, ep_size, quant_c f"got {len(te_wq)}" ) - expert_amaxes = [] + per_expert_amax = [] for i, expert in enumerate(seq_mlp.local_experts): te_amax = te_wq[i].amax seq_amax = getattr(expert, linear_name).weight_quantizer.amax + expected = _te_grouped_expert_magnitude(linear_name, i) assert te_amax is not None + + # (a) grouped and sequential agree per expert (cross-implementation parity). 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)." - ) + # (b) causal: this expert's amax equals ITS OWN set magnitude, proving the amax + # was computed from that expert's weights (no cross-expert leakage). + assert torch.allclose( + te_amax, torch.full_like(te_amax, expected), rtol=1e-3, atol=1e-3 + ), ( + f"{linear_name} expert {i}: amax {te_amax.reshape(-1)[0].item():.6f} " + f"!= set magnitude {expected}" + ) + # (c) each expert owns a distinct quantizer instance, not a shared one. + for j in range(i): + assert te_wq[i] is not te_wq[j], ( + f"{linear_name}: experts {i} and {j} share a quantizer object" + ) + per_expert_amax.append(te_amax.reshape(-1)[0]) + + # Divergence is now guaranteed by construction (distinct set magnitudes). + stacked = torch.stack(per_expert_amax) + assert (stacked.max() - stacked.min()).item() > 1e-4, ( + f"{linear_name}: per-expert amax did not diverge despite distinct set magnitudes" + ) @pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG]) From 8902e046e4f7661db47cf207526d30bcf25bd80c Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Tue, 28 Jul 2026 18:39:06 -0700 Subject: [PATCH 08/16] docs(changelog): document TEGroupedMLP per-expert quantization + breaking change Add a 0.47 New Features entry for per-expert weight quantization of TEGroupedMLP (GroupedQuantizer, one amax per expert) plus the torch.compile opt-in, and a Backward Breaking Changes note that pre-0.47 quantized TEGroupedMLP checkpoints are incompatible with 0.47's per-expert amax layout. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jennifer Chen --- CHANGELOG.rst | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 329e5d21f05..819066bde73 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,8 +6,13 @@ Changelog **New Features** +- Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD. +- 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. + **Backward Breaking Changes** +- Transformer Engine ``TEGroupedMLP`` (fused MoE experts) now uses **per-expert** weight quantization (one ``amax`` per expert) instead of a single shared ``amax`` across all experts. As a result, ModelOpt checkpoints containing quantized ``TEGroupedMLP`` modules saved before 0.47 are **not compatible** with 0.47: the per-expert ``weight_quantizer`` amax layout differs from the previous single-quantizer layout. Re-run PTQ (or re-quantize) with 0.47 to regenerate compatible checkpoints. + **Deprecations** **Bug Fixes** From 8bd4685a4b2f85113d6e0b02a051be199fd4f96d Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Tue, 28 Jul 2026 19:22:25 -0700 Subject: [PATCH 09/16] feat(quant): make TEGroupedMLP per-expert weight quantizers opt-in + review fixes Add QuantizeConfig.te_per_expert_quantizers (default False = legacy single shared weight quantizer per TEGroupedLinear); True installs a GroupedQuantizer with one quantizer per fused expert. Plumbed via convert_to_quantized_model to the env var the TE plugin reads; _setup and the per-expert methods fall back to the single shared quantizer when off. Adds a parametrized toggle unit test (validated on GPU). Also addresses PR review comments: - transformer_engine post_restore: only a genuine amax/weight shape mismatch falls through to the max|W| recompute; CUDA/OOM/other errors re-raise, and any recompute now warns instead of silently discarding MSE/static/QAD amax. - vllm_fakequant_hf disable loop: handle GroupedQuantizer so the widened _check_all_weight_quantizers_disabled(AnyQuantizer) check passes. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jennifer Chen --- .../torch/export/plugins/vllm_fakequant_hf.py | 6 +- modelopt/torch/quantization/config.py | 11 ++++ modelopt/torch/quantization/conversion.py | 7 ++ .../plugins/transformer_engine.py | 64 ++++++++++++++++--- .../quantization/plugins/test_megatron.py | 54 ++++++++++++++++ 5 files changed, 133 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py index 9168f988146..e1f84e54190 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py @@ -31,6 +31,7 @@ from modelopt.torch.quantization.model_calib import enable_stats_collection, finish_stats_collection from modelopt.torch.quantization.nn import ( AnyQuantizer, + GroupedQuantizer, QuantModule, SequentialQuantizer, TensorQuantizer, @@ -628,7 +629,10 @@ def export_hf_vllm_fq_checkpoint( for attr_name, quantizer in module.named_children(): if not (attr_name.endswith("weight_quantizer") and quantizer.is_enabled): continue - if isinstance(quantizer, SequentialQuantizer): + if isinstance(quantizer, (SequentialQuantizer, GroupedQuantizer)): + # GroupedQuantizer (per-expert TEGroupedLinear) and SequentialQuantizer + # both hold sub-quantizers; disable each so the widened + # _check_all_weight_quantizers_disabled(AnyQuantizer) check passes. quantizer.disable() for sub in quantizer: wqs_to_restore.append((sub, sub._rotate)) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 1e70ec62aab..fdb11e219e9 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1559,6 +1559,17 @@ class QuantizeConfig(ModeloptBaseConfig): validate_default=True, ) + te_per_expert_quantizers: bool = ModeloptField( + default=False, + title="Per-expert weight quantizers for Transformer Engine grouped experts", + description=( + "If True, each fused expert of a Transformer Engine ``TEGroupedLinear`` " + "(``TEGroupedMLP``) gets its own weight quantizer with an independent ``amax``. " + "If False (default, legacy behavior), a single weight quantizer is shared across all " + "experts in the layer." + ), + ) + effective_bits: float | None = ModeloptField( default=None, title="Effective bits per element (autoquant cost override)", diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 00187d291c0..8f9897a81a9 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -16,6 +16,7 @@ """Quantization conversion/restore utilities.""" import fnmatch +import os import re import warnings from collections.abc import Callable @@ -64,6 +65,12 @@ def convert_to_quantized_model(model: ModelLikeModule, config: QuantizeConfig) - # initialize the true module if necessary model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + # TEGroupedLinear per-expert weight quantizers are an opt-in structural choice made in the + # module's _setup (during replace_quant_module), before quant_cfg is applied. Bridge the config + # flag to the env var the TE plugin reads so per-expert quantizers are created only when + # requested; the default keeps the legacy single shared quantizer. + if config.get("te_per_expert_quantizers"): + os.environ["MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER"] = "1" replace_quant_module(model, version=ModeloptStateManager(model).state_version) set_quantizer_by_cfg(model, config.get("quant_cfg", [])) diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index 57affbe3ac1..1c2144d28d5 100644 --- a/modelopt/torch/quantization/plugins/transformer_engine.py +++ b/modelopt/torch/quantization/plugins/transformer_engine.py @@ -35,6 +35,17 @@ _TE_VERSION = Version(te.__version__) _COMPILE_TEGROUPED_WEIGHT_LOOP_ENV = "MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP" +_PER_EXPERT_QUANTIZER_ENV = "MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER" + + +def _te_per_expert_quantizers_enabled() -> bool: + """Whether ``TEGroupedLinear`` gives each fused expert its own weight quantizer (opt-in). + + Default (env unset / ``"0"``) is the legacy single shared weight quantizer for all experts. + ``QuantizeConfig.te_per_expert_quantizers=True`` sets this env var during ``convert``; it can + also be set directly. + """ + return os.getenv(_PER_EXPERT_QUANTIZER_ENV, "0") == "1" def _assert_te_fp8_enabled(): @@ -148,9 +159,14 @@ def _setup(self): # Remove self.weight after setup. delattr(self, "weight") - # 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. + # Opt-in (``QuantizeConfig.te_per_expert_quantizers``): each fused expert gets its own + # weight quantizer (independent amax), stored in a GroupedQuantizer (an nn.ModuleList) + # surfaced as ``weight_quantizer.{i}`` so the fused-experts name normalizer maps them to + # ``*weight_quantizer`` and the stock configs apply. Default (legacy) keeps the single + # shared weight quantizer that ``super()._setup()`` installed above. + if not _te_per_expert_quantizers_enabled(): + return + self.weight_quantizer = GroupedQuantizer( *(copy.deepcopy(self.weight_quantizer) for _ in range(self.num_gemms)) ) @@ -185,6 +201,11 @@ def modelopt_post_restore(self, prefix: str = ""): # static recipes on export and overwrites frozen amax on every QAD resume. Only # re-calibrate a quantizer whose loaded amax is shape-INCOMPATIBLE with its weight # (a genuine TP/EP change between save and restore); otherwise keep it as-is. + # Per-expert amax preservation only applies when each expert has its own quantizer; with + # the legacy single shared quantizer, super().modelopt_post_restore already handled it. + if not isinstance(self.weight_quantizer, GroupedQuantizer): + return + from modelopt.torch.quantization.model_calib import max_calibrate for i in range(self.num_gemms): @@ -200,21 +221,44 @@ def modelopt_post_restore(self, prefix: str = ""): try: wq_i(weight_i) # dry-run: succeeds iff the loaded amax fits this weight shape_ok = True - except Exception: + except Exception as e: + # Only a genuine amax/weight shape mismatch (a TP/EP change between save and + # restore) may fall through to the max|W| recompute below. A CUDA/OOM/device or + # any other error must NOT be silently turned into a recompute -- that would + # discard the stored MSE/static/QAD amax this block exists to preserve. Re-raise + # anything that is not clearly a shape mismatch. + msg = str(e).lower() + is_shape_mismatch = ( + isinstance(e, RuntimeError) + and any( + k in msg for k in ("size", "shape", "must match", "broadcast", "dimension") + ) + and not any(k in msg for k in ("cuda", "out of memory", "device-side", "nccl")) + ) + if not is_shape_mismatch: + raise shape_ok = False finally: q._fake_quant = prev_fake if shape_ok: continue # loaded amax is valid -> keep it, do NOT recompute + # Recompute is lossy for static recipes; never do it silently. + warnings.warn( + f"{type(self).__name__}: restored amax {tuple(q._amax.shape)} for expert {i} " + f"weight_quantizer is shape-incompatible with weight {tuple(weight_i.shape)} " + f"(likely a TP/EP change); recomputing as max|W| and discarding the stored " + f"MSE/static/QAD amax for this expert." + ) 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.""" + grouped = isinstance(self.weight_quantizer, GroupedQuantizer) 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[i] + yield weight_i, (self.weight_quantizer[i] if grouped else self.weight_quantizer) @staticmethod def te_grouped_quantized_linear_fn(package, func_name, self, *args): @@ -245,14 +289,18 @@ def te_grouped_quantized_linear_fn(package, func_name, self, *args): new_args[inp_pos] = self.input_quantizer(args[inp_pos]) 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 + grouped = isinstance(self.weight_quantizer, GroupedQuantizer) + use_compiled_loop = ( + grouped + and 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) + (self.weight_quantizer[gemm_idx] if grouped else self.weight_quantizer)(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 diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index bfaa9fd7974..8ba07470f36 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -15,6 +15,7 @@ import copy import math +import os import re from contextlib import nullcontext from functools import partial @@ -66,6 +67,7 @@ ) from modelopt.torch.quantization.plugins.transformer_engine import ( _COMPILE_TEGROUPED_WEIGHT_LOOP_ENV, + _PER_EXPERT_QUANTIZER_ENV, ) from modelopt.torch.quantization.utils import is_quantized_linear from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -919,6 +921,58 @@ def build(): destroy_model_parallel() +@pytest.mark.parametrize("per_expert", [False, True]) +def test_te_grouped_per_expert_quantizer_toggle(distributed_setup_size_1, monkeypatch, per_expert): + """``QuantizeConfig.te_per_expert_quantizers`` toggles per-expert TEGroupedLinear quantizers. + + Default (False) keeps the legacy single shared weight quantizer for all fused experts; True + installs a GroupedQuantizer per TEGroupedLinear with one quantizer per fused expert. + """ + # Clean env so the default case is not polluted by a prior enabled run. + monkeypatch.delenv(_PER_EXPERT_QUANTIZER_ENV, raising=False) + + 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 + + cfg = copy.deepcopy(mtq.INT8_DEFAULT_CFG) + if per_expert: + cfg["te_per_expert_quantizers"] = True + mtq.quantize(model, cfg, forward) + + grouped_linears = [ + getattr(mlp, name) + for mlp in model.modules() + if isinstance(mlp, TEGroupedMLP) + for name in ("linear_fc1", "linear_fc2") + ] + assert grouped_linears + for gl in grouped_linears: + wq = gl.weight_quantizer + if per_expert: + assert isinstance(wq, mtq.nn.GroupedQuantizer), ( + "te_per_expert_quantizers=True should install a per-expert GroupedQuantizer" + ) + assert len(wq) == gl.num_gemms + else: + assert not isinstance(wq, mtq.nn.GroupedQuantizer) and isinstance( + wq, mtq.nn.TensorQuantizer + ), "default (te_per_expert_quantizers=False) should keep one shared weight quantizer" + + # convert() sets the env var directly; reset it so the enabled case cannot leak to other tests. + os.environ.pop(_PER_EXPERT_QUANTIZER_ENV, None) + destroy_model_parallel() + + def _te_grouped_expert_magnitude(linear_name, local_idx): """Distinct, known weight magnitude for each (linear, local-expert) pair. From 39074f2c655646ef4342b237b7ceae0d180f33e7 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Wed, 29 Jul 2026 09:07:27 -0700 Subject: [PATCH 10/16] feat(quant): make TEGroupedMLP per-expert weight quantizers the default Remove the te_per_expert_quantizers opt-in flag and always give each fused expert of a TEGroupedLinear its own weight quantizer (a GroupedQuantizer with one TensorQuantizer per expert, independent amax). TP>1 works for dynamic quant (no stored amax to shard), so there is no reason to keep the legacy single-shared-quantizer path. - config.py: drop the QuantizeConfig.te_per_expert_quantizers field - conversion.py: drop the config->env bridge (and now-unused import os) - transformer_engine.py: drop _PER_EXPERT_QUANTIZER_ENV and the _te_per_expert_quantizers_enabled() gate; _setup unconditionally installs the per-expert GroupedQuantizer - test: replace the parametrized toggle test with test_te_grouped_per_expert_quantizer_default (GPU-validated: 1 passed) Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jennifer Chen --- modelopt/torch/quantization/config.py | 11 ------ modelopt/torch/quantization/conversion.py | 7 ---- .../plugins/transformer_engine.py | 26 +++----------- .../quantization/plugins/test_megatron.py | 34 +++++-------------- 4 files changed, 14 insertions(+), 64 deletions(-) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index fdb11e219e9..1e70ec62aab 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1559,17 +1559,6 @@ class QuantizeConfig(ModeloptBaseConfig): validate_default=True, ) - te_per_expert_quantizers: bool = ModeloptField( - default=False, - title="Per-expert weight quantizers for Transformer Engine grouped experts", - description=( - "If True, each fused expert of a Transformer Engine ``TEGroupedLinear`` " - "(``TEGroupedMLP``) gets its own weight quantizer with an independent ``amax``. " - "If False (default, legacy behavior), a single weight quantizer is shared across all " - "experts in the layer." - ), - ) - effective_bits: float | None = ModeloptField( default=None, title="Effective bits per element (autoquant cost override)", diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 8f9897a81a9..00187d291c0 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -16,7 +16,6 @@ """Quantization conversion/restore utilities.""" import fnmatch -import os import re import warnings from collections.abc import Callable @@ -65,12 +64,6 @@ def convert_to_quantized_model(model: ModelLikeModule, config: QuantizeConfig) - # initialize the true module if necessary model = model.init_modellike() if isinstance(model, ModelLikeModule) else model - # TEGroupedLinear per-expert weight quantizers are an opt-in structural choice made in the - # module's _setup (during replace_quant_module), before quant_cfg is applied. Bridge the config - # flag to the env var the TE plugin reads so per-expert quantizers are created only when - # requested; the default keeps the legacy single shared quantizer. - if config.get("te_per_expert_quantizers"): - os.environ["MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER"] = "1" replace_quant_module(model, version=ModeloptStateManager(model).state_version) set_quantizer_by_cfg(model, config.get("quant_cfg", [])) diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index 1c2144d28d5..4819dfdd2bc 100644 --- a/modelopt/torch/quantization/plugins/transformer_engine.py +++ b/modelopt/torch/quantization/plugins/transformer_engine.py @@ -35,17 +35,6 @@ _TE_VERSION = Version(te.__version__) _COMPILE_TEGROUPED_WEIGHT_LOOP_ENV = "MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP" -_PER_EXPERT_QUANTIZER_ENV = "MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER" - - -def _te_per_expert_quantizers_enabled() -> bool: - """Whether ``TEGroupedLinear`` gives each fused expert its own weight quantizer (opt-in). - - Default (env unset / ``"0"``) is the legacy single shared weight quantizer for all experts. - ``QuantizeConfig.te_per_expert_quantizers=True`` sets this env var during ``convert``; it can - also be set directly. - """ - return os.getenv(_PER_EXPERT_QUANTIZER_ENV, "0") == "1" def _assert_te_fp8_enabled(): @@ -159,14 +148,10 @@ def _setup(self): # Remove self.weight after setup. delattr(self, "weight") - # Opt-in (``QuantizeConfig.te_per_expert_quantizers``): each fused expert gets its own - # weight quantizer (independent amax), stored in a GroupedQuantizer (an nn.ModuleList) - # surfaced as ``weight_quantizer.{i}`` so the fused-experts name normalizer maps them to - # ``*weight_quantizer`` and the stock configs apply. Default (legacy) keeps the single - # shared weight quantizer that ``super()._setup()`` installed above. - if not _te_per_expert_quantizers_enabled(): - return - + # Each fused expert gets its own weight quantizer (independent amax), stored in a + # GroupedQuantizer (an nn.ModuleList) surfaced as ``weight_quantizer.{i}`` so the + # fused-experts name normalizer maps them to ``*weight_quantizer`` and the stock configs + # apply. This replaces the single shared weight quantizer ``super()._setup()`` installed. self.weight_quantizer = GroupedQuantizer( *(copy.deepcopy(self.weight_quantizer) for _ in range(self.num_gemms)) ) @@ -201,8 +186,7 @@ def modelopt_post_restore(self, prefix: str = ""): # static recipes on export and overwrites frozen amax on every QAD resume. Only # re-calibrate a quantizer whose loaded amax is shape-INCOMPATIBLE with its weight # (a genuine TP/EP change between save and restore); otherwise keep it as-is. - # Per-expert amax preservation only applies when each expert has its own quantizer; with - # the legacy single shared quantizer, super().modelopt_post_restore already handled it. + # weight_quantizer is a GroupedQuantizer (one per expert) after _setup; guard defensively. if not isinstance(self.weight_quantizer, GroupedQuantizer): return diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 8ba07470f36..1782d8d9a85 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -15,7 +15,6 @@ import copy import math -import os import re from contextlib import nullcontext from functools import partial @@ -67,7 +66,6 @@ ) from modelopt.torch.quantization.plugins.transformer_engine import ( _COMPILE_TEGROUPED_WEIGHT_LOOP_ENV, - _PER_EXPERT_QUANTIZER_ENV, ) from modelopt.torch.quantization.utils import is_quantized_linear from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -921,16 +919,12 @@ def build(): destroy_model_parallel() -@pytest.mark.parametrize("per_expert", [False, True]) -def test_te_grouped_per_expert_quantizer_toggle(distributed_setup_size_1, monkeypatch, per_expert): - """``QuantizeConfig.te_per_expert_quantizers`` toggles per-expert TEGroupedLinear quantizers. +def test_te_grouped_per_expert_quantizer_default(distributed_setup_size_1): + """TEGroupedLinear installs a per-expert GroupedQuantizer (one quantizer per fused expert). - Default (False) keeps the legacy single shared weight quantizer for all fused experts; True - installs a GroupedQuantizer per TEGroupedLinear with one quantizer per fused expert. + Per-expert weight quantization is unconditional: every ``TEGroupedLinear`` gets a + ``GroupedQuantizer`` with ``num_gemms`` independent quantizers, not a single shared one. """ - # Clean env so the default case is not polluted by a prior enabled run. - monkeypatch.delenv(_PER_EXPERT_QUANTIZER_ENV, raising=False) - initialize_for_megatron(seed=SEED) model = _gpt_model_provider( tp_size=1, @@ -944,10 +938,7 @@ def test_te_grouped_per_expert_quantizer_toggle(distributed_setup_size_1, monkey if isinstance(module, TopKRouter): module.topk = module.num_experts - cfg = copy.deepcopy(mtq.INT8_DEFAULT_CFG) - if per_expert: - cfg["te_per_expert_quantizers"] = True - mtq.quantize(model, cfg, forward) + mtq.quantize(model, copy.deepcopy(mtq.INT8_DEFAULT_CFG), forward) grouped_linears = [ getattr(mlp, name) @@ -958,18 +949,11 @@ def test_te_grouped_per_expert_quantizer_toggle(distributed_setup_size_1, monkey assert grouped_linears for gl in grouped_linears: wq = gl.weight_quantizer - if per_expert: - assert isinstance(wq, mtq.nn.GroupedQuantizer), ( - "te_per_expert_quantizers=True should install a per-expert GroupedQuantizer" - ) - assert len(wq) == gl.num_gemms - else: - assert not isinstance(wq, mtq.nn.GroupedQuantizer) and isinstance( - wq, mtq.nn.TensorQuantizer - ), "default (te_per_expert_quantizers=False) should keep one shared weight quantizer" + assert isinstance(wq, mtq.nn.GroupedQuantizer), ( + "TEGroupedLinear should install a per-expert GroupedQuantizer" + ) + assert len(wq) == gl.num_gemms - # convert() sets the env var directly; reset it so the enabled case cannot leak to other tests. - os.environ.pop(_PER_EXPERT_QUANTIZER_ENV, None) destroy_model_parallel() From 761d1c14878145dbea46c664ef5be842ff4d63ed Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Wed, 29 Jul 2026 10:08:06 -0700 Subject: [PATCH 11/16] test(quant): cap real torch.compile TEGrouped test at 90s Add an explicit @pytest.mark.timeout(90) to test_te_grouped_real_compile_weight_quantizer_loop so a runaway inductor recompile is bounded tighter than the 120s gpu_megatron group default, while keeping headroom over a cold first-compile so CI doesn't flake. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jennifer Chen --- tests/gpu_megatron/torch/quantization/plugins/test_megatron.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 1782d8d9a85..c05da642a87 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -848,6 +848,7 @@ def compiled(*args): destroy_model_parallel() +@pytest.mark.timeout(90) # real torch.compile: cap runaway inductor recompiles without flaking CI def test_te_grouped_real_compile_weight_quantizer_loop(distributed_setup_size_1, monkeypatch): """Real (unpatched) torch.compile parity for the per-expert weight-quantizer loop. From 57fe9631cf19a0d10005a4728e114835d54e86f7 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Thu, 30 Jul 2026 08:05:07 -0700 Subject: [PATCH 12/16] fix(quant/export): restore output_layer static amax + robust TE per-expert export Static-NVFP4 output_layer (lm_head) amax was dropped on save and never re-applied on restore, so TE-spec PTQ->HF export crashed on the lm_head with "Weight quantizer does not have attribute amax". Two-sided fix (ported from the LOCAL cherry-new-loss-qad branch), plus export-path robustness. - quantization/plugins/megatron.py: register modelopt get/set_extra_state for EVERY QuantModule incl. output_layer (old is_enabled gate ran pre-replacement so it always skipped output_layer). EP-downsize fallback rewrites only the expert index after weight_quantizer (preserves SequentialQuantizer suffixes). - opt/plugins/mcore_dist_checkpointing.py: after load_state_dict, explicitly call set_extra_state for modules with modelopt callbacks. - quantization/plugins/transformer_engine.py: modelopt_post_restore skips the CUDA-only fp4 dry-run when the weight is on CPU (export loads on CPU). - export/unified_export_megatron.py: assert weight_quantizer is GroupedQuantizer or None; warn on TP/EP-mismatch clamp; revert temporary export amax in finally. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jennifer Chen --- .../torch/export/unified_export_megatron.py | 25 +++++++++++++++---- .../opt/plugins/mcore_dist_checkpointing.py | 10 ++++++++ .../torch/quantization/plugins/megatron.py | 13 +++++----- .../plugins/transformer_engine.py | 2 ++ 4 files changed, 38 insertions(+), 12 deletions(-) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 813638390d0..ace6c2825df 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -35,7 +35,7 @@ from modelopt import __version__ from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer -from modelopt.torch.utils import import_plugin +from modelopt.torch.utils import import_plugin, warn_rank_0 from .convert_hf_config import convert_hf_quant_config_format from .model_config import ( @@ -1058,7 +1058,17 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): has_weight = hasattr(module, "weight") grouped_wq = getattr(module, "weight_quantizer", None) - per_expert_wq = isinstance(grouped_wq, GroupedQuantizer) + # Quantized TE grouped experts must be per-expert (GroupedQuantizer); None = unquantized MLP. + assert grouped_wq is None or isinstance(grouped_wq, GroupedQuantizer), ( + f"TEGroupedLinear.weight_quantizer must be GroupedQuantizer or None, got " + f"{type(grouped_wq).__name__}; pre-0.47 single-quantizer checkpoints are not supported." + ) + if grouped_wq is not None and num_experts > len(grouped_wq): + warn_rank_0( + f"TEGroupedMLP has {num_experts} local experts but only {len(grouped_wq)} " + f"per-expert weight quantizers; experts >= {len(grouped_wq)} reuse expert " + f"{len(grouped_wq) - 1}'s scales (TP/EP-mismatch fallback)." + ) ep_size = ( get_expert_model_parallel_world_size() if torch.distributed.is_initialized() else 1 @@ -1114,6 +1124,9 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): local_expert_state: dict[str, torch.Tensor] = {} seen_qformat = None seen_block_size = None + # Dynamic quantizers we populate a temporary export-only amax on; reset in finally so + # export leaves module state unchanged (else a dynamic-NVFP4 quantizer keeps a stale max|W|). + temp_amax_wqs: list = [] try: for local_id in range(num_experts): global_id = local_expert_indices[local_id] @@ -1121,8 +1134,7 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): weight_key = f"weight{local_id}" module.weight = getattr(module, weight_key) - if per_expert_wq: - assert isinstance(grouped_wq, GroupedQuantizer) + if grouped_wq is not None: module.weight_quantizer = grouped_wq[min(local_id, len(grouped_wq) - 1)] # Dynamic-NVFP4 per-expert quantizers carry no stored amax, but # weight_scale_2 derivation asserts one. Max-calibration weight amax @@ -1130,6 +1142,7 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): _wq = module.weight_quantizer if getattr(_wq, "_amax", None) is None and getattr(_wq, "is_enabled", False): _wq.amax = module.weight.detach().abs().max().float() + temp_amax_wqs.append(_wq) name_to_value, qformat, block_size = self._get_quantized_state( module, self.dtype, prefix=prefix @@ -1168,7 +1181,9 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): continue local_expert_state[expert_prefix + key] = val.detach().cpu().clone() finally: - if per_expert_wq: + for _wq in temp_amax_wqs: + _wq.reset_amax() + if grouped_wq is not None: module.weight_quantizer = grouped_wq if not has_weight and hasattr(module, "weight"): delattr(module, "weight") diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index e3bc6e49d8e..2bae53fa6f1 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -172,6 +172,16 @@ def _load_extra_state_from_sharded_checkpoint( extra_state_dict_no_prefix[k[len(prefix) :]] = v model.load_state_dict(extra_state_dict_no_prefix, strict=False) + # PyTorch load_state_dict calls set_extra_state only when the CLASS overrides it; modelopt registers it at + # instance level, so bare output_layer/ColumnParallelLinear is skipped -> saved static amax not applied. + # Invoke it explicitly here (idempotent via allow_post_restore). + for name, module in model.named_modules(): + key = f"{name}._extra_state" if name else "_extra_state" + if key in extra_state_dict_no_prefix and hasattr( + module, "modelopt_set_extra_state_callbacks" + ): + module.set_extra_state(extra_state_dict_no_prefix[key]) + def restore_sharded_modelopt_state( model: list[torch.nn.Module], diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index ea4ddcd3572..6e2c5905e71 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -238,8 +238,10 @@ def quant_module_set_extra_state(self, state: Any): # module loaded at smaller EP (e.g. EP1 export from an EP16 ckpt) has more # experts than the saved state. Per-expert properties are uniform across # experts (amax rides separately as globally-indexed sharded tensors), so - # fall back to expert 0's state. - fallback = re.sub(r"\.\d+$", ".0", name) + # fall back to expert 0's state. Rewrite only the expert index right after + # weight_quantizer, preserving deeper suffixes (e.g. a SequentialQuantizer level + # weight_quantizer..); non-expert names are left unchanged. + fallback = re.sub(r"(weight_quantizer)\.\d+", r"\1.0", name) quantizer_substate = quantizer_state.get(fallback) if quantizer_substate is None: continue @@ -305,11 +307,8 @@ def _configure_attention_for_kv_cache_quant(module: Attention): def _register_extra_state_callbacks(model: torch.nn.Module): for name, module in model.named_modules(): if type(module) in QuantModuleRegistry: - # Skip output_layer w/o enabled weight_quantizer - if name.endswith("output_layer") and not getattr( - getattr(module, "weight_quantizer", None), "is_enabled", False - ): - continue + # Register for EVERY QuantModule incl. output_layer: the old is_enabled gate ran + # pre-replacement (weight_quantizer None) so it skipped output_layer -> static amax dropped on save. register_modelopt_extra_state_callbacks( module, quant_module_get_extra_state, diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index 4819dfdd2bc..8d36722405d 100644 --- a/modelopt/torch/quantization/plugins/transformer_engine.py +++ b/modelopt/torch/quantization/plugins/transformer_engine.py @@ -196,6 +196,8 @@ def modelopt_post_restore(self, prefix: str = ""): weight_i = getattr(self, f"weight{i}", None) if weight_i is None: continue + if weight_i.device.type != "cuda": + continue # export loads weights on CPU; the fp4 dry-run needs CUDA — keep loaded amax wq_i = self.weight_quantizer[i] q = wq_i[0] if isinstance(wq_i, SequentialQuantizer) else wq_i if not hasattr(q, "_amax") or q._amax is None: From 790a87114fa536af0e2630f0561ab5be12b6b880 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Fri, 31 Jul 2026 14:39:39 -0700 Subject: [PATCH 13/16] fix(quant): emit empty extra_state for unquantized modules After 57fe9631cf registered get/set_extra_state on every QuantModule (incl. output_layer), quant_module_get_extra_state still always returned a non-empty modelopt_quantizer_state (it iterates all TensorQuantizers, including disabled ones). An *unquantized* output_layer therefore emitted a non-empty _extra_state, and Megatron-Bridge's save_megatron_model -> GPTModel.sharded_state_dict (which asserts output_layer._extra_state is empty) failed with "Boolean value of Tensor with more than one value is ambiguous". Gate quant_module_get_extra_state to return {} when the module has no enabled TensorQuantizer and is not a compressed RealQuantLinear; the aggregator then returns None, satisfying Megatron's assert, while a genuinely quantized lm_head still saves its amax. Fixes tests/examples/megatron_bridge/test_quantize_export.py::test_quantize_and_export Signed-off-by: Jennifer Chen Co-Authored-By: Claude Opus 4.8 (1M context) --- modelopt/torch/quantization/plugins/megatron.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 6e2c5905e71..280ed71e2b5 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -141,6 +141,12 @@ def quant_module_get_extra_state(self) -> dict: QuantModule's extra_state with QuantModule.get_extra_state() which avoids the need to store the full module name. """ + # Nothing quantized here -> return {} so unquantized output_layer._extra_state stays empty (Megatron asserts empty). + if not isinstance(self, RealQuantLinear) and not any( + isinstance(m, TensorQuantizer) and m.is_enabled for m in self.modules() + ): + return {} + extra_state = {} quantizer_state = {} From 72d3cb5e714d7d8dfe6aa9adb0fc2a93583ebb64 Mon Sep 17 00:00:00 2001 From: James Shen Date: Mon, 3 Aug 2026 10:42:18 -0700 Subject: [PATCH 14/16] feat(megatron-bridge): Nemotron-Nano-3 W4A16 NVFP4 four_over_six PTQ/QAD support Also fixes a ModelOpt bug that affects ANY model with an untied lm_head quantized through Megatron-Bridge, not just Nemotron: the output layer was silently exported as BF16. An untied `lm_head` (`output_layer`) was silently exported as BF16 instead of NVFP4 whenever the model was built with Megatron-Bridge, even though the recipe enabled `*output_layer*weight_quantizer`. The same recipe under Megatron-LM quantized it correctly, so this was not a configuration problem. `_MegatronParallelLinear.sharded_state_dict()` special-cases `output_layer` and asks `megatron.training.get_args()` whether embeddings are untied. Megatron-Bridge has no global args store, so the call raises and the handler falls back to "tied", taking the early return that drops all quantizer state. Fixing that alone is not sufficient: the dist-checkpoint loader silently skips any checkpoint key the model does not advertise, and `sharded_state_dict()` can only advertise a buffer that already exists, so the calibrated scales in the checkpoint had nowhere to land. * `_resolve_output_layer_untied()` reads `share_embeddings_and_output_weights` off the model, which Megatron-Core carries under both frameworks, and records it on the config so `sharded_state_dict()` can consult it. `get_args()` remains the fallback, so Megatron-LM behavior is unchanged, and an unknown result still means "tied". * Materialize missing weight-quantizer scale buffers before the load plan is built. `_amax` must be allocated flat as `[numel // block, 1]`: `_process_quantizer_amax` exposes it to the checkpoint as an `[out_features, blocks]` view over the same storage, so the loader writes straight through. Allocating the viewed shape loads successfully but leaves the wrong in-memory rank, which then breaks the exporter's scale math. `_global_amax` is registered directly because its property lives on `StaticBlockScaleQuantizer` and the module is still a plain `TensorQuantizer` here. The export example now fails loudly instead of quietly emitting BF16 when an enabled NVFP4 weight quantizer is missing either scale, naming the module and the attribute. Static-block NVFP4 needs both; an `_amax`-only check let a half-restored quantizer through, which failed much later inside `NVFP4QTensor.quantize` where `scale * scale_2` broadcasts `[N, 1]` against `[N]`. `MODELOPT_ALLOW_UNCALIBRATED_NVFP4=1` restores the previous behavior, reported rather than silent. Also included: Megatron-Bridge PTQ/QAD enablement for Nemotron-style hybrid MoE models (gate the non-grouped MoE spec to the quantized student so the BF16 teacher is built with its natural spec, SFT-masked distillation, student initialization from a Megatron checkpoint, calibration random offset). Verified end to end against a Megatron-LM-produced reference: `lm_head.weight` is now `U8 [131072, 1344]` with `weight_scale` and `weight_scale_2`, 18487 keys and 72 excluded modules, matching the reference exactly (previously BF16 `[131072, 2688]`, 18485 keys, 73 excluded with `lm_head` among them). Two Megatron-Bridge compatibility shims were removed after being shown unnecessary: a `DistillationProvider.to_cfg_dict` monkeypatch (a 5-iteration distillation run trains and checkpoints cleanly without it) and an `InferenceCudaGraphScope` enum stub added for a Megatron-LM-PTQ import path that is not used (zero occurrences across a full PTQ/QAD/export run). Signed-off-by: James Shen The two example scripts no longer hard-code model-shape assumptions. MoE expert grouping is a `--grouped_experts` flag on both `quantize.py` and the exporter, defaulting to non-grouped so existing behavior is unchanged; per-block NVFP4 requires non-grouped because TEGroupedLinear can only represent a per-tensor scale, while per-tensor recipes can now opt into faster grouped GEMM. The exporter reads `mtp_num_layers` from the checkpoint's run_config.yaml instead of assuming 0, and `quantize.py` now warns when it drops MTP heads, matching prune_minitron.py. Expert grouping is deliberately NOT derived from run_config.yaml: a MambaModelProvider sets the layout via mamba_stack_spec, so a non-grouped checkpoint still records `moe_grouped_gemm: true` and trusting it would build a mismatched model. Signed-off-by: James Shen Review fixes: the export guard no longer excludes StaticBlockScaleQuantizer, which is the class that owns `_global_amax` -- excluding it skipped exactly the case the guard exists to catch -- and it now filters on `is_enabled` to match the message it prints. The scale-buffer materialization checks `in_features % block_size`, matching the `view(weight.shape[0], -1)` performed later rather than total element count, and warns instead of silently leaving the buffers unallocated. Signed-off-by: James Shen --- examples/megatron_bridge/distill.py | 299 +++++++++++++----- .../export_quantized_megatron_to_hf.py | 102 ++++++ examples/megatron_bridge/quantize.py | 39 +++ modelopt/torch/distill/plugins/megatron.py | 46 +++ .../torch/quantization/plugins/megatron.py | 103 +++++- modelopt/torch/utils/dataset_utils.py | 20 +- .../utils/plugins/megatron_calibration.py | 4 + 7 files changed, 509 insertions(+), 104 deletions(-) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index dfd404fab30..66e636c1ca0 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -25,15 +25,18 @@ import os import torch -from _distillation_provider import convert_to_distillation_provider -from export_distilled_megatron_to_hf import export_llm_to_hf, save_vlm_to_hf from megatron.bridge import AutoBridge +from megatron.bridge.models.distillation_provider import ( + DistillationProvider, + convert_to_distillation_provider, +) from megatron.bridge.recipes.utils.optimizer_utils import ( distributed_fused_adam_with_cosine_annealing, ) from megatron.bridge.training.config import ( CheckpointConfig, ConfigContainer, + FinetuningDatasetConfig, GPTDatasetConfig, LoggerConfig, MockGPTDatasetConfig, @@ -46,18 +49,83 @@ from megatron.bridge.training.post_training.distillation import ModelOptDistillConfig from megatron.core.datasets.utils import get_blend_from_list from megatron.core.distributed import DistributedDataParallelConfig -from megatron.core.utils import unwrap_model from transformers import AutoConfig import modelopt.torch.distill as mtd +import modelopt.torch.distill.plugins.megatron as mtd_mcore import modelopt.torch.utils.distributed as dist -from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0 +from modelopt.torch.utils import print_args, print_rank_0 from modelopt.torch.utils.plugins.mbridge import load_modelopt_megatron_checkpoint with contextlib.suppress(ModuleNotFoundError): import modelopt.torch.puzzletron.plugins.mbridge # noqa: F401 +# TODO: Megatron-Bridge does not (yet) expose a hook to initialize the student before the +# knowledge-distillation conversion, so we patch ``DistillationProvider.provide`` to do it. Replace +# this block once a first-class mechanism is available upstream. +# +# Maps id(distill_provider) -> megatron_checkpoint_path for providers whose student should be +# initialized from a Megatron checkpoint. A registry is used (instead of an instance attribute) +# because a DistillationProvider proxies attribute assignment to its teacher once the teacher is +# set, so anything stored on the instance would leak onto the teacher. +_MEGATRON_STUDENT_CKPT_PATHS: dict[int, str] = {} + +_original_distill_provide = DistillationProvider.provide + + +def _distill_provide_with_megatron_student( + self, pre_process=None, post_process=None, vp_stage=None +): + """Replacement for ``DistillationProvider.provide`` that can initialize the student from a ckpt. + + For providers registered in ``_MEGATRON_STUDENT_CKPT_PATHS``, the student is built and its weights + (plus, for a quantized checkpoint, the ModelOpt quantize mode) are restored from the Megatron + checkpoint *before* the knowledge-distillation conversion -- otherwise the quantize mode is lost, + since ``restore_sharded_modelopt_state`` is a no-op once a model is already converted. The rest + mirrors the upstream implementation. Patched at the class level (not the instance) to avoid the + teacher-proxying issue described on ``_MEGATRON_STUDENT_CKPT_PATHS``. + """ + if vp_stage is not None: + raise ValueError("ModelOpt KD currently does not support virtual-pipeline parallel.") + + megatron_path = _MEGATRON_STUDENT_CKPT_PATHS.get(id(self)) + if megatron_path is None: + # If a path was registered (for some provider) but this provide() call doesn't match, + # the provider was likely copied/wrapped between convert_to_distillation_provider() and now, + # so the id()-keyed lookup silently misses. Fail loudly rather than train an uninitialized + # student (this script only ever builds one DistillationProvider). + if _MEGATRON_STUDENT_CKPT_PATHS: + raise RuntimeError( + "DistillationProvider.provide() found no registered Megatron-student checkpoint path " + "for this provider, but one was registered for a different provider id -- the provider " + "was likely copied/wrapped. Update this workaround." + ) + return _original_distill_provide(self, pre_process, post_process, vp_stage) + + student_model = self._super_class.provide(self, pre_process, post_process, vp_stage) + print_rank_0(f"Loading student weights from Megatron checkpoint {megatron_path}") + load_modelopt_megatron_checkpoint([student_model], megatron_path) + # Hack to get teacher's pre-wrap hooks called to potentially load HF weights + teacher_model = self.teacher.provide_distributed_model( + wrap_with_ddp=False, mixed_precision_wrapper=None + )[0] + kd_cfg = mtd_mcore.setup_distillation_config( + self.kd_config, student_model.config, teacher_model.config + ) + modelopt_cfg = { + "teacher_model": teacher_model, + "criterion": kd_cfg.criterion, + "loss_balancer": kd_cfg.loss_balancer, + } + kd_model = mtd.convert(student_model, mode=[("kd_loss", modelopt_cfg)]) + mtd_mcore.adjust_distillation_model_for_mcore(kd_model, kd_cfg) + return kd_model + + +DistillationProvider.provide = _distill_provide_with_megatron_student + + def get_args(): """Parse command-line arguments.""" parser = argparse.ArgumentParser(description="Distillation for Megatron-Bridge.") @@ -75,6 +143,16 @@ def get_args(): help="HuggingFace model name or path for the teacher (e.g. Qwen/Qwen3-8B)", ) parser.add_argument("--trust_remote_code", action="store_true", help="Trust remote code") + parser.add_argument( + "--student_nongrouped_experts", + action="store_true", + help=( + "Build the quantized student with non-grouped MoE experts. Required for STATIC-BLOCK " + "NVFP4 recipes (e.g. four_over_six): TEGroupedLinear only supports per-tensor scales, " + "not per-block. Leave OFF (default) for dynamic NVFP4 / grouped-expert checkpoints " + "(e.g. Nemotron-3-Nano). Never applied to the BF16 teacher." + ), + ) parser.add_argument( "--student_megatron_path", type=str, @@ -110,6 +188,20 @@ def get_args(): parser.add_argument( "--use_mock_data", action="store_true", help="Use mock data instead of --data_paths" ) + parser.add_argument( + "--sft", + action="store_true", + help="SFT-masked distillation: read raw prompt-completion jsonl from --sft_dataset_root and " + "mask the loss to the completion (assistant response) tokens. Uses GPTSFTDatasetConfig + the " + "real (HuggingFace) tokenizer instead of the pretraining GPTDataset + NullTokenizer.", + ) + parser.add_argument( + "--sft_dataset_root", + type=str, + default=None, + help="Directory containing training.jsonl / validation.jsonl with prompt-completion " + '{"input": , "output": } records (used with --sft).', + ) # Training & Eval arguments parser.add_argument( "--output_dir", type=str, required=True, help="Folder for logging and checkpoint saving" @@ -188,18 +280,20 @@ def get_args(): type=str, required=False, default=None, - help="Reference HF model with a homogeneous architecture, used as the export template for a " - "heterogeneous (Puzzletron/NAS) student's weights. Defaults to --student_hf_path, which is " - "correct for homogeneous students; unused for VLMs.", + help="HuggingFace model ID to use as template for export (e.g., Qwen/Qwen3-0.6B). " + "Should match the base architecture of the student model if --hf_export_path is provided.", ) args = parser.parse_args() # Sanity checks - if not args.use_mock_data and not args.data_paths: + if args.sft: + if not args.sft_dataset_root: + raise ValueError("--sft requires --sft_dataset_root (dir with training.jsonl/validation.jsonl).") + elif not args.use_mock_data and not args.data_paths: raise ValueError("Must provide either --data_paths or set --use_mock_data.") - if args.student_hf_model is None: - args.student_hf_model = args.student_hf_path + if args.hf_export_path and not args.student_hf_model: + raise ValueError("Must provide --student_hf_model if --hf_export_path is provided.") print_args(args) @@ -211,7 +305,7 @@ def main(args: argparse.Namespace): tensorboard_dir = os.path.join(args.output_dir, "tb_logs") # Build student and teacher model providers - def _build_model_provider(hf_path, load_weights=True): + def _build_model_provider(hf_path, load_weights=True, quantized=True): bridge = AutoBridge.from_hf_pretrained(hf_path, trust_remote_code=args.trust_remote_code) provider = bridge.to_megatron_provider(load_weights=load_weights) @@ -224,6 +318,32 @@ def _build_model_provider(hf_path, load_weights=True): provider.expert_model_parallel_size = args.ep_size provider.expert_tensor_parallel_size = 1 # Expert tensor parallelism is not supported provider.seq_length = args.seq_length + # Match the PTQ/quantize.py setup: MTP is not supported during QAD, and NVFP4 per-block + # quantization requires non-grouped experts (TEGroupedLinear only supports per-tensor). + # For a hybrid Mamba provider the layer SPEC must be rebuilt with moe_grouped_gemm=False -- + # setting the flag alone does not propagate. Mirror modelopt's load_mbridge_model_from_hf. + provider.mtp_num_layers = 0 + from modelopt.torch.nas.plugins.megatron import get_te_mamba_stack_spec + + if quantized and args.student_nongrouped_experts: + # Static-block NVFP4 students need non-grouped experts (TEGroupedLinear can't do per-block + # scales). OFF by default = grouped = committed behavior (works for dynamic NVFP4 like + # Nano-3). NEVER applied to the BF16 teacher (would misplace its MoE experts). + if hasattr(provider, "mamba_stack_spec"): + provider.mamba_stack_spec = get_te_mamba_stack_spec(moe_grouped_gemm=False) + elif (getattr(provider, "num_moe_experts", 0) or 0) > 0: + provider.moe_grouped_gemm = False + # Regularize the MoE router during QAD so it does not degenerate. Jenny's working Megatron-LM + # QAD uses `--moe-aux-loss-coeff 1e-4 --moe-router-load-balancing-type seq_aux_loss`; without + # it our router weights drifted the most (~8% vs ~1% elsewhere) and the MoE broke. Applies to + # both providers, but only the (trained) student's aux loss affects optimization. + if (getattr(provider, "num_moe_experts", 0) or 0) > 0: + provider.moe_router_load_balancing_type = "seq_aux_loss" + provider.moe_aux_loss_coeff = 1e-4 + if args.sft: + # Finetuning (SFT) with context parallel (CP>1) requires per-token loss so the + # response loss-mask reduces correctly across the CP ranks. + provider.calculate_per_token_loss = os.environ.get("FORCE_NO_PER_TOKEN_LOSS", "0") != "1" if args.recompute_granularity is not None: provider.recompute_granularity = args.recompute_granularity provider.recompute_method = args.recompute_method @@ -245,52 +365,25 @@ def _build_model_provider(hf_path, load_weights=True): # Gradient accumulation fusion is not supported with ModelOpt quantized models. Disable it # before the model is built so the student's linear layers are constructed accordingly. student_provider.gradient_accumulation_fusion = False - teacher_provider = _build_model_provider(args.teacher_hf_path) + teacher_provider = _build_model_provider(args.teacher_hf_path, quantized=False) + # Wrap into DistillationProvider kd_config = ModelOptDistillConfig( skip_lm_loss=not args.no_skip_lm_loss, kd_loss_scale=args.kd_loss_scale ) - - # VLM detection convention: HF VLM configs expose a ``vision_config``, and Megatron-Bridge nests - # the text model under the ``language_model`` submodule (used as ``distill_submodule`` below). If a - # future model breaks either convention, the ``getattr(model, "language_model")`` in the provider - # will error loudly rather than silently distilling the wrong module. - is_vlm = hasattr( - AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code), - "vision_config", - ) - - if is_vlm: - warn_rank_0( - "VLM detected: distilling model.language_model only (vision tower / projector untouched). " - "To export megatron non-quantized checkpoint, use export_distilled_megatron_to_hf.py" - ) distill_provider = convert_to_distillation_provider( - student_provider, - teacher_provider, - kd_config, - distill_submodule="language_model" if is_vlm else None, + student_provider, teacher_provider, kd_config ) if args.student_megatron_path: - # QAD: restore the quantized student weights + ModelOpt state before the KD conversion (a no-op - # once converted). Prepend so this runs before the provider's KD-conversion pre-wrap hook. if student_has_modelopt_state: print_rank_0( f"Detected ModelOpt state in {args.student_megatron_path}; " "restoring quantizers for Quantization Aware Distillation (QAD)." ) - - def _restore_student_hook(model_chunks): - print_rank_0( - f"Loading student weights from Megatron checkpoint {args.student_megatron_path}" - ) - load_modelopt_megatron_checkpoint( - [unwrap_model(model_chunks[0])], args.student_megatron_path - ) - return model_chunks - - distill_provider.register_pre_wrap_hook(_restore_student_hook, prepend=True) + # Register so the patched DistillationProvider.provide initializes this provider's student + # from the Megatron checkpoint (see _distill_provide_with_megatron_student). + _MEGATRON_STUDENT_CKPT_PATHS[id(distill_provider)] = args.student_megatron_path # Build optimizer and scheduler optimizer_config, scheduler_config = distributed_fused_adam_with_cosine_annealing( @@ -301,24 +394,48 @@ def _restore_student_hook(model_chunks): ) # Build dataset config - dataset_kwargs = { - "seq_length": args.seq_length, - "path_to_cache": args.data_path_to_cache, - "random_seed": args.seed, - "reset_attention_mask": False, - "reset_position_ids": False, - "eod_mask_loss": False, - "num_dataset_builder_threads": 1, - "data_sharding": True, - "dataloader_type": "single", - "skip_getting_attention_mask_from_dataset": True, - } - if args.use_mock_data: - dataset_config = MockGPTDatasetConfig(**dataset_kwargs) + if args.sft: + # SFT-masked (Quantization-Aware) distillation via the container's Bridge FinetuningDatasetConfig + # -> NeMo-style GPTSFTDataset. `dataset_root` holds training.jsonl / validation.jsonl with + # {"input": , "output": } records. prompt_template="{input}{output}" tokenizes + # input+output verbatim (adjacent placeholders, no separator) matching the identity-formatted + # source; label_key="output" + answer_only_loss=True mask the loss to the assistant response only + # (answer_start_idx == len(context_ids)); truncation_field="input" truncates the context if needed. + dataset_config = FinetuningDatasetConfig( + seq_length=args.seq_length, + dataset_root=args.sft_dataset_root, + seed=args.seed, + dataloader_type="batch", + do_validation=True, + do_test=False, + dataset_kwargs={ + "prompt_template": "{input}{output}", + "label_key": "output", + "truncation_field": "input", + "answer_only_loss": True, + "add_bos": False, + "add_eos": True, + }, + ) else: - # Convert flat CLI list (e.g. ["1.0", "/path/data"]) to Megatron blend format - blend = get_blend_from_list(args.data_paths) - dataset_config = GPTDatasetConfig(blend=blend, split="99,1,0", **dataset_kwargs) + dataset_kwargs = { + "seq_length": args.seq_length, + "path_to_cache": args.data_path_to_cache, + "random_seed": args.seed, + "reset_attention_mask": False, + "reset_position_ids": False, + "eod_mask_loss": False, + "num_dataset_builder_threads": 1, + "data_sharding": True, + "dataloader_type": "single", + "skip_getting_attention_mask_from_dataset": True, + } + if args.use_mock_data: + dataset_config = MockGPTDatasetConfig(**dataset_kwargs) + else: + # Convert flat CLI list (e.g. ["1.0", "/path/data"]) to Megatron blend format + blend = get_blend_from_list(args.data_paths) + dataset_config = GPTDatasetConfig(blend=blend, split="99,1,0", **dataset_kwargs) # Assemble ConfigContainer and run distillation config = ConfigContainer( @@ -341,7 +458,9 @@ def _restore_student_hook(model_chunks): grad_reduce_in_fp32=True, overlap_grad_reduce=True, overlap_param_gather=True, - average_in_collective=True, + # Finetuning (SFT) with CP>1 requires per-token loss (set on the provider) and + # average_in_collective=False (the per-token loss is summed, not averaged, in the collective). + average_in_collective=(not args.sft) or os.environ.get("FORCE_NO_PER_TOKEN_LOSS", "0") == "1", use_distributed_optimizer=True, ), dataset=dataset_config, @@ -354,22 +473,36 @@ def _restore_student_hook(model_chunks): wandb_entity=args.wandb_entity, # optional wandb_exp_name=args.wandb_exp_name, ), - tokenizer=TokenizerConfig( - tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size + tokenizer=( + TokenizerConfig( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model=args.student_hf_path, + hf_tokenizer_kwargs={"trust_remote_code": args.trust_remote_code}, + ) + if args.sft + else TokenizerConfig( + tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size + ) ), checkpoint=CheckpointConfig( save_interval=args.eval_interval, save=checkpoint_dir, load=checkpoint_dir, # Resume from this directory (if exists) - most_recent_k=5, # Keeps 5 most recent checkpoints (not metric-based) + most_recent_k=2, # Keeps 2 most recent checkpoints (each ~413GB here; fs1 near quota) ckpt_format="torch_dist", - async_save=True, + async_save=False, # sync save: async writer repeatedly corrupted iter-400 ckpt (inline_container) fully_parallel_save=True, ), rng=RNGConfig(seed=args.seed), mixed_precision="bf16_mixed", ) + # QAD with NVFP4 fake-quant makes the first optimizer step (many grad-accum microbatches) + # very slow; raise the NCCL process-group timeout above the default (~10 min) so the initial + # step does not trip the collective watchdog. Guarded in case the config field is renamed. + if hasattr(config, "dist") and hasattr(config.dist, "distributed_timeout_minutes"): + config.dist.distributed_timeout_minutes = 60 + print_rank_0("\nStarting distillation...") distill(config) print_rank_0( @@ -377,22 +510,7 @@ def _restore_student_hook(model_chunks): " in megatron distributed checkpoint format.\n" ) - if args.hf_export_path and is_vlm: - # Only the language model was distilled; export it back into the full VLM. - print_rank_0(f"Exporting distilled VLM to HF format to {args.hf_export_path}") - # ``distill`` tore down the model-parallel groups on exit, so rebuild them. - distill_provider.initialize_model_parallel(seed=args.seed) - full_student = distill_provider.full_model - # Strip the distillation wrapper -> plain trained language model (in place; reassign to be safe). - full_student.language_model = mtd.export(full_student.language_model) - save_vlm_to_hf( - full_student, - args.hf_export_path, - args.student_hf_path, - trust_remote_code=args.trust_remote_code, - ) - print_rank_0(f"Saved distilled VLM to {args.hf_export_path} in HF format") - elif args.hf_export_path: + if args.hf_export_path: print_rank_0(f"Exporting final distilled ckpt to HF format to {args.hf_export_path}") # Save rank before destroying process group (dist.rank() won't work after destruction) is_rank_0 = dist.rank() == 0 @@ -402,13 +520,20 @@ def _restore_student_hook(model_chunks): dist.cleanup() if is_rank_0: - export_llm_to_hf( + export_bridge = AutoBridge.from_hf_pretrained( + args.student_hf_model, trust_remote_code=args.trust_remote_code + ) + # Copy weights and remote code + export_bridge.export_ckpt( megatron_path=f"{checkpoint_dir}/iter_{args.train_iters:07d}", - hf_export_path=args.hf_export_path, - student_hf_path=args.student_hf_path, - template_hf=args.student_hf_model, - trust_remote_code=args.trust_remote_code, + hf_path=args.hf_export_path, + show_progress=True, + strict=True, ) + # Copy config.json from student_hf_path (handles both local paths and HF model IDs) + AutoConfig.from_pretrained( + args.student_hf_path, trust_remote_code=args.trust_remote_code + ).save_pretrained(args.hf_export_path) if __name__ == "__main__": diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py index 17db5e6da34..a0b2e8be407 100644 --- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py +++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py @@ -36,6 +36,9 @@ """ import argparse +import yaml +import pathlib +import os import torch from megatron.bridge.models.hf_pretrained.utils import is_safe_repo @@ -44,6 +47,10 @@ import modelopt.torch.utils.distributed as dist from modelopt.torch.export import export_mcore_gpt_to_hf from modelopt.torch.utils import print_args, print_rank_0 +from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( + StaticBlockScaleQuantizer, + TensorQuantizer, +) from modelopt.torch.utils.plugins.mbridge import ( load_mbridge_model_from_hf, load_modelopt_megatron_checkpoint, @@ -52,6 +59,13 @@ def get_args() -> argparse.Namespace: parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument( + "--grouped_experts", + action="store_true", + help="Build MoE experts grouped (GroupedMLP). Default is non-grouped, which per-block " + "NVFP4 checkpoints require. Set this only when the checkpoint was saved with grouped " + "experts; the layout must match or the weights will not load.", + ) parser.add_argument( "--hf_model_name_or_path", type=str, @@ -99,7 +113,39 @@ def get_args() -> argparse.Namespace: return args +def _provider_overrides_from_checkpoint(megatron_path: str) -> dict: + """Read ``mtp_num_layers`` from the checkpoint so the exporter matches how it was saved. + + Only ``mtp_num_layers`` is taken from here. ``moe_grouped_gemm`` is deliberately NOT derived: + for a ``MambaModelProvider`` the expert layout is set by ``mamba_stack_spec``, so a checkpoint + saved with non-grouped experts still records ``moe_grouped_gemm: true`` and trusting it would + build a mismatched model. + """ + defaults = {"mtp_num_layers": 0} + run_config = next(iter(sorted(pathlib.Path(megatron_path).glob("*/run_config.yaml"))), None) + if run_config is None: + run_config = pathlib.Path(megatron_path) / "run_config.yaml" + if not run_config.exists(): + print_rank_0(f"No run_config.yaml under {megatron_path}; using defaults {defaults}.") + return defaults + try: + cfg = yaml.safe_load(run_config.read_text()) or {} + except Exception as exc: + print_rank_0(f"Could not parse {run_config} ({exc}); using defaults {defaults}.") + return defaults + model_cfg = cfg.get("model") if isinstance(cfg.get("model"), dict) else cfg + resolved = { + key: model_cfg.get(key, default) + for key, default in defaults.items() + if isinstance(model_cfg, dict) + } + resolved = {**defaults, **resolved} + print_rank_0(f"Model shape from {run_config.name}: {resolved}") + return resolved + + def main(args: argparse.Namespace): + _ckpt_shape = _provider_overrides_from_checkpoint(args.megatron_path) trust_remote_code = is_safe_repo( trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_name_or_path ) @@ -116,8 +162,11 @@ def main(args: argparse.Namespace): "num_layers_in_first_pipeline_stage": args.num_layers_in_first_pipeline_stage, "num_layers_in_last_pipeline_stage": args.num_layers_in_last_pipeline_stage, "pipeline_dtype": torch.bfloat16, + "mtp_num_layers": _ckpt_shape["mtp_num_layers"], }, init_model_parallel=True, + # Default non-grouped, matching quantize.py; the layout must match the checkpoint. + moe_grouped_gemm=args.grouped_experts, load_weights=False, # The weights come from the Megatron checkpoint, so HF weights are not loaded ) @@ -127,6 +176,59 @@ def main(args: argparse.Namespace): load_modelopt_megatron_checkpoint(model, args.megatron_path) unwrapped_model = unwrap_model(model[0]) + # Static-NVFP4 export guard. + # + # An *enabled* NVFP4 weight quantizer that reaches the exporter without its calibrated scales + # means the values stored in the checkpoint were not restored. Exporting such a weight silently + # falls back to BF16: the result is larger than the recipe specifies and no longer matches it, + # with nothing in the logs to say so. Fail loudly instead. + # + # Static-block NVFP4 needs BOTH ``_amax`` (per block) and ``_global_amax`` (per tensor). A + # missing ``_global_amax`` slips past an ``_amax``-only check and then fails much later inside + # ``NVFP4QTensor.quantize``, where ``scale * scale_2`` broadcasts [N, 1] against [N] into an + # N x N allocation. Naming the attribute here turns that into an actionable message. + # + # Set MODELOPT_ALLOW_UNCALIBRATED_NVFP4=1 to keep the previous behavior (disable the quantizer + # and emit BF16), which is then reported rather than silent. + uncalibrated: list[tuple[str, str]] = [] + for name, module in unwrapped_model.named_modules(): + # StaticBlockScaleQuantizer must be INCLUDED: `_global_amax` is defined on it, so excluding + # it would skip exactly the case this guard exists to catch. Only report enabled quantizers, + # matching the message. + if not isinstance(module, TensorQuantizer) or not getattr(module, "is_enabled", False): + continue + block_sizes = getattr(module, "_block_sizes", None) + is_nvfp4 = getattr(module, "_num_bits", None) == (2, 1) and ( + isinstance(block_sizes, dict) and block_sizes.get("scale_bits") == (4, 3) + ) + if not is_nvfp4: + continue + if getattr(module, "_amax", None) is None: + uncalibrated.append((name, "_amax")) + elif ( + isinstance(module, StaticBlockScaleQuantizer) + or block_sizes.get("type") == "static" + ) and getattr(module, "_global_amax", None) is None: + uncalibrated.append((name, "_global_amax")) + + if uncalibrated: + detail = ", ".join(f"{name}.{attr}" for name, attr in uncalibrated[:8]) + if len(uncalibrated) > 8: + detail += ", ..." + message = ( + f"{len(uncalibrated)} enabled NVFP4 weight quantizer(s) are missing calibrated scales " + f"after loading {args.megatron_path}: {detail}. These weights would be exported as " + "BF16 instead of NVFP4. Re-run PTQ with a ModelOpt that saves and restores this " + "quantizer state, or set MODELOPT_ALLOW_UNCALIBRATED_NVFP4=1 to export them as BF16." + ) + if os.environ.get("MODELOPT_ALLOW_UNCALIBRATED_NVFP4") != "1": + raise RuntimeError(message) + print_rank_0(f"WARNING (MODELOPT_ALLOW_UNCALIBRATED_NVFP4=1): {message}") + for name, _ in uncalibrated: + unwrapped_model.get_submodule(name).disable() + else: + print_rank_0("All enabled NVFP4 weight quantizers have calibrated scales.") + # Extra modules (Medusa / EAGLE / MTP) only exist on the last pipeline stage. Use an all-reduce # MAX over all ranks (rather than a broadcast from a hard-coded source rank) so the decision is # correct regardless of pipeline placement / global rank ordering. diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index 3454da00441..fb898c99933 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -174,6 +174,19 @@ def get_args() -> argparse.Namespace: parser.add_argument( "--calib_num_samples", type=int, default=1024, help="Number of samples for calibration" ) + parser.add_argument( + "--grouped_experts", + action="store_true", + help="Build MoE experts grouped (GroupedMLP). Default is non-grouped (SequentialMLP), " + "which per-block NVFP4 requires because TEGroupedLinear can only represent per-tensor " + "scales. Set this for per-tensor recipes on MoE models, where grouped GEMM is faster. " + "The export must use the matching layout.", + ) + parser.add_argument( + "--calib_random_offset", + action="store_true", + help="Drop a random leading-token offset before packing calib windows (Megatron-LM --calib-use-random-offset).", + ) parser.add_argument("--calib_batch_size", type=int, default=1, help="Calibration batch size") parser.add_argument( "--seq_length", @@ -275,6 +288,18 @@ def get_quant_config(args: argparse.Namespace) -> dict: return mtq_config +_MTP_HF_CONFIG_FIELDS = ("num_nextn_predict_layers", "mtp_num_hidden_layers", "mtp_num_layers") + + +def _hf_config_has_mtp(hf_cfg) -> bool: + """Whether an HF config declares MTP heads (checked top-level and under ``text_config``).""" + return any( + cfg is not None and getattr(cfg, field, 0) + for cfg in (getattr(hf_cfg, "text_config", None), hf_cfg) + for field in _MTP_HF_CONFIG_FIELDS + ) + + def main(args: argparse.Namespace): bridge, _provider, model, unwrapped_model, tokenizer = load_mbridge_model_from_hf( hf_model_name_or_path=args.hf_model_name_or_path, @@ -284,14 +309,27 @@ def main(args: argparse.Namespace): "pipeline_model_parallel_size": args.pp_size, "expert_model_parallel_size": args.ep_size, "context_parallel_size": args.cp_size, + "mtp_num_layers": 0, # MTP not supported during calibration "expert_tensor_parallel_size": 1, # Expert tensor parallelism is not supported "pipeline_dtype": torch.bfloat16, "seq_length": args.seq_length, "gradient_accumulation_fusion": False, # not supported }, init_model_parallel=True, + # Default non-grouped: per-block NVFP4 needs it (TEGroupedLinear is per-tensor only). + # Opt into grouped for per-tensor recipes, where grouped GEMM is faster. + moe_grouped_gemm=args.grouped_experts, ) + # `mtp_num_layers=0` above drops MTP heads: calibration does not support them. Say so rather + # than silently shipping a checkpoint without a head the model declares. + if _hf_config_has_mtp(bridge.hf_pretrained.config): + warn_rank_0( + "Dropping Multi-Token Prediction (MTP): calibration does not support it. The exported " + "checkpoint will not contain MTP weights and standard autoregressive inference is " + "unaffected. To use MTP speculative decoding, run a separate phase with mtp_num_layers>0." + ) + # Only the language model is quantized (vision tower + projector stay full precision) language_model = getattr(unwrapped_model, "language_model", unwrapped_model) is_vlm = language_model is not unwrapped_model @@ -367,6 +405,7 @@ def main(args: argparse.Namespace): seq_length=args.seq_length, batch_size=args.calib_batch_size, pack=True, # Megatron pretraining-style global-stream document packing + random_offset=args.calib_random_offset, ) # Run text prefill on the language model: we quantize the root (a VLM root forward expects diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index 7e81c21a462..e02d65eeef0 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -608,6 +608,52 @@ def _set_input_tensor(self, input_tensors: list[Tensor]): # HACK: Concatenate output tensors when PP>1 so they can be passed between ranks. def _forward(self, *args, **kwargs): + # Static-block NVFP4: promote the student's weight quantizers once, after the + # checkpoint amax/scales have been loaded, so the training forward takes the + # StaticBlockScaleQuantizer path rather than the generic FP8 (E4M3) path. Promotion + # cannot happen at build time because the scales only exist after the load. + # + # NOTE: in practice this converts exactly ONE module -- ``output_layer``. A measured run + # reports ``already promoted 460, converted 1, skipped 0``: every other quantizer is + # already a StaticBlockScaleQuantizer by the time training starts. So this is a workaround + # for output_layer being the one module the restore path does not promote (the same + # asymmetry behind its weight-quantizer scales not being restored). The better fix is to + # promote it on the normal path; until then, without this block the output projection + # would train through the generic FP8 path instead of static-block NVFP4. + if not getattr(self, "_modelopt_nvfp4_promoted", False): + from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( + StaticBlockScaleQuantizer, + TensorQuantizer, + ) + + n_promoted = n_skipped = 0 + for name, module in self.named_modules(): + if not isinstance(module, TensorQuantizer) or isinstance( + module, StaticBlockScaleQuantizer + ): + continue + block_sizes = getattr(module, "_block_sizes", None) + is_nvfp4 = getattr(module, "_num_bits", None) == (2, 1) and ( + isinstance(block_sizes, dict) and block_sizes.get("scale_bits") == (4, 3) + ) + if not is_nvfp4: + continue + amax = getattr(module, "_amax", None) + if amax is None: + # Uncalibrated: leave it alone rather than silently changing precision. + logger.warning(f"NVFP4 weight quantizer {name} has no _amax; not promoted.") + n_skipped += 1 + continue + StaticBlockScaleQuantizer.from_tensor_quantizer( + module, global_amax=amax.detach().float().abs().max() + ) + n_promoted += 1 + if n_promoted or n_skipped: + logger.info( + f"Promoted {n_promoted} NVFP4 weight quantizer(s) to " + f"StaticBlockScaleQuantizer ({n_skipped} skipped)." + ) + self._modelopt_nvfp4_promoted = True with torch.no_grad(): self._teacher_model.eval() teacher_output = self._teacher_model(*args, **kwargs) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 752dd801a6e..2f742b87c5a 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -248,6 +248,22 @@ def _incompatible_method(self, *args, **kwargs): return _incompatible_method +def _resolve_output_layer_untied(model: torch.nn.Module) -> bool | None: + """Whether ``output_layer`` weights are untied from the input embeddings, or None if unknown. + + Megatron-Core models carry ``share_embeddings_and_output_weights`` (Megatron-Bridge sets it + from the HF config, Megatron-LM from ``--untie-embeddings-and-output-weights``), so reading + it off the model works under both frameworks. ``megatron.training.get_args()`` does not: + Bridge has no global args store, and defaulting to "tied" there silently drops the + ``output_layer`` weight-quantizer state from the sharded checkpoint. + """ + for _, module in model.named_modules(): + shared = getattr(module, "share_embeddings_and_output_weights", None) + if shared is not None: + return not bool(shared) + return None + + def megatron_replace_quant_module_hook(model: torch.nn.Module): """Configure Megatron-Core model quantization support. @@ -261,6 +277,8 @@ def megatron_replace_quant_module_hook(model: torch.nn.Module): 3. For Attention modules, we configure them to use core_attention path for KV cache quantization. """ + untied = _resolve_output_layer_untied(model) + def _configure_attention_for_kv_cache_quant(module: Attention): """Configure Attention module for KV cache quantization compatibility.""" # Disable flash_decode if enabled - it bypasses core_attention (only called during inference) @@ -287,11 +305,19 @@ def _configure_attention_for_kv_cache_quant(module: Attention): def _register_extra_state_callbacks(model: torch.nn.Module): for name, module in model.named_modules(): if type(module) in QuantModuleRegistry: - # Skip output_layer w/o enabled weight_quantizer - if name.endswith("output_layer") and not getattr( - getattr(module, "weight_quantizer", None), "is_enabled", False - ): - continue + # Skip output_layer w/o enabled weight_quantizer. This hook also runs BEFORE + # QuantModule replacement (e.g. on restore), when ``weight_quantizer`` does not + # exist yet -- the old check then always skipped, so output_layer never received + # ModelOpt extra-state callbacks and its quantizer state (promotion to + # StaticBlockScaleQuantizer, ``_amax``, ``_global_amax``) was never restored. + # Fall back to the tying flag: an untied output_layer is quantizable. + if name.endswith("output_layer"): + _wq = getattr(module, "weight_quantizer", None) + _skip = ( + not getattr(_wq, "is_enabled", False) if _wq is not None else not untied + ) + if _skip: + continue register_modelopt_extra_state_callbacks( module, quant_module_get_extra_state, @@ -307,6 +333,10 @@ def _register_extra_state_callbacks(model: torch.nn.Module): if "vision_model" not in name: # We only enable hetereogenous_dist_checkpoint for language model, vision model is not quantized module.config.hetereogenous_dist_checkpoint = True + if untied is not None: + # Carried on the config so _MegatronParallelLinear.sharded_state_dict can read + # it without Megatron-LM global args (absent under Megatron-Bridge). + module.config.modelopt_output_layer_untied = untied _register_extra_state_callbacks(module) @@ -374,18 +404,63 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): # output_layer.input_quantizer._amax but TP-only does not. This lead to # state_dict mismatch. if prefix.endswith("output_layer."): - try: - from megatron.training import get_args as _mlm_get_args - - _untied = bool( - getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False) - ) - except Exception as e: - warn_rank_0(f"Failed to get Megatron arg untie_embeddings_and_output_weights: {e}") - _untied = False + # Prefer the model-derived flag (set by megatron_replace_quant_module_hook); it is the + # only source available under Megatron-Bridge, which has no global args store. + _untied = getattr(self.config, "modelopt_output_layer_untied", None) + if _untied is None: + try: + from megatron.training import get_args as _mlm_get_args + + _untied = bool( + getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False) + ) + except Exception as e: + warn_rank_0( + f"Failed to get Megatron arg untie_embeddings_and_output_weights: {e}" + ) + _untied = False if not _untied: return super().sharded_state_dict(prefix, sharded_offsets, metadata) + # Materialize missing weight-quantizer scale buffers so their keys appear in the load + # plan -- the dist-checkpoint loader SILENTLY SKIPS any checkpoint key the model does + # not advertise, which leaves output_layer uncalibrated and exports it as BF16. + # ``_amax`` must be allocated FLAT ``[numel // block, 1]``: that is the in-memory + # layout every other block-quantized layer uses, and ``_process_quantizer_amax`` below + # exposes it to the checkpoint as a ``[out_features, blocks]`` VIEW sharing the same + # storage, so the loader writes straight through. Allocating the viewed shape instead + # loads fine but leaves the wrong in-memory shape, which breaks the export scale math. + _wq = getattr(self, "weight_quantizer", None) + if _wq is not None and getattr(_wq, "is_enabled", False): + _block_sizes = getattr(_wq, "_block_sizes", None) or {} + _block = _block_sizes.get(-1) or _block_sizes.get(1) + # `_process_quantizer_amax` later does `v.view(weight.shape[0], -1)`, which + # requires in_features (not just numel) to divide evenly by the block size. + if _block and self.weight.shape[-1] % int(_block) == 0: + if getattr(_wq, "_amax", None) is None: + _wq.amax = torch.zeros( + self.weight.numel() // int(_block), + 1, + dtype=torch.float32, + device=self.weight.device, + ) + # register_buffer directly: the ``global_amax`` property lives on + # StaticBlockScaleQuantizer, and on restore this is still a plain + # TensorQuantizer (promotion happens later), so the setter is unavailable. + if getattr(_wq, "_global_amax", None) is None: + _wq.register_buffer( + "_global_amax", + torch.zeros((), dtype=torch.float32, device=self.weight.device), + ) + else: + # Leaving the buffers unallocated is the silent-drop failure this block exists + # to prevent, so say so rather than proceeding quietly. + warn_rank_0( + f"{prefix}weight_quantizer: cannot materialize scale buffers " + f"(block_size={_block}, in_features={self.weight.shape[-1]}); its " + "calibrated scales will not be restored from the checkpoint." + ) + quantizer_state_dict = {} for k, v in self.state_dict(prefix="", keep_vars=True).items(): if "_quantizer" in k and "_amax" in k: diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index fd9b1e2f55e..e980383e730 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -679,7 +679,11 @@ def __len__(self): def _pack_documents_into_rows( - samples: list[str], tokenizer: "PreTrainedTokenizerBase", seq_length: int, num_rows: int + samples: list[str], + tokenizer: "PreTrainedTokenizerBase", + seq_length: int, + num_rows: int, + random_offset: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Global-stream document packing (Megatron-LM pretraining style). @@ -696,14 +700,23 @@ def _pack_documents_into_rows( eos_id = tokenizer.eos_token_id pad_id = tokenizer.pad_token_id has_eos_sep = eos_id is not None + # With random_offset (Megatron-LM --calib-use-random-offset), build one extra window + # of headroom, then drop a random number of leading tokens so the window grid shifts and + # calibration samples mid-document positions differently (relevant for long-context KV stats). + target_len = num_rows * seq_length + (seq_length if random_offset else 0) token_stream: list[int] = [] for s in samples: token_stream.extend(tokenizer.encode(s, add_special_tokens=False)) if has_eos_sep: token_stream.append(eos_id) - if len(token_stream) >= num_rows * seq_length: + if len(token_stream) >= target_len: break + if random_offset: + max_off = min(seq_length, max(0, len(token_stream) - num_rows * seq_length)) + if max_off > 0: + token_stream = token_stream[random.randint(0, max_off):] + n_full = min(num_rows, len(token_stream) // seq_length) rows_ids: list[list[int]] = [ token_stream[i * seq_length : (i + 1) * seq_length] for i in range(n_full) @@ -749,6 +762,7 @@ def get_dataset_dataloader( include_labels: bool = False, apply_chat_template: bool = False, pack: bool = False, + random_offset: bool = False, distributed: bool = False, sampler_kwargs: dict | None = None, ) -> DataLoader: @@ -858,7 +872,7 @@ def get_dataset_dataloader( if pack: total_rows = sum(num_samples) input_ids, attention_mask = _pack_documents_into_rows( - all_samples, tokenizer, max_sample_length, total_rows + all_samples, tokenizer, max_sample_length, total_rows, random_offset=random_offset ) if input_ids.shape[0] < total_rows: warn_rank_0( diff --git a/modelopt/torch/utils/plugins/megatron_calibration.py b/modelopt/torch/utils/plugins/megatron_calibration.py index 4da38858209..a069f1505a4 100644 --- a/modelopt/torch/utils/plugins/megatron_calibration.py +++ b/modelopt/torch/utils/plugins/megatron_calibration.py @@ -50,6 +50,7 @@ def get_megatron_calibration_dataloader( device: torch.device | str | None = "cuda", apply_chat_template: bool = True, pack: bool = False, + random_offset: bool = False, ) -> torch.utils.data.DataLoader: """Build a DP-sharded calibration dataloader for Megatron-Core models. @@ -76,6 +77,7 @@ def get_megatron_calibration_dataloader( device=device, apply_chat_template=apply_chat_template, pack=pack, + random_offset=random_offset, distributed=dp_size > 1, sampler_kwargs={ "num_replicas": dp_size, @@ -95,6 +97,7 @@ def get_megatron_calibration_forward_loop( device: torch.device | str | None = "cuda", apply_chat_template: bool = True, pack: bool = False, + random_offset: bool = False, ) -> Callable[[torch.nn.Module], None]: """Build a Megatron-Core calibration ``forward_loop(model)``. @@ -116,6 +119,7 @@ def get_megatron_calibration_forward_loop( device=device, apply_chat_template=apply_chat_template, pack=pack, + random_offset=random_offset, ) def _forward_loop(model: torch.nn.Module) -> None: From 972a57ac14601894052b525cb094ab30c693a1fb Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Tue, 4 Aug 2026 09:15:16 -0700 Subject: [PATCH 15/16] remove deepcopy Signed-off-by: Jennifer Chen --- modelopt/torch/quantization/plugins/transformer_engine.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index 8d36722405d..6294f8d9889 100644 --- a/modelopt/torch/quantization/plugins/transformer_engine.py +++ b/modelopt/torch/quantization/plugins/transformer_engine.py @@ -15,7 +15,6 @@ """Support quantization for Transformer Engine layers.""" -import copy import inspect import os import warnings @@ -29,7 +28,7 @@ from modelopt.torch.quantization.utils import replace_function -from ..nn import GroupedQuantizer, QuantModuleRegistry, SequentialQuantizer +from ..nn import GroupedQuantizer, QuantModuleRegistry, SequentialQuantizer, TensorQuantizer from .custom import _ParallelLinear _TE_VERSION = Version(te.__version__) @@ -153,7 +152,7 @@ def _setup(self): # fused-experts name normalizer maps them to ``*weight_quantizer`` and the stock configs # apply. This replaces the single shared weight quantizer ``super()._setup()`` installed. self.weight_quantizer = GroupedQuantizer( - *(copy.deepcopy(self.weight_quantizer) for _ in range(self.num_gemms)) + *(TensorQuantizer(self.default_quant_desc_weight) for _ in range(self.num_gemms)) ) # Compile only the per-expert quantizer loop. The surrounding TE grouped GEMM remains From 5f2f3de3caedbe932b76d2099644eb68638bb6af Mon Sep 17 00:00:00 2001 From: James Shen Date: Wed, 5 Aug 2026 00:01:30 -0700 Subject: [PATCH 16/16] fix(quantization): resolve default_quant_desc_weight in TEGroupedLinear setup PR #1550 builds the new per-expert GroupedQuantizer in _QuantTEGroupedLinear._setup() from `self.default_quant_desc_weight`, but that attribute does not resolve on the class: _ParallelLinear derives from QuantModule, not _QuantLinear, and _ParallelLinear._setup() itself references the class attribute directly (see plugins/custom.py). Constructing any grouped-experts model therefore fails with AttributeError: QuantTEColumnParallelGroupedLinear object has no attribute default_quant_desc_weight Use _QuantLinear.default_quant_desc_weight, matching what _ParallelLinear._setup() already does for the non-grouped path. Verified on Nemotron-Nano-3 W4A16 NVFP4 four_over_six with TEGroupedMLP: PTQ (6382 quantizers) -> QAD (200 iters, logits-KD 3.37e-2 -> 1.91e-2) -> HF export (18487 keys; routed-expert scales stay per-block, e.g. [1856, 116] with block 16) -> compressed-tensors -> served on stock vLLM 0.26.0 at TP=2. Signed-off-by: James Shen --- modelopt/torch/quantization/plugins/transformer_engine.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index 6294f8d9889..0ed2d5108c2 100644 --- a/modelopt/torch/quantization/plugins/transformer_engine.py +++ b/modelopt/torch/quantization/plugins/transformer_engine.py @@ -29,6 +29,7 @@ from modelopt.torch.quantization.utils import replace_function from ..nn import GroupedQuantizer, QuantModuleRegistry, SequentialQuantizer, TensorQuantizer +from ..nn.modules.quant_linear import _QuantLinear from .custom import _ParallelLinear _TE_VERSION = Version(te.__version__) @@ -152,7 +153,10 @@ def _setup(self): # fused-experts name normalizer maps them to ``*weight_quantizer`` and the stock configs # apply. This replaces the single shared weight quantizer ``super()._setup()`` installed. self.weight_quantizer = GroupedQuantizer( - *(TensorQuantizer(self.default_quant_desc_weight) for _ in range(self.num_gemms)) + *( + TensorQuantizer(_QuantLinear.default_quant_desc_weight) + for _ in range(self.num_gemms) + ) ) # Compile only the per-expert quantizer loop. The surrounding TE grouped GEMM remains