Autoquant and GPTQ in support in Megatron-Core [OMNIML-3151]

.github/workflows/gpu_tests.yml

@@ -42,7 +42,7 @@ jobs:
             timeout: 60
             container_image: nvcr.io/nvidia/pytorch:26.05-py3
           - example: gpu_megatron
-            timeout: 60
+            timeout: 75
             container_image: nvcr.io/nvidia/nemo:26.06
           - example: gpu_trtllm
             timeout: 15

CHANGELOG.rst

@@ -76,6 +76,7 @@ Changelog
 - Add shared Megatron-Core calibration forward loop: ``modelopt.torch.utils.plugins.megatron_calibration.get_megatron_calibration_forward_loop`` produces the ``forward_loop`` callable expected by ``mtq.quantize`` / ``mtp.prune``. Replaces the bespoke calibration loops in Megatron-LM and Megatron-Bridge for quantization and pruning with a single canonical implementation.
 - Support Megatron-Core checkpoint restore and export for MSE ``NVFP4StaticQuantizer``.
 - Add mixed-precision FP8 + NVFP4 export for Megatron-Core: per-layer ``quant_algo`` recorded under ``quantized_layers`` in ``hf_quant_config.json``, PP-aware ``kv_cache_dtype`` gather, fused-QKV exclude split into per-HF-name ``q/k/v_proj`` entries.
+- Add AutoQuant and GPTQ support for Megatron-Core models, including MCore-specific AutoQuant hooks and decoder-layer discovery for GPTQ layerwise calibration.
 - Add support for ``active_params`` (for MoE models) and ``memory_mb`` constraints in Minitron pruning on top of existing ``params`` constraint. You can also provide multiple constraints. See `examples/pruning/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/pruning>`_ for more details. The underlying utility functions ``mcore_param_count``, ``mcore_memory_footprint_mb``, and ``print_mcore_model_stats`` in ``modelopt.torch.nas.plugins.megatron_model_stats`` are also available for standalone use to compute parameter counts and memory footprints (weights + KV-cache + Mamba state) for any Megatron-Core model.
 - Add Minitron pruning support for Megatron-Bridge Gemma3 models.
 - Add end-to-end optimization tutorial for Minitron pruning + two-phase distillation (80B @ 8K + 20B @ 32K long-context = 100B tokens) + FP8 PTQ + vLLM deployment for Nemotron-3-Nano-30B-A3B-BF16 (MoE + Mamba-Transformer hybrid) → Pruned 22B/A3.0B active params, along with data blend preparation steps (with tool-calling data) and detailed pruning / data-blend / long-context ablations. See `examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/>`_ for details.

modelopt/torch/quantization/algorithms.py

@@ -378,13 +378,14 @@ def get_score(self, recipe: QuantRecipe) -> float:
                 total_score += importance.cpu().item()
                 continue
 
-            if parallel_state.expert_model_parallel_group.is_initialized():
-                # TODO: Support expert model parallelism for score estimation
-                warnings.warn("AutoQuantize does not support expert model parallelism yet.")
             importance = importance.cpu()
             importance = DistributedProcessGroup.get_dist_syncd_obj(
                 importance,
-                [parallel_state.tensor_parallel_group, parallel_state.data_parallel_group],
+                [
+                    parallel_state.tensor_parallel_group,
+                    parallel_state.data_parallel_group,
+                    parallel_state.expert_model_parallel_group,
+                ],
                 sum,
             )
             total_score += importance.item()
@@ -408,13 +409,12 @@ def get_cost(self, recipe: QuantRecipe, cost_weight: float | None = None) -> flo
                 cost += weight_size * recipe.compression
                 continue
 
-            if parallel_state.expert_model_parallel_group.is_initialized():
-                # TODO: Support expert model parallelism
-                warnings.warn("AutoQuantize does not support expert model parallelism yet.")
-
             weight_size = DistributedProcessGroup.get_dist_syncd_obj(
                 weight_size,
-                [parallel_state.tensor_parallel_group],
+                [
+                    parallel_state.tensor_parallel_group,
+                    parallel_state.expert_model_parallel_group,
+                ],
                 sum,
             )
 
@@ -466,6 +466,8 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
         # gate_proj, up_proj, down_proj for Qwen3 like MoE models
         r"^(.*?\.mlp\.experts)\.\d+\.(gate_proj|up_proj|down_proj)$",
         r"^(.*?\.mixer\.experts)\.\d+\.(up_proj|down_proj)$",  # NemotronH MoE experts
+        # NemotronH MoE experts in MCore naming (linear_fc1=gate+up fused, linear_fc2=down)
+        r"^(.*?\.mlp\.experts\.local_experts)\.\d+\.(linear_fc1|linear_fc2)$",
         r"^(.*?)\.(gate_proj|up_proj)$",  # gate_proj, up_proj for llama like models
         r"^(.*?)\.(\d+\.(w1|w2|w3))$",  # mixtral experts
         r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$",  # dbrx experts
@@ -875,6 +877,15 @@ def _get_total_weight_size(modules):
             for module in modules
         )
 
+    @staticmethod
+    def _get_total_weight_size_from_candidate_stats(candidate_stats):
+        no_quant_recipe = QuantRecipe(quant_cfg=None)
+        total_weight_size = 0
+        for candidate_stat in candidate_stats.values():
+            no_quant_idx = candidate_stat["formats"].index(no_quant_recipe)
+            total_weight_size += candidate_stat["costs"][no_quant_idx]
+        return total_weight_size
+
     def _get_constraints_for_search(self, max_weight_size, lower_bound=None):
         constraints = {
             "weight_size_after_compression": (
@@ -905,9 +916,14 @@ def run_search(self):
         )
 
         compression = self._get_formatted_weight_compression_constraint()
-        total_weight_size = self._cost_model.total_weight_size(
-            self.model.named_modules(), self._is_auto_quantize_module, self.config["cost"]
-        )
+        if self.candidate_stats:
+            total_weight_size = self._get_total_weight_size_from_candidate_stats(
+                self.candidate_stats
+            )
+        else:
+            total_weight_size = self._cost_model.total_weight_size(
+                self.model.named_modules(), self._is_auto_quantize_module, self.config["cost"]
+            )
         self.cost_denominator = total_weight_size
         max_weight_size = total_weight_size * compression
         if verbose:
@@ -928,12 +944,16 @@ def run_search(self):
         best_recipe = {}
         best_constraints, best_scores = 0, 0
         for name, best_hparam_recipe_info in best_recipe_info.items():
-            # Solvers could give different solutions for the same layer across DP/TP groups even though
-            # the scores and costs are the same. Lets make sure the same recipe is selected across DP/TP
+            # Solvers could give different solutions for the same layer across DP/TP/EP groups even though
+            # the scores and costs are the same. Lets make sure the same recipe is selected across DP/TP/EP
             _ps = self.model.get_submodule(name.split(".quant_recipe")[0]).parallel_state
             best_format = DistributedProcessGroup.get_dist_syncd_obj(
                 best_hparam_recipe_info["format"],
-                [_ps.data_parallel_group, _ps.tensor_parallel_group],
+                [
+                    _ps.data_parallel_group,
+                    _ps.tensor_parallel_group,
+                    _ps.expert_model_parallel_group,
+                ],
                 lambda a: a[0],
             )
 

modelopt/torch/quantization/model_calib.py

@@ -57,14 +57,6 @@
 )
 from .utils.calib_utils import _GPTQ_HELPER_REGISTRY, GPTQHelper
 
-try:
-    from .plugins.megatron import _check_nvfp4_static_tp_supported
-except ImportError:
-
-    def _check_nvfp4_static_tp_supported(model: nn.Module) -> None:  # no-op without megatron
-        return
-
-
 __all__ = [
     "CalibratorFactory",
     "awq",
@@ -305,7 +297,12 @@ def max_calibrate(
             module.layer_sync_moe_local_experts_amax(sync_weight_amax=sync_expert_weight_amax)
 
     # Fail fast on NVFP4 static-block with TP>1 (sharded_state_dict treats _amax as replicated).
-    _check_nvfp4_static_tp_supported(model)
+    try:
+        from .plugins.megatron import _check_nvfp4_static_tp_supported
+    except ImportError:
+        pass
+    else:
+        _check_nvfp4_static_tp_supported(model)
 
     if not distributed_sync:
         # Single-process: _amax is final.

modelopt/torch/quantization/plugins/megatron.py

@@ -16,6 +16,7 @@
 """Support quantization for megatron linear layers."""
 
 import types
+from contextlib import contextmanager
 from typing import Any
 
 import megatron.core.parallel_state as mcore_parallel
@@ -39,11 +40,13 @@
 from modelopt.torch.utils import warn_rank_0
 from modelopt.torch.utils.distributed import ParallelState
 
+from ..algorithms import AutoQuantizeGradientSearcher
 from ..conversion import maybe_promote_nvfp4_static_quantizer
 from ..nn import QuantModule, QuantModuleRegistry, SequentialQuantizer, TensorQuantizer
 from ..nn.modules.quant_linear import RealQuantLinear
 from ..qtensor import QTensorWrapper
 from ..utils import sync_moe_expert_amax
+from ..utils.layerwise_calib import LayerActivationCollector
 from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear
 
 try:
@@ -612,10 +615,11 @@ def _setup(self):
                 expert_model_parallel_group=mcore_parallel.get_expert_model_parallel_group(),
             )
 
-        # Initialize parallel state for submodules local_experts.*.linear_fc1 and local_experts.*.linear_fc2
+        # These child linears are still native MCore modules here. Seed `_parallel_state`
+        # directly so the later QuantModule conversion sees the intended parallel state.
         for expert in self.local_experts:
-            expert.linear_fc1.parallel_state = self.parallel_state
-            expert.linear_fc2.parallel_state = self.parallel_state
+            expert.linear_fc1._parallel_state = self.parallel_state
+            expert.linear_fc2._parallel_state = self.parallel_state
 
     def layer_sync_moe_local_experts_amax(self, sync_weight_amax=False):
         """Sync quantizer amax across local experts in a SequentialMLP.
@@ -721,9 +725,10 @@ def _setup(self):
                     tensor_parallel_group=mcore_parallel.get_expert_tensor_parallel_group(),
                     expert_model_parallel_group=mcore_parallel.get_expert_model_parallel_group(),
                 )
-            # initialize parallel state for submodules linear_fc1 and linear_fc2
-            self.linear_fc1.parallel_state = self.parallel_state
-            self.linear_fc2.parallel_state = self.parallel_state
+            # These child linears are still native MCore modules here. Seed `_parallel_state`
+            # directly so the later QuantModule conversion sees the intended parallel state.
+            self.linear_fc1._parallel_state = self.parallel_state
+            self.linear_fc2._parallel_state = self.parallel_state
 
     @QuantModuleRegistry.register({TEDotProductAttention: "TEDotProductAttention"})
     class _QuantTEDotProductAttention(QuantModule):
@@ -779,3 +784,43 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
             # Affine KVCache Quant bias vector.
             state_dict = self.state_dict(prefix="", keep_vars=True)
             return make_sharded_tensors_for_checkpoint(state_dict, prefix, {}, sharded_offsets)
+
+
+def _is_supported_megatron_model(model: torch.nn.Module) -> bool:
+    return isinstance(model, MegatronModule)
+
+
+@contextmanager
+def _megatron_grad_ckpt_context(model: torch.nn.Module):
+    # Megatron configures activation recompute at model build time via TransformerConfig,
+    # so there is no runtime flag to flip here.
+    yield
+
+
+def _is_param_grad_enabled_for_megatron(pname: str, model: torch.nn.Module) -> bool:
+    return "weight" in pname
+
+
+AutoQuantizeGradientSearcher.register_custom_support(
+    _is_supported_megatron_model,
+    _megatron_grad_ckpt_context,
+    _is_param_grad_enabled_for_megatron,
+)
+
+
+def get_mcore_layerwise_calibration_layers(
+    model: torch.nn.Module,
+) -> list[torch.nn.Module] | torch.nn.ModuleList | None:
+    if not hasattr(model, "decoder") or not hasattr(model.decoder, "layers"):
+        return None
+    decoder_layers = model.decoder.layers
+    if getattr(model, "output_layer", None) is None:
+        return decoder_layers
+    layers = list(decoder_layers)
+    layers.append(model.output_layer)
+    return layers
+
+
+LayerActivationCollector.register_decoder_layer_support(
+    _is_supported_megatron_model, get_mcore_layerwise_calibration_layers
+)

modelopt/torch/quantization/utils/calib_utils.py

@@ -63,6 +63,9 @@ def update_hessian(input, hessian, n_samples):
     input_flat = input.reshape(-1, input.shape[-1]).t().float()
     batch_size = input_flat.shape[1]
 
+    if batch_size == 0:  # in MOEs some experts receive no tokens
+        return hessian, n_samples
+
     # Incremental averaging: scale down old hessian
     hessian *= n_samples / (n_samples + batch_size)
     n_samples += batch_size

modelopt/torch/utils/plugins/megatron_calibration.py

@@ -34,12 +34,13 @@
     from transformers import PreTrainedTokenizerBase, ProcessorMixin
 
 __all__ = [
+    "get_megatron_calibration_dataloader",
     "get_megatron_calibration_forward_loop",
     "get_megatron_vlm_calibration_forward_loop",
 ]
 
 
-def get_megatron_calibration_forward_loop(
+def get_megatron_calibration_dataloader(
     tokenizer: "PreTrainedTokenizerBase",
     *,
     dataset_name: str | list[str] = "cnn_dailymail",
@@ -49,17 +50,15 @@ def get_megatron_calibration_forward_loop(
     device: torch.device | str | None = "cuda",
     apply_chat_template: bool = True,
     pack: bool = False,
-) -> Callable[[torch.nn.Module], None]:
-    """Build a Megatron-Core calibration ``forward_loop(model)``.
+) -> torch.utils.data.DataLoader:
+    """Build a DP-sharded calibration dataloader for Megatron-Core models.
 
-    Iterates a packed dataloader built via ``get_dataset_dataloader(pack=True)``
-    and drives a logits-free prefill pass through the model so activation hooks
-    fire on every layer. All kwargs except ``seq_length`` are forwarded
-    1:1 — see :func:`get_dataset_dataloader` for their semantics. ``seq_length``
-    maps to that function's ``max_sample_length``.
+    Each batch is a dict with at least ``input_ids`` and ``attention_mask`` tensors
+    on ``device``. The dataloader is suitable as the ``data_loader`` argument to
+    ``mtq.auto_quantize`` or any other API that iterates batches directly.
 
-    Returns:
-        A ``forward_loop(model)`` callable to pass into ``mtq.quantize``, ``mtp.prune``, or other such APIs.
+    All kwargs are forwarded to :func:`get_dataset_dataloader`; ``seq_length``
+    maps to that function's ``max_sample_length``.
     """
     # Deepcopy before mutating pad_token so the caller's tokenizer isn't silently changed.
     if getattr(tokenizer, "pad_token", None) is None:
@@ -68,7 +67,7 @@ def get_megatron_calibration_forward_loop(
 
     # Shard calibration data across DP ranks; amax is max-reduced across DP inside ``mtq``.
     dp_size = mpu.get_data_parallel_world_size()
-    dataloader = get_dataset_dataloader(
+    return get_dataset_dataloader(
         dataset_name=dataset_name,
         tokenizer=tokenizer,
         batch_size=batch_size,
@@ -85,6 +84,40 @@ def get_megatron_calibration_forward_loop(
         },
     )
 
+
+def get_megatron_calibration_forward_loop(
+    tokenizer: "PreTrainedTokenizerBase",
+    *,
+    dataset_name: str | list[str] = "cnn_dailymail",
+    batch_size: int = 1,
+    num_samples: int | list[int] = 512,
+    seq_length: int = 512,
+    device: torch.device | str | None = "cuda",
+    apply_chat_template: bool = True,
+    pack: bool = False,
+) -> Callable[[torch.nn.Module], None]:
+    """Build a Megatron-Core calibration ``forward_loop(model)``.
+
+    Iterates a dataloader built via :func:`get_megatron_calibration_dataloader`
+    and drives a logits-free prefill pass through the model so activation hooks
+    fire on every layer. All kwargs are forwarded 1:1 — see
+    :func:`get_megatron_calibration_dataloader` for their semantics.
+
+    Returns:
+        A ``forward_loop(model)`` callable to pass into ``mtq.quantize``,
+        ``mtp.prune``, or other such APIs.
+    """
+    dataloader = get_megatron_calibration_dataloader(
+        tokenizer,
+        dataset_name=dataset_name,
+        batch_size=batch_size,
+        num_samples=num_samples,
+        seq_length=seq_length,
+        device=device,
+        apply_chat_template=apply_chat_template,
+        pack=pack,
+    )
+
     def _forward_loop(model: torch.nn.Module) -> None:
         cp_size = mpu.get_context_parallel_world_size()
         cp_group = mpu.get_context_parallel_group()

tests/conftest.py

@@ -54,7 +54,7 @@ def pytest_addoption(parser):
     "gpu_trtllm": 60,
     "gpu_vllm": 60,
     "regression": 180,
-    "unit": 60,
+    "unit": 120 if platform.system() == "Windows" else 60,
 }