Fix gradient reduction issue when EP=1, EP=TP, and EGTP != GTP

megatron/core/distributed/param_and_grad_buffer.py

@@ -894,8 +894,9 @@ def group_params_for_buffers(
     - param_dtype: storage dtype (torch.uint8 for FP8/NVFP4 parameters, else param.dtype).
     - grad_dtype: gradient reduction dtype (torch.float if grad_reduce_in_fp32, else param.dtype).
     - is_expert_parallel: whether the parameter uses the expert topology (param.allreduce == False),
-      which requires a separate buffer for the expert data-parallel group.  This is true for experts
-      when expert-parallelism > 1 or expert-tensor-parallelism != tensor-parallelism.
+      which requires a separate buffer for the expert data-parallel group. This is true for experts
+      when expert-parallelism > 1, expert-tensor-parallelism != tensor-parallelism, or expert-GTP
+      != GTP.
 
     The param_indices track each parameter's position among same-dtype params (using
     the "fake" high-precision dtype for FP8/NVFP4 params), needed for loading non-native-fp8

megatron/core/extensions/transformer_engine.py

@@ -93,7 +93,7 @@ def _set_expert_parameter_attributes(
 ) -> None:
     """Set process-group and tensor-partition metadata on an expert TE module.
 
-    ``allreduce=False`` selects EDP for gradient reduction.
+    ``allreduce=False`` selects the expert topology, including EDP for gradient reduction.
 
     Weights and biases, including TEGroupedLinear's numbered parameters, are also marked as
     TP-partitioned according to ``parallel_mode``; row-parallel biases remain replicated.
@@ -104,7 +104,7 @@ def _set_expert_parameter_attributes(
     Args:
         module: Transformer Engine module whose direct parameters should be marked.
         parallel_mode: Tensor-parallel mode used by the module (``"column"``, ``"row"``, or None).
-        use_expert_pgs: Whether to use EP/ETP/EDP process groups instead of TP/CP/DP.
+        use_expert_pgs: Whether to use EP/ETP/EGTP/EDP process groups instead of TP/GTP/CP/DP.
     """
     for name, param in module.named_parameters(recurse=False):
         param.allreduce = not use_expert_pgs
@@ -982,6 +982,7 @@ def __init__(
         use_expert_pgs = is_expert and (
             self.expert_parallel
             or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size
+            or self.config.expert_gtp_weight_remat_size != self.config.gtp_weight_remat_size
         )
         if is_expert:
             rng_tracker_name = get_expert_parallel_rng_tracker_name()
@@ -1475,6 +1476,7 @@ def __init__(
             use_expert_pgs = (
                 config.expert_model_parallel_size > 1
                 or config.expert_tensor_parallel_size != config.tensor_model_parallel_size
+                or config.expert_gtp_weight_remat_size != config.gtp_weight_remat_size
             )
             _set_expert_parameter_attributes(self, "column", use_expert_pgs)
 
@@ -1731,6 +1733,7 @@ def __init__(
             use_expert_pgs = (
                 config.expert_model_parallel_size > 1
                 or config.expert_tensor_parallel_size != config.tensor_model_parallel_size
+                or config.expert_gtp_weight_remat_size != config.gtp_weight_remat_size
             )
             _set_expert_parameter_attributes(self, "row", use_expert_pgs)
 
@@ -2163,6 +2166,7 @@ def __init__(
             use_expert_pgs = is_expert and (
                 self.expert_parallel
                 or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size
+                or self.config.expert_gtp_weight_remat_size != self.config.gtp_weight_remat_size
             )
             if is_expert:
                 extra_kwargs["rng_tracker_name"] = get_expert_parallel_rng_tracker_name()

megatron/core/post_training/modelopt/layers.py

@@ -163,6 +163,7 @@ def __init__(
                     self.config.expert_model_parallel_size > 1
                     or self.config.expert_tensor_parallel_size
                     != self.config.tensor_model_parallel_size
+                    or self.config.expert_gtp_weight_remat_size != self.config.gtp_weight_remat_size
                 )
                 setattr(param, "allreduce", not use_expert_groups)
             else:

megatron/core/tensor_parallel/generalized_tensor_parallelism.py

@@ -501,11 +501,14 @@ def _gtp_slice_one_param(param, gtp_remat_group, *, name="<unnamed>"):
     shard = tensor[gtp_rank * shard_size : (gtp_rank + 1) * shard_size]
     gtp_shard = GTPShardedParam(shard.clone())
     gtp_shard.pad_length = pad_length
-    # Preserve the source weight's TP attributes (dropped when wrapping into GTPShardedParam),
-    # so param_is_not_tensor_parallel_duplicate still classifies it without GTP-specific code.
-    from megatron.core.tensor_parallel import copy_tensor_model_parallel_attributes
+    # Preserve duplicate-filtering metadata dropped when wrapping into GTPShardedParam.
+    from megatron.core.tensor_parallel import (
+        copy_gtp_attributes,
+        copy_tensor_model_parallel_attributes,
+    )
 
     copy_tensor_model_parallel_attributes(gtp_shard, param)
+    copy_gtp_attributes(gtp_shard, param)
     return gtp_shard
 
 
@@ -539,10 +542,14 @@ def _gtp_wrap_bf16_shard(module, name, param):
     :func:`_gtp_slice_one_param`, which slices a full weight — this only wraps it, no slicing.
     Returns the new param (also swapped into the module).
     """
-    from megatron.core.tensor_parallel import copy_tensor_model_parallel_attributes
+    from megatron.core.tensor_parallel import (
+        copy_gtp_attributes,
+        copy_tensor_model_parallel_attributes,
+    )
 
     gtp_shard = GTPShardedParam(param.data)
     copy_tensor_model_parallel_attributes(gtp_shard, param)
+    copy_gtp_attributes(gtp_shard, param)
     delattr(module, name)
     module._parameters[name] = gtp_shard
     return gtp_shard

megatron/core/tensor_parallel/layers.py

@@ -100,7 +100,7 @@ def param_is_not_tensor_parallel_duplicate(param, tp_group=None, expert_tp_group
     """
     if hasattr(param, "tensor_model_parallel") and param.tensor_model_parallel:
         return True
-    # allreduce=False marks parameters reduced over expert DP, so filter their duplicates over ETP.
+    # allreduce=False marks parameters using the expert topology, so filter duplicates over ETP.
     if not getattr(param, "allreduce", True) and expert_tp_group is not None:
         tp_group = expert_tp_group
     # Prefer provided tp_group when available (new explicit path).
@@ -961,6 +961,7 @@ def __init__(
         use_expert_pgs = self.is_expert and (
             self.expert_parallel
             or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size
+            or self.config.expert_gtp_weight_remat_size != self.config.gtp_weight_remat_size
         )
         self.output_size_per_partition = divide(output_size, world_size)
 
@@ -1379,6 +1380,7 @@ def __init__(
         use_expert_pgs = self.is_expert and (
             self.expert_parallel
             or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size
+            or self.config.expert_gtp_weight_remat_size != self.config.gtp_weight_remat_size
         )
         setattr(self.weight, "allreduce", not use_expert_pgs)
 

tests/unit_tests/distributed/test_grad_sync_with_expert_parallel.py

@@ -142,26 +142,44 @@ def _build_expert_linear(implementation: str, config: TransformerConfig) -> torc
 
 
 @pytest.mark.parametrize(
-    ("tensor_model_parallel_size", "expert_tensor_parallel_size"), [(2, 1), (1, 2)]
+    (
+        "tensor_model_parallel_size",
+        "expert_tensor_parallel_size",
+        "gtp_weight_remat_size",
+        "expert_gtp_weight_remat_size",
+    ),
+    [(2, 1, 1, 1), (1, 2, 1, 1), (1, 1, 1, 2)],
 )
 @pytest.mark.parametrize(
     "implementation", ["native", "transformer_engine", "transformer_engine_grouped"]
 )
 def test_expert_grad_sync_uses_expert_data_parallel_group(
-    implementation: str, tensor_model_parallel_size: int, expert_tensor_parallel_size: int
+    implementation: str,
+    tensor_model_parallel_size: int,
+    expert_tensor_parallel_size: int,
+    gtp_weight_remat_size: int,
+    expert_gtp_weight_remat_size: int,
 ):
-    """Expert gradients must not be reduced over ordinary DP when ETP differs from TP."""
+    """Expert gradients must use expert DP when the expert and dense topologies differ."""
+    if expert_gtp_weight_remat_size > 1:
+        from megatron.core.tensor_parallel.gtp_api import HAVE_GTP
+
+        if not HAVE_GTP:
+            pytest.skip("GTP requires TransformerEngine >= 2.19")
     if Utils.world_size < 4 or Utils.world_size % 4 != 0:
         pytest.skip("Test requires a world size divisible by four")
     if Utils.world_size > 16:
         pytest.skip("Rank-encoded gradients are intended for small unit-test world sizes")
 
-    Utils.initialize_model_parallel(
-        tensor_model_parallel_size=tensor_model_parallel_size,
-        expert_model_parallel_size=1,
-        expert_tensor_parallel_size=expert_tensor_parallel_size,
-    )
     try:
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tensor_model_parallel_size,
+            expert_model_parallel_size=1,
+            expert_tensor_parallel_size=expert_tensor_parallel_size,
+            gtp_remat_size=gtp_weight_remat_size,
+            expert_gtp_remat_size=expert_gtp_weight_remat_size,
+        )
+
         # Per-token loss leaves DDP's pre-collective gradient scaling at one.
         config = TransformerConfig(
             num_layers=1,
@@ -174,6 +192,10 @@ def test_expert_grad_sync_uses_expert_data_parallel_group(
             tensor_model_parallel_size=tensor_model_parallel_size,
             expert_model_parallel_size=1,
             expert_tensor_parallel_size=expert_tensor_parallel_size,
+            tensor_parallel_num_weight_shards=(tensor_model_parallel_size * gtp_weight_remat_size),
+            expert_tensor_parallel_num_weight_shards=(
+                expert_tensor_parallel_size * expert_gtp_weight_remat_size
+            ),
             calculate_per_token_loss=True,
             gradient_accumulation_fusion=False,
             perform_initialization=False,
@@ -194,7 +216,7 @@ def test_expert_grad_sync_uses_expert_data_parallel_group(
         )
 
         expert_dp_group = parallel_state.get_expert_data_parallel_group(
-            partial_expert_data_parallel=True
+            with_gtp_remat=False, partial_expert_data_parallel=True
         )
         ordinary_dp_group = parallel_state.get_data_parallel_group(
             with_context_parallel=True, partial_data_parallel=True
@@ -222,6 +244,10 @@ def test_expert_grad_sync_uses_expert_data_parallel_group(
         assert len(model.expert_parallel_buffers) == 1
         assert all(param.allreduce is False for param in model.parameters())
     finally:
+        if expert_gtp_weight_remat_size > 1:
+            from megatron.core.tensor_parallel.generalized_tensor_parallelism import reset_gtp_state
+
+            reset_gtp_state()
         Utils.destroy_model_parallel()
 
 

tests/unit_tests/training/test_param_norm.py

@@ -19,6 +19,8 @@ def _build_tiny_moe_gpt(
     tensor_parallel_size: int,
     expert_parallel_size: int,
     expert_tensor_parallel_size: int,
+    tensor_parallel_num_weight_shards: int | None = None,
+    expert_tensor_parallel_num_weight_shards: int | None = None,
     bf16: bool = False,
     add_bias_linear: bool = False,
 ) -> GPTModel:
@@ -36,6 +38,8 @@ def _build_tiny_moe_gpt(
         tensor_model_parallel_size=tensor_parallel_size,
         expert_model_parallel_size=expert_parallel_size,
         expert_tensor_parallel_size=expert_tensor_parallel_size,
+        tensor_parallel_num_weight_shards=tensor_parallel_num_weight_shards,
+        expert_tensor_parallel_num_weight_shards=expert_tensor_parallel_num_weight_shards,
         sequence_parallel=tensor_parallel_size > 1,
         use_cpu_initialization=True,
         add_bias_linear=add_bias_linear,
@@ -119,22 +123,116 @@ def test_moe_param_norm_counts_each_logical_parameter_once(
         Utils.destroy_model_parallel()
 
 
+def test_moe_param_norm_uses_expert_gtp_topology_when_it_differs_from_dense_gtp(monkeypatch):
+    """Expert parameters must use EGTP even when EP, TP, and ETP alone do not distinguish them."""
+    from megatron.core.tensor_parallel.generalized_tensor_parallelism import (
+        GTP_CONFIG,
+        GTPShardedParam,
+        reset_gtp_state,
+        update_gtp_config,
+    )
+    from megatron.core.tensor_parallel.gtp_api import HAVE_GTP
+
+    if not HAVE_GTP:
+        pytest.skip("GTP requires TransformerEngine >= 2.19")
+    if Utils.world_size < 2 or Utils.world_size % 2 != 0:
+        pytest.skip("test requires an even world size")
+
+    monkeypatch.setattr(
+        common_utils, "get_args", lambda: SimpleNamespace(use_megatron_fsdp=False, bf16=False)
+    )
+    # Keep the all-ones assertion focused on topology rather than physical GTP padding.
+    original_pad_for_alignment = GTP_CONFIG.pad_for_alignment
+    update_gtp_config(pad_for_alignment=0)
+
+    try:
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1,
+            expert_model_parallel_size=1,
+            expert_tensor_parallel_size=1,
+        )
+        reference_model = _build_tiny_moe_gpt(
+            tensor_parallel_size=1, expert_parallel_size=1, expert_tensor_parallel_size=1
+        )
+        _fill_parameters_with_ones(reference_model)
+        expected_numel = sum(param.numel() for param in reference_model.parameters())
+        expected_norm = math.sqrt(expected_numel)
+        del reference_model
+
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1,
+            expert_model_parallel_size=1,
+            expert_tensor_parallel_size=1,
+            gtp_remat_size=1,
+            expert_gtp_remat_size=2,
+        )
+        model = _build_tiny_moe_gpt(
+            tensor_parallel_size=1,
+            expert_parallel_size=1,
+            expert_tensor_parallel_size=1,
+            tensor_parallel_num_weight_shards=1,
+            expert_tensor_parallel_num_weight_shards=2,
+        )
+        _fill_parameters_with_ones(model)
+
+        expert_params = [param for name, param in model.named_parameters() if ".experts." in name]
+        assert any(isinstance(param, GTPShardedParam) for param in expert_params)
+
+        actual_norm = common_utils.calc_params_l2_norm(model)
+
+        assert actual_norm == pytest.approx(expected_norm)
+    finally:
+        update_gtp_config(pad_for_alignment=original_pad_for_alignment)
+        reset_gtp_state()
+        Utils.destroy_model_parallel()
+
+
 @pytest.mark.parametrize("use_distributed_optimizer", (False, True), ids=("optimizer", "distopt"))
 @pytest.mark.parametrize(
-    ("tensor_parallel_size", "expert_parallel_size", "expert_tensor_parallel_size"),
-    ((2, 2, 1), (2, 1, 2), (4, 1, 2), (2, 1, 4)),
-    ids=("expert-parallel", "expert-tensor-parallel", "tp-larger-than-etp", "etp-larger-than-tp"),
+    (
+        "tensor_parallel_size",
+        "expert_parallel_size",
+        "expert_tensor_parallel_size",
+        "gtp_weight_remat_size",
+        "expert_gtp_weight_remat_size",
+    ),
+    ((2, 2, 1, 1, 1), (2, 1, 2, 1, 1), (4, 1, 2, 1, 1), (2, 1, 4, 1, 1), (1, 1, 1, 1, 2)),
+    ids=(
+        "expert-parallel",
+        "expert-tensor-parallel",
+        "tp-larger-than-etp",
+        "etp-larger-than-tp",
+        "expert-gtp-differs-from-dense-gtp",
+    ),
 )
 def test_moe_gradient_stats_and_clipping_count_each_logical_gradient_once(
     tensor_parallel_size: int,
     expert_parallel_size: int,
     expert_tensor_parallel_size: int,
+    gtp_weight_remat_size: int,
+    expert_gtp_weight_remat_size: int,
     use_distributed_optimizer: bool,
 ):
     """Gradient norm, clipping, and zero count should include each logical gradient once."""
+    if expert_gtp_weight_remat_size > 1:
+        from megatron.core.tensor_parallel.gtp_api import HAVE_GTP
+
+        if not HAVE_GTP:
+            pytest.skip("GTP requires TransformerEngine >= 2.19")
     if Utils.world_size < 4 or Utils.world_size % 4 != 0:
         pytest.skip("test requires a world size divisible by four")
 
+    original_pad_for_alignment = None
+    if expert_gtp_weight_remat_size > 1:
+        from megatron.core.tensor_parallel.generalized_tensor_parallelism import (
+            GTP_CONFIG,
+            update_gtp_config,
+        )
+
+        # Keep the all-ones assertion focused on topology rather than physical GTP padding.
+        original_pad_for_alignment = GTP_CONFIG.pad_for_alignment
+        update_gtp_config(pad_for_alignment=0)
+
     try:
         Utils.initialize_model_parallel(
             tensor_model_parallel_size=1,
@@ -152,11 +250,17 @@ def test_moe_gradient_stats_and_clipping_count_each_logical_gradient_once(
             tensor_model_parallel_size=tensor_parallel_size,
             expert_model_parallel_size=expert_parallel_size,
             expert_tensor_parallel_size=expert_tensor_parallel_size,
+            gtp_remat_size=gtp_weight_remat_size,
+            expert_gtp_remat_size=expert_gtp_weight_remat_size,
         )
         model = _build_tiny_moe_gpt(
             tensor_parallel_size=tensor_parallel_size,
             expert_parallel_size=expert_parallel_size,
             expert_tensor_parallel_size=expert_tensor_parallel_size,
+            tensor_parallel_num_weight_shards=(tensor_parallel_size * gtp_weight_remat_size),
+            expert_tensor_parallel_num_weight_shards=(
+                expert_tensor_parallel_size * expert_gtp_weight_remat_size
+            ),
             bf16=True,
         )
         ddp_config = DistributedDataParallelConfig(
@@ -211,6 +315,14 @@ def test_moe_gradient_stats_and_clipping_count_each_logical_gradient_once(
             grads_checked += 1
         assert grads_checked > 0
     finally:
+        if expert_gtp_weight_remat_size > 1:
+            from megatron.core.tensor_parallel.generalized_tensor_parallelism import (
+                reset_gtp_state,
+                update_gtp_config,
+            )
+
+            update_gtp_config(pad_for_alignment=original_pad_for_alignment)
+            reset_gtp_state()
         Utils.destroy_model_parallel()