From 34d1ef39fb909aa2cf2b668319be00ea6e24aae8 Mon Sep 17 00:00:00 2001
From: Jennifer Chen <jennifchen@nvidia.com>
Date: Tue, 12 May 2026 20:28:12 -0700
Subject: [PATCH 1/8] te_qad_debug: per-expert TEGrouped quantizer + logits
 debug print + infra fixes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

modelopt/torch/quantization/plugins/transformer_engine.py:
  MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER=1 opts into per-gemm
  weight_quantizer_0..N-1 inside _QuantTEGroupedLinear (deepcopied from
  the shared weight_quantizer). Lets TEGroupedMLP recover per-expert
  amax granularity, matching SequentialMLP's default behavior.

modelopt/torch/distill/plugins/megatron.py:
  LogitsKLLoss.forward prints student/teacher logit stats (mean/std/
  min/max/shape) on rank 0 each call. Diagnostic for the QAD loss-spike
  investigation — confirms which spec produces which logits without
  changing the KL math.

tests/gpu_megatron/torch/quantization/plugins/test_megatron.py:
  New test_te_grouped_vs_sequential_default_amax + ..._default_loss
  cover the structural amax asymmetry between TEGroupedMLP and
  SequentialMLP (TEGrouped per-linear amax = max-over-Sequential-experts
  amax) and a finiteness sanity check on the resulting quant error.

tools/launcher/common/service_utils.sh:
  - Fall back to SLURM_PROCID / SLURM_LOCALID when PMIX_*/OMPI_* are
    unset, so `[[ "$mpi_local_rank" -eq 0 ]]` doesn't silently pass on
    every rank under plain srun.
  - util_install_extra_dep: per-node marker so concurrent ranks wait
    for rank 0 to finish installing (concurrent pip on a shared FS
    leaves a broken state); also installs nvidia-resiliency-ext.

Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
---
 modelopt/torch/distill/plugins/megatron.py    |  16 +++
 .../plugins/transformer_engine.py             |  27 +++-
 .../quantization/plugins/test_megatron.py     | 126 ++++++++++++++++++
 tools/launcher/common/service_utils.sh        |   2 +
 4 files changed, 166 insertions(+), 5 deletions(-)

diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py
index 7e81c21a462..4c2d98eb6a5 100644
--- a/modelopt/torch/distill/plugins/megatron.py
+++ b/modelopt/torch/distill/plugins/megatron.py
@@ -318,6 +318,22 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor:
         """
         predictions, targets = self.pre_forward(predictions, targets)
 
+        if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
+            if not hasattr(self, "_dbg_call"):
+                self._dbg_call = 0
+            with torch.no_grad():
+                s = predictions.float()
+                t = targets.float()
+                print(
+                    f"[LogitsKLLoss call={self._dbg_call}] "
+                    f"student: mean={s.mean().item():.5f} std={s.std().item():.5f} "
+                    f"min={s.min().item():.3f} max={s.max().item():.3f} shape={tuple(predictions.shape)} | "
+                    f"teacher: mean={t.mean().item():.5f} std={t.std().item():.5f} "
+                    f"min={t.min().item():.3f} max={t.max().item():.3f}",
+                    flush=True,
+                )
+            self._dbg_call += 1
+
         # Division by temp should happen prior to finding max for both student and teacher.
         output_teacher = targets.float() / self._temperature
         output_student = predictions.float() / self._temperature
diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py
index e670141f79a..7e675197e4c 100644
--- a/modelopt/torch/quantization/plugins/transformer_engine.py
+++ b/modelopt/torch/quantization/plugins/transformer_engine.py
@@ -15,7 +15,11 @@
 
 """Support quantization for Transformer Engine layers."""
 
+import copy
+import os
 import inspect
+import copy
+import os
 import warnings
 
 import torch
@@ -33,6 +37,11 @@
 _TE_VERSION = Version(te.__version__)
 
 
+def _per_expert_weight_quantizer_enabled() -> bool:
+    """Opt-in MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER=1: per-gemm weight_quantizer in TEGroupedLinear."""
+    return os.environ.get("MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER", "0") == "1"
+
+
 def _assert_te_fp8_enabled():
     """Check if Transformer Engine FP8 autocast is enabled and raise error if so."""
     try:
@@ -137,8 +146,10 @@ 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.
+        self._per_expert_weight_quantizer = _per_expert_weight_quantizer_enabled()
+        if self._per_expert_weight_quantizer:
+            for i in range(self.num_gemms):
+                self.add_module(f"weight_quantizer_{i}", copy.deepcopy(self.weight_quantizer))
 
     def modelopt_post_restore(self, prefix: str = ""):
         # GroupedMLP stores the weights as weight0, weight1, etc. To run post_restore in order to
@@ -150,12 +161,17 @@ def modelopt_post_restore(self, prefix: str = ""):
         # Remove self.weight after post_restore.
         delattr(self, "weight")
 
+    def _get_weight_quantizer(self, gemm_idx: int):
+        if getattr(self, "_per_expert_weight_quantizer", False):
+            return getattr(self, f"weight_quantizer_{gemm_idx}")
+        return self.weight_quantizer
+
     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._get_weight_quantizer(i)
 
     @staticmethod
     def te_grouped_quantized_linear_fn(package, func_name, self, *args):
@@ -182,8 +198,9 @@ 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])
+        for gemm_idx in range(num_gemms):
+            pos = weights_start + gemm_idx
+            new_args[pos] = self._get_weight_quantizer(gemm_idx)(args[pos])
         output = getattr(package, func_name)(*new_args)
         # TE 2.15+ returns `(out, new_workspaces)`; TE <= 2.14 returns just `out`.
         # Only the activation tensor participates in output quantization.
diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
index 35cb967deac..069d2195664 100644
--- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
+++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import copy
+import math
 from contextlib import nullcontext
 from functools import partial
 
@@ -712,6 +713,131 @@ def test_te_grouped_vs_sequential_quantize(dist_workers_size_4, quant_cfg):
     )
 
 
+def _test_te_grouped_vs_sequential_default_amax_helper(tp_size, ep_size, quant_cfg, rank, size):
+    """TEGrouped per-linear amax should equal max-over-Sequential-experts 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
+
+    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_amax = getattr(te_mlp, linear_name).weight_quantizer.amax
+            assert te_amax is not None and te_amax.numel() == 1
+
+            seq_amaxes = torch.stack([
+                getattr(expert, linear_name).weight_quantizer.amax.view(())
+                for expert in seq_mlp.local_experts
+            ])
+            seq_max = seq_amaxes.max()
+
+            assert torch.allclose(te_amax.view(()), seq_max, atol=1e-5, rtol=1e-5), (
+                f"TEGrouped per-linear amax ({te_amax.item()}) != "
+                f"max-over-Sequential-experts ({seq_max.item()}) for {linear_name}"
+            )
+
+            if (seq_amaxes.max() - seq_amaxes.min()).item() > 1e-5:
+                saw_per_expert_divergence = True
+
+    assert saw_per_expert_divergence, (
+        "Expected per-expert weight amax to diverge across SequentialMLP experts."
+    )
+
+
+@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG])
+def test_te_grouped_vs_sequential_default_amax(dist_workers_size_4, quant_cfg):
+    dist_workers_size_4.run(
+        partial(_test_te_grouped_vs_sequential_default_amax_helper, 1, 2, quant_cfg)
+    )
+
+
+def _test_te_grouped_vs_sequential_default_loss_helper(tp_size, ep_size, quant_cfg, rank, size):
+    """TEGrouped quantized output should diverge from BF16 reference 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)
+    )
+
+
 @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):
diff --git a/tools/launcher/common/service_utils.sh b/tools/launcher/common/service_utils.sh
index 50bd530ed3e..917d48bdae1 100755
--- a/tools/launcher/common/service_utils.sh
+++ b/tools/launcher/common/service_utils.sh
@@ -20,6 +20,8 @@ native_mpi_local_rank=$OMPI_COMM_WORLD_LOCAL_RANK
 # Works with Slurm launching with `--mpi=pmix`
 mpi_rank=${PMIX_RANK:-${native_mpi_rank:-${SLURM_PROCID:-0}}}
 mpi_local_rank=${PMIX_LOCAL_RANK:-${native_mpi_local_rank:-${SLURM_LOCALID:-0}}}
+mpi_rank=${PMIX_RANK:-${native_mpi_rank:-${SLURM_PROCID:-0}}}
+mpi_local_rank=${PMIX_LOCAL_RANK:-${native_mpi_local_rank:-${SLURM_LOCALID:-0}}}
 
 FAIL=0
 FAIL_EXIT=0

From e36213ae3eec8dd916285a72b3f13a36146890c7 Mon Sep 17 00:00:00 2001
From: Jennifer Chen <jennifchen@nvidia.com>
Date: Wed, 27 May 2026 15:47:21 -0700
Subject: [PATCH 2/8] revert logging

Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
---
 modelopt/torch/distill/plugins/megatron.py | 16 ----------------
 1 file changed, 16 deletions(-)

diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py
index 4c2d98eb6a5..7e81c21a462 100644
--- a/modelopt/torch/distill/plugins/megatron.py
+++ b/modelopt/torch/distill/plugins/megatron.py
@@ -318,22 +318,6 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor:
         """
         predictions, targets = self.pre_forward(predictions, targets)
 
-        if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
-            if not hasattr(self, "_dbg_call"):
-                self._dbg_call = 0
-            with torch.no_grad():
-                s = predictions.float()
-                t = targets.float()
-                print(
-                    f"[LogitsKLLoss call={self._dbg_call}] "
-                    f"student: mean={s.mean().item():.5f} std={s.std().item():.5f} "
-                    f"min={s.min().item():.3f} max={s.max().item():.3f} shape={tuple(predictions.shape)} | "
-                    f"teacher: mean={t.mean().item():.5f} std={t.std().item():.5f} "
-                    f"min={t.min().item():.3f} max={t.max().item():.3f}",
-                    flush=True,
-                )
-            self._dbg_call += 1
-
         # Division by temp should happen prior to finding max for both student and teacher.
         output_teacher = targets.float() / self._temperature
         output_student = predictions.float() / self._temperature

From 60326ce6328385fb73a6355b12bc42e3fbe32777 Mon Sep 17 00:00:00 2001
From: Jennifer Chen <jennifchen@nvidia.com>
Date: Wed, 27 May 2026 15:57:42 -0700
Subject: [PATCH 3/8] te_per_expert: dedup imports + add TP-change post-restore
 caveat
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- transformer_engine.py: dedup `import copy`/`import os` left over from the
  rebase, sort the four imports alphabetically.
- transformer_engine.py: comment near the per-expert weight_quantizer setup
  explaining that base modelopt_post_restore won't re-calibrate the
  weight_quantizer_{i} modules, so save/restore is only safe when TP/EP is
  unchanged. Per-channel _amax shape depends on the TP-sliced output dim.
- service_utils.sh: drop the duplicated mpi_rank / mpi_local_rank
  re-assignments — main already carries the SLURM fallback, the extra two
  lines were leftover rebase noise.

Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
---
 modelopt/torch/quantization/plugins/transformer_engine.py | 8 ++++++--
 tools/launcher/common/service_utils.sh                    | 2 --
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py
index 7e675197e4c..3e06b08f6ec 100644
--- a/modelopt/torch/quantization/plugins/transformer_engine.py
+++ b/modelopt/torch/quantization/plugins/transformer_engine.py
@@ -16,9 +16,7 @@
 """Support quantization for Transformer Engine layers."""
 
 import copy
-import os
 import inspect
-import copy
 import os
 import warnings
 
@@ -150,6 +148,12 @@ def _setup(self):
         if self._per_expert_weight_quantizer:
             for i in range(self.num_gemms):
                 self.add_module(f"weight_quantizer_{i}", copy.deepcopy(self.weight_quantizer))
+            # NOTE: base modelopt_post_restore only re-calibrates self.weight_quantizer.
+            # The per-expert weight_quantizer_{i} _amax buffers ride through state_dict
+            # unchanged, which works when TP/EP is identical between save and restore.
+            # If parallelism changes (per-channel _amax shape depends on the TP-sliced
+            # output dim), the loaded _amax won't match the new weight shape — needs a
+            # custom modelopt_post_restore that re-runs max_calibrate per expert.
 
     def modelopt_post_restore(self, prefix: str = ""):
         # GroupedMLP stores the weights as weight0, weight1, etc. To run post_restore in order to
diff --git a/tools/launcher/common/service_utils.sh b/tools/launcher/common/service_utils.sh
index 917d48bdae1..50bd530ed3e 100755
--- a/tools/launcher/common/service_utils.sh
+++ b/tools/launcher/common/service_utils.sh
@@ -20,8 +20,6 @@ native_mpi_local_rank=$OMPI_COMM_WORLD_LOCAL_RANK
 # Works with Slurm launching with `--mpi=pmix`
 mpi_rank=${PMIX_RANK:-${native_mpi_rank:-${SLURM_PROCID:-0}}}
 mpi_local_rank=${PMIX_LOCAL_RANK:-${native_mpi_local_rank:-${SLURM_LOCALID:-0}}}
-mpi_rank=${PMIX_RANK:-${native_mpi_rank:-${SLURM_PROCID:-0}}}
-mpi_local_rank=${PMIX_LOCAL_RANK:-${native_mpi_local_rank:-${SLURM_LOCALID:-0}}}
 
 FAIL=0
 FAIL_EXIT=0

From ea0ab1bc09668d4b0fb4af519db9166280d3fc16 Mon Sep 17 00:00:00 2001
From: Jennifer Chen <jennifchen@nvidia.com>
Date: Wed, 27 May 2026 16:07:33 -0700
Subject: [PATCH 4/8] sync weight_quantizer_{i} during post restore

Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
---
 .../plugins/transformer_engine.py             | 26 ++++++++++++++-----
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py
index 3e06b08f6ec..e3a87927fd3 100644
--- a/modelopt/torch/quantization/plugins/transformer_engine.py
+++ b/modelopt/torch/quantization/plugins/transformer_engine.py
@@ -29,7 +29,7 @@
 
 from modelopt.torch.quantization.utils import replace_function
 
-from ..nn import QuantModuleRegistry
+from ..nn import QuantModuleRegistry, SequentialQuantizer
 from .custom import _ParallelLinear
 
 _TE_VERSION = Version(te.__version__)
@@ -148,12 +148,6 @@ def _setup(self):
         if self._per_expert_weight_quantizer:
             for i in range(self.num_gemms):
                 self.add_module(f"weight_quantizer_{i}", copy.deepcopy(self.weight_quantizer))
-            # NOTE: base modelopt_post_restore only re-calibrates self.weight_quantizer.
-            # The per-expert weight_quantizer_{i} _amax buffers ride through state_dict
-            # unchanged, which works when TP/EP is identical between save and restore.
-            # If parallelism changes (per-channel _amax shape depends on the TP-sliced
-            # output dim), the loaded _amax won't match the new weight shape — needs a
-            # custom modelopt_post_restore that re-runs max_calibrate per expert.
 
     def modelopt_post_restore(self, prefix: str = ""):
         # GroupedMLP stores the weights as weight0, weight1, etc. To run post_restore in order to
@@ -165,6 +159,24 @@ def modelopt_post_restore(self, prefix: str = ""):
         # Remove self.weight after post_restore.
         delattr(self, "weight")
 
+        # Base post_restore only re-calibrates self.weight_quantizer; the per-expert
+        # weight_quantizer_{i} also need re-calibration so a TP/EP change between save
+        # and restore produces correctly shaped per-channel _amax. Mirror base behavior:
+        # only re-calibrate quantizers whose loaded state had _amax (skip unused ones).
+        if getattr(self, "_per_expert_weight_quantizer", False):
+            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._get_weight_quantizer(i)
+                q = wq_i[0] if isinstance(wq_i, SequentialQuantizer) else wq_i
+                if not hasattr(q, "_amax"):
+                    continue
+                wq_i.reset_amax()
+                max_calibrate(wq_i, lambda wq, w=weight_i: wq(w), distributed_sync=False)
+
     def _get_weight_quantizer(self, gemm_idx: int):
         if getattr(self, "_per_expert_weight_quantizer", False):
             return getattr(self, f"weight_quantizer_{gemm_idx}")

From c333f773a6e2b7b3bf0d0fc70c4ddd5f0df319d6 Mon Sep 17 00:00:00 2001
From: Jenny Chen <jennifchen@nvidia.com>
Date: Tue, 30 Jun 2026 13:34:21 -0700
Subject: [PATCH 5/8] Apply suggestion from @realAsma

Co-authored-by: realAsma <86726418+realAsma@users.noreply.github.com>
Signed-off-by: Jenny Chen <jc4686@columbia.edu>
---
 modelopt/torch/quantization/plugins/transformer_engine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py
index e3a87927fd3..e9b1ecb982b 100644
--- a/modelopt/torch/quantization/plugins/transformer_engine.py
+++ b/modelopt/torch/quantization/plugins/transformer_engine.py
@@ -147,7 +147,7 @@ def _setup(self):
         self._per_expert_weight_quantizer = _per_expert_weight_quantizer_enabled()
         if self._per_expert_weight_quantizer:
             for i in range(self.num_gemms):
-                self.add_module(f"weight_quantizer_{i}", copy.deepcopy(self.weight_quantizer))
+                self.add_module(f"weight_quantizer_{i}", TensorQuantizer())
 
     def modelopt_post_restore(self, prefix: str = ""):
         # GroupedMLP stores the weights as weight0, weight1, etc. To run post_restore in order to

From 16714c2d5e82beb652cc10fbc93730b6176eafc7 Mon Sep 17 00:00:00 2001
From: Jennifer Chen <jennifchen@nvidia.com>
Date: Wed, 1 Jul 2026 11:18:54 -0700
Subject: [PATCH 6/8] feat: per-expert weight quantizer for TEGroupedMLP via
 GroupedQuantizer

Fused TEGroupedLinear experts previously shared a single weight quantizer.
Add GroupedQuantizer (an nn.ModuleList of TensorQuantizers) and have
_QuantTEGroupedLinear build one quantizer per expert in _setup, so each
fused expert gets an independent weight amax by default, matching
SequentialMLP. This replaces the MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER
env-var gate.

Storing the quantizers in a ModuleList named weight_quantizer surfaces them
as weight_quantizer.{i}, which the fused-experts name normalizer maps to
*weight_quantizer so the stock quant configs reach every expert. Teach the
calibration/sync paths (model_calib DP/EP/TP sync, megatron NVFP4-static TP
check, custom restore probe) to recurse into the container so per-expert
amaxes sync correctly.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
---
 modelopt/torch/quantization/config.py         |  4 +-
 modelopt/torch/quantization/model_calib.py    | 18 +++--
 .../nn/modules/tensor_quantizer.py            | 65 +++++++++++++++++++
 modelopt/torch/quantization/plugins/custom.py | 20 +++++-
 .../torch/quantization/plugins/megatron.py    | 10 ++-
 .../plugins/transformer_engine.py             | 60 +++++++----------
 6 files changed, 129 insertions(+), 48 deletions(-)

diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py
index 344292264fa..a98d008ddaf 100644
--- a/modelopt/torch/quantization/config.py
+++ b/modelopt/torch/quantization/config.py
@@ -852,8 +852,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 24984179791..859c0b74774 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -41,7 +41,13 @@
 
 from .calib import MseCalibrator, NVFP4MSECalibrator, _Calibrator
 from .conversion import create_and_replace_svdquant_linear_on_the_fly, set_quantizer_by_cfg_context
-from .nn import NVFP4StaticQuantizer, QuantModule, SequentialQuantizer, TensorQuantizer
+from .nn import (
+    GroupedQuantizer,
+    NVFP4StaticQuantizer,
+    QuantModule,
+    SequentialQuantizer,
+    TensorQuantizer,
+)
 from .utils import (
     SHARED_PATTERNS,
     SharedWeightGlobalAmaxState,
@@ -212,7 +218,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
@@ -316,12 +322,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
@@ -335,7 +341,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
@@ -357,7 +363,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 c50804dd2a9..c625b3380d4 100644
--- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
+++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
@@ -65,6 +65,7 @@
 from ..functional import normalized_hadamard_transform
 
 __all__ = [
+    "GroupedQuantizer",
     "HardDisabledTensorQuantizer",
     "NVFP4StaticQuantizer",
     "SequentialQuantizer",
@@ -1594,3 +1595,67 @@ 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 __getattr__(self, name):
+        """Delegate property reads to the first member and method calls to all members."""
+        if name in self._delegated_properties:
+            return getattr(self[0], name)
+
+        if name in self._delegated_methods:
+
+            def method_wrapper(*args, **kwargs):
+                return [getattr(quantizer, name)(*args, **kwargs) for quantizer in self]
+
+            return method_wrapper
+
+        return super().__getattr__(name)
+
+    def __setattr__(self, name, value):
+        if name in self._delegated_properties:
+            for quantizer in self:
+                setattr(quantizer, name, value)
+        else:
+            super().__setattr__(name, value)
+
+    def set_from_attribute_config(self, attributes):
+        """Set the attributes of contained quantizers; a single config broadcasts to all."""
+        if not isinstance(attributes, (list, tuple)):
+            attributes = [attributes] * len(self)
+        for attribute, quantizer in zip(attributes, self):
+            quantizer.set_from_attribute_config(attribute)
+
+    def get_modelopt_state(self) -> dict[str, Any]:
+        """Get meta state to be saved in checkpoint."""
+        return {"num_quantizers": len(self), "is_grouped_quantizer": True}
diff --git a/modelopt/torch/quantization/plugins/custom.py b/modelopt/torch/quantization/plugins/custom.py
index f480d245daa..d16596908c9 100644
--- a/modelopt/torch/quantization/plugins/custom.py
+++ b/modelopt/torch/quantization/plugins/custom.py
@@ -24,7 +24,13 @@
 
 from modelopt.torch.utils.distributed import ParallelState
 
-from ..nn import NVFP4StaticQuantizer, QuantModule, SequentialQuantizer, TensorQuantizer
+from ..nn import (
+    GroupedQuantizer,
+    NVFP4StaticQuantizer,
+    QuantModule,
+    SequentialQuantizer,
+    TensorQuantizer,
+)
 from ..nn.modules.quant_linear import _QuantLinear
 from ..utils import multi_context, replace_function
 
@@ -134,11 +140,19 @@ def _check_unsupported_states(quantizer: TensorQuantizer):
 
         def _has_state(quantizer, name):
             # Handling for SequentialQuantizer
-            quantizer = quantizer[0] if isinstance(quantizer, SequentialQuantizer) else quantizer
+            quantizer = (
+                quantizer[0]
+                if isinstance(quantizer, (SequentialQuantizer, GroupedQuantizer))
+                else quantizer
+            )
             return hasattr(quantizer, name)
 
         def _has_complete_static_nvfp4_weight_state(quantizer, weight):
-            quantizer = quantizer[0] if isinstance(quantizer, SequentialQuantizer) else quantizer
+            quantizer = (
+                quantizer[0]
+                if isinstance(quantizer, (SequentialQuantizer, GroupedQuantizer))
+                else quantizer
+            )
             if not isinstance(quantizer, NVFP4StaticQuantizer):
                 return False
             amax = getattr(quantizer, "_amax", None)
diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py
index e18e3cd064b..0289a94e483 100644
--- a/modelopt/torch/quantization/plugins/megatron.py
+++ b/modelopt/torch/quantization/plugins/megatron.py
@@ -40,7 +40,13 @@
 from modelopt.torch.utils.distributed import ParallelState
 
 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
@@ -87,7 +93,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):
diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py
index e9b1ecb982b..8571eb7eb20 100644
--- a/modelopt/torch/quantization/plugins/transformer_engine.py
+++ b/modelopt/torch/quantization/plugins/transformer_engine.py
@@ -17,7 +17,6 @@
 
 import copy
 import inspect
-import os
 import warnings
 
 import torch
@@ -29,17 +28,12 @@
 
 from modelopt.torch.quantization.utils import replace_function
 
-from ..nn import QuantModuleRegistry, SequentialQuantizer
+from ..nn import GroupedQuantizer, QuantModuleRegistry, SequentialQuantizer
 from .custom import _ParallelLinear
 
 _TE_VERSION = Version(te.__version__)
 
 
-def _per_expert_weight_quantizer_enabled() -> bool:
-    """Opt-in MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER=1: per-gemm weight_quantizer in TEGroupedLinear."""
-    return os.environ.get("MODELOPT_TEGROUPED_PER_EXPERT_QUANTIZER", "0") == "1"
-
-
 def _assert_te_fp8_enabled():
     """Check if Transformer Engine FP8 autocast is enabled and raise error if so."""
     try:
@@ -144,10 +138,12 @@ def _setup(self):
         # Remove self.weight after setup.
         delattr(self, "weight")
 
-        self._per_expert_weight_quantizer = _per_expert_weight_quantizer_enabled()
-        if self._per_expert_weight_quantizer:
-            for i in range(self.num_gemms):
-                self.add_module(f"weight_quantizer_{i}", TensorQuantizer())
+        # 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))
+        )
 
     def modelopt_post_restore(self, prefix: str = ""):
         # GroupedMLP stores the weights as weight0, weight1, etc. To run post_restore in order to
@@ -159,35 +155,29 @@ def modelopt_post_restore(self, prefix: str = ""):
         # Remove self.weight after post_restore.
         delattr(self, "weight")
 
-        # Base post_restore only re-calibrates self.weight_quantizer; the per-expert
-        # weight_quantizer_{i} also need re-calibration so a TP/EP change between save
-        # and restore produces correctly shaped per-channel _amax. Mirror base behavior:
-        # only re-calibrate quantizers whose loaded state had _amax (skip unused ones).
-        if getattr(self, "_per_expert_weight_quantizer", False):
-            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._get_weight_quantizer(i)
-                q = wq_i[0] if isinstance(wq_i, SequentialQuantizer) else wq_i
-                if not hasattr(q, "_amax"):
-                    continue
-                wq_i.reset_amax()
-                max_calibrate(wq_i, lambda wq, w=weight_i: wq(w), distributed_sync=False)
-
-    def _get_weight_quantizer(self, gemm_idx: int):
-        if getattr(self, "_per_expert_weight_quantizer", False):
-            return getattr(self, f"weight_quantizer_{gemm_idx}")
-        return self.weight_quantizer
+        # Base post_restore only re-calibrates the first member; the other per-expert
+        # quantizers also need re-calibration so a TP/EP change between save and restore
+        # produces correctly shaped per-channel _amax. Mirror base behavior: only
+        # re-calibrate quantizers whose loaded state had _amax (skip unused ones).
+        from modelopt.torch.quantization.model_calib import max_calibrate
+
+        for i in range(self.num_gemms):
+            weight_i = getattr(self, f"weight{i}", None)
+            if weight_i is None:
+                continue
+            wq_i = self.weight_quantizer[i]
+            q = wq_i[0] if isinstance(wq_i, SequentialQuantizer) else wq_i
+            if not hasattr(q, "_amax"):
+                continue
+            wq_i.reset_amax()
+            max_calibrate(wq_i, lambda wq, w=weight_i: wq(w), distributed_sync=False)
 
     def iter_weights_for_calibration(self):
         """Yield ``(weight_i, weight_quantizer)`` for each of the ``num_gemms`` grouped weights."""
         for i in range(self.num_gemms):
             weight_i = getattr(self, f"weight{i}", None)
             if weight_i is not None:
-                yield weight_i, self._get_weight_quantizer(i)
+                yield weight_i, self.weight_quantizer[i]
 
     @staticmethod
     def te_grouped_quantized_linear_fn(package, func_name, self, *args):
@@ -216,7 +206,7 @@ def te_grouped_quantized_linear_fn(package, func_name, self, *args):
         new_args[inp_pos] = self.input_quantizer(args[inp_pos])
         for gemm_idx in range(num_gemms):
             pos = weights_start + gemm_idx
-            new_args[pos] = self._get_weight_quantizer(gemm_idx)(args[pos])
+            new_args[pos] = self.weight_quantizer[gemm_idx](args[pos])
         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.

From 096d430e4e6e97a2d3608b77a2d048b153adab70 Mon Sep 17 00:00:00 2001
From: Jennifer Chen <jennifchen@nvidia.com>
Date: Wed, 1 Jul 2026 11:19:05 -0700
Subject: [PATCH 7/8] test: per-expert TEGroupedMLP quantization assertions

Update test_te_grouped_vs_sequential_default_amax to assert each fused
expert's weight amax matches its SequentialMLP counterpart (and that they
diverge across experts) now that TEGroupedMLP quantizes per-expert by
default, instead of comparing a single shared amax. Drop the
sync_expert_weight_amax=True override in test_te_grouped_vs_sequential_quantize
so both models quantize per-expert and produce identical outputs.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
---
 .../quantization/plugins/test_megatron.py     | 77 ++++++++++++-------
 1 file changed, 49 insertions(+), 28 deletions(-)

diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
index 069d2195664..1c96cd93db5 100644
--- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
+++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
@@ -690,10 +690,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)
@@ -714,7 +714,8 @@ def test_te_grouped_vs_sequential_quantize(dist_workers_size_4, quant_cfg):
 
 
 def _test_te_grouped_vs_sequential_default_amax_helper(tp_size, ep_size, quant_cfg, rank, size):
-    """TEGrouped per-linear amax should equal max-over-Sequential-experts under default sync=False."""
+    """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,
@@ -722,14 +723,22 @@ def _test_te_grouped_vs_sequential_default_amax_helper(tp_size, ep_size, quant_c
     )
 
     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,
+        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",
+        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)
 
@@ -750,25 +759,29 @@ def _test_te_grouped_vs_sequential_default_amax_helper(tp_size, ep_size, quant_c
     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_amax = getattr(te_mlp, linear_name).weight_quantizer.amax
-            assert te_amax is not None and te_amax.numel() == 1
-
-            seq_amaxes = torch.stack([
-                getattr(expert, linear_name).weight_quantizer.amax.view(())
-                for expert in seq_mlp.local_experts
-            ])
-            seq_max = seq_amaxes.max()
-
-            assert torch.allclose(te_amax.view(()), seq_max, atol=1e-5, rtol=1e-5), (
-                f"TEGrouped per-linear amax ({te_amax.item()}) != "
-                f"max-over-Sequential-experts ({seq_max.item()}) for {linear_name}"
+            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)}"
             )
 
-            if (seq_amaxes.max() - seq_amaxes.min()).item() > 1e-5:
+            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 SequentialMLP experts."
+        "Expected per-expert weight amax to diverge across experts (proves no cross-expert sharing)."
     )
 
 
@@ -780,7 +793,7 @@ def test_te_grouped_vs_sequential_default_amax(dist_workers_size_4, 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 reference more than SequentialMLP under default sync=False."""
+    """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,
@@ -788,14 +801,22 @@ def _test_te_grouped_vs_sequential_default_loss_helper(tp_size, ep_size, quant_c
     )
 
     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,
+        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",
+        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)
 

From 16c2f0b909997fa4a35799409fbb563b94301de1 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Mon, 6 Jul 2026 14:15:31 -0700
Subject: [PATCH 8/8] Compile grouped expert quantization loop

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 CHANGELOG.rst                                 |  1 +
 .../nn/modules/tensor_quantizer.py            |  4 +
 .../torch/quantization/plugins/megatron.py    | 10 ++-
 .../plugins/transformer_engine.py             | 41 ++++++++-
 .../quantization/plugins/test_megatron.py     | 89 ++++++++++++++++++-
 .../quantization/test_tensor_quantizer_cpu.py | 15 ++++
 6 files changed, 154 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 5002918175d..9db8df9d4a0 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -17,6 +17,7 @@ Changelog
 
 **New Features**
 
+- Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
 - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 <https://arxiv.org/abs/2605.18810>`_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change).
 - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred.
 - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``.
diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
index c625b3380d4..87f7f06a16c 100644
--- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
+++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py
@@ -1628,6 +1628,10 @@ def __init__(self, *quantizers: "TensorQuantizer | SequentialQuantizer"):
             "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:
diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py
index 9cd63200fac..66f43801eca 100644
--- a/modelopt/torch/quantization/plugins/megatron.py
+++ b/modelopt/torch/quantization/plugins/megatron.py
@@ -15,6 +15,7 @@
 
 """Support quantization for megatron linear layers."""
 
+import re
 import types
 from contextlib import contextmanager
 from typing import Any
@@ -703,8 +704,13 @@ def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
             return super()._load_from_state_dict(filtered_state_dict, prefix, *args, **kwargs)
 
         def _process_quantizer_amax(self, k, v, quantizer_state_dict):
-            assert v.numel() == 1, "TEGroupedLinear only supports per-tensor quantization"
-            quantizer_state_dict[k] = v.view(-1)
+            # Per-expert quantizers have independent checkpoint keys. Preserve their native
+            # scalar, channel, or block shape instead of flattening them through the legacy
+            # single-quantizer path.
+            if re.match(r"weight_quantizer\.\d+\..+_amax$", k):
+                quantizer_state_dict[k] = v
+            else:
+                quantizer_state_dict[k] = v.view(-1) if v.numel() == 1 else v
 
     @QuantModuleRegistry.register(
         {TEColumnParallelGroupedLinear: "megatron_TEColumnParallelGroupedLinear"}
diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py
index 8571eb7eb20..f411715519c 100644
--- a/modelopt/torch/quantization/plugins/transformer_engine.py
+++ b/modelopt/torch/quantization/plugins/transformer_engine.py
@@ -17,6 +17,7 @@
 
 import copy
 import inspect
+import os
 import warnings
 
 import torch
@@ -33,6 +34,8 @@
 
 _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."""
@@ -49,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
@@ -145,6 +155,21 @@ def _setup(self):
             *(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
         # initialize the quantizer states, self.weight is used to extract shape, dtype etc. Assigning
@@ -204,9 +229,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 gemm_idx in range(num_gemms):
-            pos = weights_start + gemm_idx
-            new_args[pos] = self.weight_quantizer[gemm_idx](args[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
+        )
+        if use_compiled_loop:
+            quantized_weights = self._compiled_weight_quantizer_loop(*weights)
+        else:
+            quantized_weights = tuple(
+                self.weight_quantizer[gemm_idx](weight) for gemm_idx, weight in enumerate(weights)
+            )
+        for gemm_idx, quantized_weight in enumerate(quantized_weights):
+            new_args[weights_start + gemm_idx] = quantized_weight
         output = getattr(package, func_name)(*new_args)
         # TE 2.15+ returns `(out, new_workspaces)`; TE <= 2.14 returns just `out`.
         # Only the activation tensor participates in output quantization.
diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
index 3214cb46542..284cdc488ac 100644
--- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
+++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
@@ -59,9 +59,13 @@
 from modelopt.torch.quantization.algorithms import QuantRecipe, _AutoQuantizeBaseSearcher
 from modelopt.torch.quantization.nn import QuantModuleRegistry
 from modelopt.torch.quantization.plugins.megatron import (
+    _QuantMegatronTEGroupedLinear,
     _QuantTEMCoreRowParallelLinear,
     get_mcore_layerwise_calibration_layers,
 )
+from modelopt.torch.quantization.plugins.transformer_engine import (
+    _COMPILE_TEGROUPED_WEIGHT_LOOP_ENV,
+)
 from modelopt.torch.quantization.utils import is_quantized_linear
 from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
 
@@ -758,6 +762,88 @@ 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)."""
@@ -903,7 +989,7 @@ def test_te_grouped_vs_sequential_default_loss(dist_workers_size_4, quant_cfg):
         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,
@@ -1112,6 +1198,7 @@ def _test_mcore_layerwise_calibration_layers_do_not_mutate_decoder(rank, size):
 def test_mcore_layerwise_calibration_layers_do_not_mutate_decoder(dist_workers_size_1):
     dist_workers_size_1.run(_test_mcore_layerwise_calibration_layers_do_not_mutate_decoder)
 
+
 @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):
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))