Fix TEGroupedMLP quantizer checkpoint resharding

modelopt/torch/opt/plugins/mcore_dist_checkpointing.py

@@ -181,6 +181,10 @@ def _load_extra_state_from_sharded_checkpoint(
             module, "modelopt_set_extra_state_callbacks"
         ):
             module.set_extra_state(extra_state_dict_no_prefix[key])
+    for module in model.modules():
+        post_load_extra_state = getattr(module, "modelopt_post_load_extra_state", None)
+        if callable(post_load_extra_state):
+            post_load_extra_state()
 
 
 def restore_sharded_modelopt_state(

modelopt/torch/quantization/plugins/megatron.py

@@ -110,6 +110,39 @@ def _check_nvfp4_static_tp_supported(model: torch.nn.Module) -> None:
         )
 
 
+def _initialize_grouped_weight_quantizer_state(module: torch.nn.Module) -> None:
+    """Create per-expert buffers as destinations for the subsequent checkpoint load."""
+    grouped_leaves = [
+        quantizer if isinstance(quantizer, torch.nn.Sequential) else [quantizer]
+        for quantizer in [module.weight_quantizer[idx] for idx in range(module.num_gemms)]
+    ]
+    for sibling_leaves in zip(*grouped_leaves):
+        eligible_leaves = [
+            quantizer
+            for quantizer in sibling_leaves
+            if (
+                quantizer.is_enabled
+                and not quantizer.is_mx_format
+                and not getattr(quantizer, "_dynamic", False)
+                and not getattr(quantizer, "_lsq", False)
+            )
+        ]
+        for state_name in ("_amax", "_global_amax"):
+            reference = next(
+                (
+                    state
+                    for quantizer in eligible_leaves
+                    if (state := getattr(quantizer, state_name, None)) is not None
+                ),
+                None,
+            )
+            if reference is None:
+                continue
+            for quantizer in eligible_leaves:
+                if getattr(quantizer, state_name, None) is None:
+                    quantizer.register_buffer(state_name, torch.zeros_like(reference))
+
+
 def real_quant_module_get_extra_state(self) -> dict:
     """Populating real_quantizer_state and q_tensor_state."""
     extra_state = {}
@@ -757,6 +790,9 @@ class _QuantTELayerNormColumnParallelLinear(
 
     # Quantized subclasses to support TEGroupedLinear quantization
     class _QuantMegatronTEGroupedLinear(_QuantTEGroupedLinear, _MegatronParallelLinear):
+        def modelopt_post_load_extra_state(self):
+            _initialize_grouped_weight_quantizer_state(self)
+
         def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
             # _sharded_state_dict_grouped adds _extra_state{gemm_idx} for gemm_idx:[1, num_gemms] in
             # sharded_state_dict which is same as _extra_state. The _extra_state{gemm_idx} is used for

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

@@ -15,7 +15,6 @@
 
 import copy
 import math
-import re
 import sys
 import types
 from contextlib import nullcontext
@@ -38,7 +37,9 @@
     get_batch,
     get_forward,
     initialize_for_megatron,
+    load_distributed_checkpoint,
     run_mcore_inference,
+    save_distributed_checkpoint,
     sharded_state_dict_test_helper,
 )
 from _test_utils.torch.misc import set_seed
@@ -52,6 +53,7 @@
 from megatron.core.parallel_state import (
     destroy_model_parallel,
     get_data_parallel_group,
+    get_expert_model_parallel_rank,
     get_tensor_model_parallel_group,
 )
 from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
@@ -62,10 +64,15 @@
 import modelopt
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
+from modelopt.torch.opt.plugins.mcore_dist_checkpointing import (
+    restore_sharded_modelopt_state,
+    save_sharded_modelopt_state,
+)
 from modelopt.torch.quantization.algorithms import QuantRecipe, _AutoQuantizeBaseSearcher
 from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer
 from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear
 from modelopt.torch.quantization.plugins.megatron import (
+    _initialize_grouped_weight_quantizer_state,
     _output_layer_untied,
     _QuantMegatronTEGroupedLinear,
     _QuantTEMCoreRowParallelLinear,
@@ -1090,100 +1097,207 @@ def test_te_grouped_vs_sequential_default_amax(dist_workers_size_1, quant_cfg):
     )
 
 
-def _te_grouped_expert_identity_from_sharded_state(module):
-    """Return {local_key: (global_expert_idx, num_global_experts)} for per-expert amax shards.
+def _set_te_grouped_weight_quantizer_state(model, ep_rank, num_local_experts):
+    """Give every local expert distinct quantizer state derived from its global index."""
+    for linear in model.modules():
+        if not isinstance(linear, _QuantMegatronTEGroupedLinear):
+            continue
+        for local_expert_idx in range(linear.num_gemms):
+            quantizer = linear.weight_quantizer[local_expert_idx]
+            leaves = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer]
+            for leaf in leaves:
+                amax = getattr(leaf, "_amax", None)
+                if amax is not None:
+                    amax.fill_(1.0 + ep_rank * num_local_experts + local_expert_idx)
+                global_amax = getattr(leaf, "_global_amax", None)
+                if global_amax is not None:
+                    global_amax.fill_(1.0 + ep_rank * num_local_experts + local_expert_idx)
+
+
+def _assert_te_grouped_weight_quantizer_state(model, expected_amax, expect_global_amax):
+    checked = 0
+    for linear in model.modules():
+        if not isinstance(linear, _QuantMegatronTEGroupedLinear):
+            continue
+        for local_expert_idx in range(linear.num_gemms):
+            quantizer = linear.weight_quantizer[local_expert_idx]
+            leaves = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer]
+            for leaf in leaves:
+                amax = getattr(leaf, "_amax", None)
+                assert amax is not None, (
+                    "TEGrouped per-expert weight quantizer amax was not restored"
+                )
+                checked += 1
+                assert torch.equal(amax, torch.full_like(amax, expected_amax[local_expert_idx]))
+                global_amax = getattr(leaf, "_global_amax", None)
+                if expect_global_amax:
+                    assert global_amax is not None, (
+                        "TEGrouped per-expert weight quantizer global_amax was not restored"
+                    )
+                    assert torch.equal(
+                        global_amax,
+                        torch.full_like(global_amax, expected_amax[local_expert_idx]),
+                    )
+    assert checked > 0, "no TEGrouped per-expert weight quantizer amax was checked"
 
-    The grouped linear must give each fused expert the same global identity the weights use:
-    the dict key keeps the local expert index (maps to the local buffer on restore) while the
-    ShardedTensor carries the global expert offset. Called with sharded_offsets=() so the expert
-    axis is the (only) prepended axis at index 0.
-    """
-    sharded_sd = module.sharded_state_dict(prefix="", sharded_offsets=(), metadata=None)
-    identity = {}
-    for key, sh_ten in sharded_sd.items():
-        if re.match(r"weight_quantizer\.\d+\..*_amax$", key):
-            assert sh_ten.prepend_axis_num >= 1, f"{key}: expected a prepended expert axis"
-            identity[key] = (int(sh_ten.global_offset[0]), int(sh_ten.global_shape[0]))
-    return identity
 
+def test_initialize_grouped_weight_quantizer_state_for_restore():
+    """Missing grouped state inherits the shape and dtype of a populated sibling."""
+    source = mtq.nn.StaticBlockScaleQuantizer.from_tensor_quantizer(
+        mtq.nn.TensorQuantizer(amax=torch.tensor([1.0, 2.0])), global_amax=torch.tensor(2.0)
+    )
+    target = mtq.nn.StaticBlockScaleQuantizer.from_tensor_quantizer(mtq.nn.TensorQuantizer())
+    disabled = mtq.nn.StaticBlockScaleQuantizer.from_tensor_quantizer(mtq.nn.TensorQuantizer())
+    disabled.disable()
+    mx = mtq.nn.TensorQuantizer(
+        mtq.config.QuantizerAttributeConfig(
+            num_bits=(4, 3), block_sizes={-1: 32, "type": "dynamic", "scale_bits": (8, 0)}
+        )
+    )
 
-def _test_te_grouped_sharded_state_dict_global_expert_identity_helper(
-    tp_size, ep_size, quant_cfg, rank, size
-):
-    """Per-expert quantizer amax must persist all num_global_experts across EP.
+    model = torch.nn.Module()
+    model.weight = torch.nn.Parameter(torch.empty(1))
+    model.num_gemms = 4
+    model.weight_quantizer = torch.nn.ModuleList([target, source, disabled, mx])
 
-    With EP>1 the base linear emitted ``weight_quantizer.{local_i}._amax`` at the local index with
-    no expert offset, so every rank wrote identical keys and torch_dist dedup collapsed them to a
-    single rank's experts. Assert each rank's fused experts now carry distinct global identities so
-    the union across ranks covers every global expert.
-    """
+    _initialize_grouped_weight_quantizer_state(model)
+
+    assert torch.equal(target.amax, torch.zeros_like(source.amax))
+    assert torch.equal(target.global_amax, torch.zeros_like(source.global_amax))
+    assert disabled.amax is None
+    assert disabled.global_amax is None
+    assert mx.amax is None
+    assert not hasattr(mx, "_amax")
+
+
+def _test_te_grouped_sharded_state_dict_reshard_helper(
+    save_tp_size,
+    save_ep_size,
+    load_tp_size,
+    load_ep_size,
+    quant_cfg,
+    expect_global_amax,
+    checkpoint_path,
+    rank,
+    size,
+):
+    """Round-trip TEGroupedMLP amax through a topology change."""
+    num_experts = 4
+    save_num_local_experts = num_experts // save_ep_size
     initialize_for_megatron(
-        tensor_model_parallel_size=tp_size,
-        expert_model_parallel_size=ep_size,
+        tensor_model_parallel_size=save_tp_size,
+        expert_model_parallel_size=save_ep_size,
         seed=SEED,
     )
-    num_experts = 4
-    num_local = num_experts // ep_size
 
-    te_grouped = _gpt_model_provider(
-        tp_size=tp_size,
-        ep_size=ep_size,
+    source = _gpt_model_provider(
+        tp_size=save_tp_size,
+        ep_size=save_ep_size,
         hidden_size=32,
         moe_grouped_gemm=True,
         transformer_impl="transformer_engine",
         num_moe_experts=num_experts,
     )
-    forward = get_forward(te_grouped, batch_size=8)
-    for module in te_grouped.modules():
+    forward = get_forward(source, batch_size=8)
+    for module in source.modules():
         if isinstance(module, TopKRouter):
             module.topk = module.num_experts
-    mtq.quantize(te_grouped, quant_cfg, forward)
+    mtq.quantize(source, copy.deepcopy(quant_cfg), forward)
+    _set_te_grouped_weight_quantizer_state(
+        source, get_expert_model_parallel_rank(), save_num_local_experts
+    )
+    save_distributed_checkpoint(checkpoint_path, source)
+    save_sharded_modelopt_state([source], checkpoint_path)
+    torch.distributed.barrier()
+    del source
+    destroy_model_parallel()
 
-    grouped_linears = [
-        m for m in te_grouped.modules() if isinstance(m, _QuantMegatronTEGroupedLinear)
-    ]
-    assert grouped_linears, "No grouped quant linears found"
-
-    expected_global = {rank * num_local + i for i in range(num_local)}
-    for linear in grouped_linears:
-        # Give each expert a distinct amax so a value mix-up would also be observable.
-        for i in range(linear.num_gemms):
-            wq = linear.weight_quantizer[i]
-            leaves = list(wq) if isinstance(wq, SequentialQuantizer) else [wq]
-            for leaf in leaves:
-                if hasattr(leaf, "_amax") and leaf._amax is not None:
-                    leaf._amax.fill_(1.0 + rank * num_local + i)
-
-        identity = _te_grouped_expert_identity_from_sharded_state(linear)
-        # One entry per local expert per amax buffer; dict keys keep the LOCAL index.
-        local_keys = {int(re.search(r"weight_quantizer\.(\d+)\.", k).group(1)) for k in identity}
-        assert local_keys == set(range(num_local)), (
-            f"Expected local expert keys {set(range(num_local))}, got {local_keys}"
-        )
-        # ShardedTensor global identity: this rank owns experts {rank*num_local + i}.
-        local_global = {gidx for gidx, _ in identity.values()}
-        assert local_global == expected_global, (
-            f"rank {rank}: expected global experts {expected_global}, got {local_global}"
-        )
-        assert all(total == num_experts for _, total in identity.values()), (
-            f"num_global_experts should be {num_experts}, got {identity}"
+    initialize_for_megatron(
+        tensor_model_parallel_size=load_tp_size,
+        expert_model_parallel_size=load_ep_size,
+        seed=SEED,
+    )
+    target = _gpt_model_provider(
+        tp_size=load_tp_size,
+        ep_size=load_ep_size,
+        hidden_size=32,
+        moe_grouped_gemm=True,
+        transformer_impl="transformer_engine",
+        num_moe_experts=num_experts,
+    )
+    target_models = [target]
+    restore_sharded_modelopt_state(target_models, checkpoint_path)
+    target = target_models[0]
+    load_distributed_checkpoint(checkpoint_path, target)
+    load_num_local_experts = num_experts // load_ep_size
+    expected_amax = tuple(
+        range(
+            1 + get_expert_model_parallel_rank() * load_num_local_experts,
+            1 + (get_expert_model_parallel_rank() + 1) * load_num_local_experts,
         )
-
-    # Gather the global expert indices across all EP ranks: the union must cover every expert.
-    gathered = [None] * size
-    torch.distributed.all_gather_object(gathered, sorted(expected_global))
-    union = set()
-    for part in gathered:
-        union.update(part)
-    assert union == set(range(num_experts)), (
-        f"Union of global experts across EP ranks should be {set(range(num_experts))}, got {union}"
     )
+    _assert_te_grouped_weight_quantizer_state(target, expected_amax, expect_global_amax)
 
 
-@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG])
-def test_te_grouped_sharded_state_dict_global_expert_identity(dist_workers_size_2, quant_cfg):
+@pytest.mark.parametrize(
+    (
+        "quant_cfg",
+        "expect_global_amax",
+        "save_tp_size",
+        "save_ep_size",
+        "load_tp_size",
+        "load_ep_size",
+    ),
+    [
+        pytest.param(mtq.FP8_DEFAULT_CFG, False, 1, 2, 1, 1, id="fp8-ep-downsize"),
+        pytest.param(mtq.FP8_DEFAULT_CFG, False, 1, 1, 1, 2, id="fp8-ep-upsize"),
+        pytest.param(mtq.FP8_DEFAULT_CFG, False, 1, 1, 2, 1, id="fp8-tp-upsize"),
+        pytest.param(mtq.FP8_DEFAULT_CFG, False, 2, 1, 1, 1, id="fp8-tp-downsize"),
+        pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 1, 2, 1, 1, id="nvfp4-ep-downsize"),
+        pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 1, 1, 1, 2, id="nvfp4-ep-upsize"),
+        pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 1, 1, 2, 1, id="nvfp4-tp-upsize"),
+        pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 2, 1, 1, 1, id="nvfp4-tp-downsize"),
+        pytest.param(
+            mtq.NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG,
+            True,
+            1,
+            2,
+            1,
+            1,
+            id="nvfp4-mse-ep-downsize",
+        ),
+        pytest.param(
+            mtq.NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG,
+            True,
+            1,
+            1,
+            1,
+            2,
+            id="nvfp4-mse-ep-upsize",
+        ),
+    ],
+)
+def test_te_grouped_sharded_state_dict_reshard(
+    dist_workers_size_2,
+    tmp_path,
+    save_tp_size,
+    save_ep_size,
+    load_tp_size,
+    load_ep_size,
+    quant_cfg,
+    expect_global_amax,
+):
     dist_workers_size_2.run(
-        partial(_test_te_grouped_sharded_state_dict_global_expert_identity_helper, 1, 2, quant_cfg)
+        partial(
+            _test_te_grouped_sharded_state_dict_reshard_helper,
+            save_tp_size,
+            save_ep_size,
+            load_tp_size,
+            load_ep_size,
+            quant_cfg,
+            expect_global_amax,
+            tmp_path,
+        )
     )