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

CHANGELOG.rst

@@ -42,6 +42,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

@@ -368,13 +368,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()
@@ -398,13 +399,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,
             )
 
@@ -456,6 +456,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
@@ -865,6 +867,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": (
@@ -895,9 +906,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:
@@ -918,12 +934,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_quant.py

@@ -519,6 +519,14 @@ def forward_backward_step(model, batch) -> None:
                 "get_auto_quantize_config() to obtain a config for mtq.quantize()."
             )
 
+    # Lazy import: avoids module-load-time circular import via plugins/megatron → algorithms.
+    try:
+        from .plugins.megatron import register_megatron_autoquant_support
+
+        register_megatron_autoquant_support()
+    except ImportError:
+        pass
+
     # Select the appropriate searcher based on method
     if method == "gradient":
         searcher = AutoQuantizeGradientSearcher()

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
@@ -44,6 +45,7 @@
 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:
@@ -779,3 +781,54 @@ 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
+
+
+_AUTOQUANT_SUPPORT_REGISTERED = False
+
+
+def register_megatron_autoquant_support() -> None:
+    """Register megatron AutoQuant hooks. Call from `auto_quantize` entry (lazy), not at module-load."""
+    global _AUTOQUANT_SUPPORT_REGISTERED
+    if _AUTOQUANT_SUPPORT_REGISTERED:
+        return
+    from ..algorithms import AutoQuantizeGradientSearcher
+
+    AutoQuantizeGradientSearcher.register_custom_support(
+        _is_supported_megatron_model,
+        _megatron_grad_ckpt_context,
+        _is_param_grad_enabled_for_megatron,
+    )
+    _AUTOQUANT_SUPPORT_REGISTERED = True
+
+
+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

tests/gpu/torch/quantization/test_gptq.py

@@ -120,6 +120,21 @@ def test_update_hessian():
     assert n_samples_accumulated == 12, f"Expected n_samples=12, got {n_samples_accumulated}"
 
 
+def test_update_hessian_zero_token_input_noops():
+    features = 4
+    hessian = torch.eye(features, dtype=torch.float32)
+    expected_hessian = hessian.clone()
+    n_samples = 7
+
+    updated_hessian, new_n_samples = update_hessian(
+        torch.empty(0, features, dtype=torch.float32), hessian, n_samples
+    )
+
+    assert updated_hessian is hessian
+    assert new_n_samples == n_samples
+    torch.testing.assert_close(hessian, expected_hessian)
+
+
 @pytest.mark.parametrize(
     ("block_size", "dim", "model_weight", "expect_weight_change"),
     [

tests/gpu_megatron/torch/quantization/plugins/test_megatron.py

@@ -28,6 +28,7 @@
 from _test_utils.torch.megatron.utils import (
     compare_amax_sync_across_expert_parallel,
     copy_weights_from_grouped_to_non_grouped,
+    get_batch,
     get_forward,
     initialize_for_megatron,
     run_mcore_inference,
@@ -54,7 +55,11 @@
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
 from modelopt.torch.quantization.nn import QuantModuleRegistry
-from modelopt.torch.quantization.plugins.megatron import _QuantTEMCoreRowParallelLinear
+from modelopt.torch.quantization.plugins.megatron import (
+    _QuantTEMCoreRowParallelLinear,
+    get_mcore_layerwise_calibration_layers,
+)
+from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
 
 try:
     from megatron.core.extensions.transformer_engine import TERowParallelLinear
@@ -696,6 +701,80 @@ def test_te_grouped_vs_sequential_quantize(dist_workers_size_4, quant_cfg):
     )
 
 
+def _test_auto_quantize_moe_ep_helper(rank, size):
+    initialize_for_megatron(
+        tensor_model_parallel_size=1,
+        expert_model_parallel_size=size,
+        seed=SEED,
+    )
+    model = _gpt_model_provider(
+        tp_size=1,
+        ep_size=size,
+        hidden_size=32,
+        num_moe_experts=4,
+        moe_grouped_gemm=False,
+        transformer_impl="modelopt",
+    )
+
+    def forward_step(model, batch):
+        input_ids, labels, position_ids, attention_mask, loss_mask = batch
+        return model.forward(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            loss_mask=loss_mask,
+        )
+
+    auto_quantize_helper(
+        model,
+        data_loader=[get_batch(model, batch_size=2) for _ in range(2)],
+        forward_step=forward_step,
+        forward_backward_step=lambda m, b: forward_step(m, b).mean().backward(),
+        quantization_formats=[mtq.NVFP4_DEFAULT_CFG, mtq.FP8_DEFAULT_CFG],
+    )
+
+
+def test_auto_quantize_moe_ep(dist_workers_size_2):
+    """auto_quantize must sum score/cost across EP ranks and pick a consistent recipe."""
+    dist_workers_size_2.run(_test_auto_quantize_moe_ep_helper)
+
+
+def test_mcore_layerwise_calibration_layers_do_not_mutate_decoder(distributed_setup_size_1):
+    if not HAS_TE:
+        pytest.skip("Transformer Engine is not installed")
+
+    initialize_for_megatron(tensor_model_parallel_size=1, seed=SEED)
+    try:
+        model = _gpt_model_provider(
+            tp_size=1,
+            hidden_size=32,
+            meta_device=True,
+            transformer_impl="modelopt",
+        )
+        decoder_layers = model.decoder.layers
+        decoder_len = len(decoder_layers)
+        output_layer = model.output_layer
+
+        discovered_layers = get_mcore_layerwise_calibration_layers(model)
+
+        assert discovered_layers is not None
+        assert len(discovered_layers) == decoder_len + 1
+        assert discovered_layers[-1] is output_layer
+        assert len(decoder_layers) == decoder_len
+        assert all(layer is not output_layer for layer in decoder_layers)
+
+        assert LayerActivationCollector.is_supported(model)
+        discovered_layers = LayerActivationCollector.get_decoder_layers(model)
+        assert discovered_layers is not None
+        assert len(discovered_layers) == decoder_len + 1
+        assert discovered_layers[-1] is output_layer
+        assert len(decoder_layers) == decoder_len
+        assert all(layer is not output_layer for layer in decoder_layers)
+    finally:
+        destroy_model_parallel()
+
+
 @pytest.mark.parametrize("ep_size", [1, 2])
 @pytest.mark.parametrize("sync_weight_amax", [True, False])
 def test_layer_sync_moe_local_experts_amax(dist_workers, ep_size, sync_weight_amax):

tests/unit/torch/quantization/test_autoquant.py

@@ -33,6 +33,7 @@
     AutoQuantizeGradientSearcher,
     QuantRecipe,
     QuantRecipeHparam,
+    _AutoQuantizeBaseSearcher,
     estimate_quant_compression,
 )
 from modelopt.torch.quantization.config import _base_disable_all, _default_disabled_quantizer_cfg
@@ -486,6 +487,39 @@ def test_data_parallel_auto_quantize(skip_on_windows):
     spawn_multiprocess_job(2, _test_data_parallel_auto_quantize, backend="gloo")
 
 
+def test_auto_quantize_budget_uses_no_quant_candidate_cost(monkeypatch):
+    class _BudgetCaptureSearcher(AutoQuantizeGradientSearcher):
+        def run_search_with_stats(self, max_weight_size, verbose=False):
+            self.max_weight_size = max_weight_size
+            return {}, True
+
+    def _raise_local_total_weight_size(modules):
+        pytest.fail("run_search should derive total weight size from candidate costs")
+
+    monkeypatch.setattr(
+        _AutoQuantizeBaseSearcher,
+        "_get_total_weight_size",
+        staticmethod(_raise_local_total_weight_size),
+    )
+
+    searcher = _BudgetCaptureSearcher()
+    searcher.reset_search()
+    searcher.model = torch.nn.Module()
+    searcher.config = {"verbose": False}
+    searcher.constraints = {"effective_bits": 8.0}
+    searcher.candidate_stats = {
+        "local_expert.quant_recipe": {
+            "formats": [QuantRecipe(mtq.NVFP4_DEFAULT_CFG), QuantRecipe(None)],
+            "scores": [1.0, 0.0],
+            "costs": [25.0, 100.0],
+        }
+    }
+
+    searcher.run_search()
+
+    assert searcher.max_weight_size == 50.0
+
+
 def test_estimate_quant_compression():
     nvfp4_affine_kv_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_AFFINE_KV_CFG)
     assert estimate_quant_compression(nvfp4_affine_kv_cfg) == 0.25