NVIDIA · yueshen2016 · Jul 22, 2026 · Jul 22, 2026 · Jul 22, 2026 · Jul 22, 2026
@@ -9,8 +9,13 @@ Experimental
 0.46 (2026-08-xx)
 ^^^^^^^^^^^^^^^^^
 
+- 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.46 are **not compatible** with 0.46: the per-expert ``weight_quantizer`` amax layout differs from the previous single-quantizer layout. Re-run PTQ (or re-quantize) with 0.46 to regenerate compatible checkpoints.
+
 - Remove the ``examples/diffusers/eval`` image-quality evaluation example (ImageReward / CLIP-IQA / CLIP metrics) and its references in ``examples/diffusers/README.md``. The example was deprecated in 0.45 and is no longer maintained.
 - Remove the deprecated ``examples/llm_autodeploy`` example (deprecated in 0.45). Use TensorRT-LLM's `AutoDeploy <https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/auto_deploy>`_ directly together with ModelOpt PTQ in ``examples/llm_ptq``.
 - Remove the deprecated ``examples/llm_qad`` Megatron-LM QAD example (deprecated in 0.45). Use the `megatron_bridge QAD example <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge#quantization-aware-distillation-qad>`_ instead, which provides a simpler Python-based interface and better model coverage.

@@ -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}
-    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}
+    model_cfg = cfg.get("model") if isinstance(cfg.get("model"), dict) else cfg
+    resolved = dict(defaults)
+    if isinstance(model_cfg, dict):
+        for key, default in defaults.items():
+            value = model_cfg.get(key)
+            resolved[key] = default if value is None else value
-    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}
+    model_cfg = cfg.get("model") if isinstance(cfg.get("model"), dict) else cfg
+    resolved = dict(defaults)
+    if isinstance(model_cfg, dict):
+        for key, default in defaults.items():
+            value = model_cfg.get(key)
+            resolved[key] = default if value is None else value
+    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.

@@ -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

@@ -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
-            for name, module in self.named_modules():
-                if not isinstance(module, TensorQuantizer) or isinstance(
-                    module, StaticBlockScaleQuantizer
-                ):
-                    continue
+            n_promoted = n_skipped = 0
+            teacher_modules = {id(m) for m in self._teacher_model.modules()}
+            for name, module in self.named_modules():
+                if id(module) in teacher_modules:
+                    continue
+                if not isinstance(module, TensorQuantizer) or isinstance(
+                    module, StaticBlockScaleQuantizer
+                ):
+                    continue
-            for name, module in self.named_modules():
-                if not isinstance(module, TensorQuantizer) or isinstance(
-                    module, StaticBlockScaleQuantizer
-                ):
-                    continue
+            n_promoted = n_skipped = 0
+            teacher_modules = {id(m) for m in self._teacher_model.modules()}
+            for name, module in self.named_modules():
+                if id(module) in teacher_modules:
+                    continue
+                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)

@@ -29,7 +29,13 @@
 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,
+    GroupedQuantizer,
+    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 +131,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 "
@@ -625,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))