fix(training): correct parameter norm reporting with GTP (#6133)

Signed-off-by: ykarnati <ykarnati@nvidia.com>

Author: yashaswikarnati

src/megatron/bridge/training/utils/train_utils.py

@@ -30,7 +30,7 @@
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.moe.moe_utils import track_moe_metrics
 from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper
-from megatron.core.utils import get_data_parallel_group_if_dtensor, to_local_if_dtensor
+from megatron.core.utils import get_data_parallel_group_if_dtensor, get_pg_rank, to_local_if_dtensor
 
 from megatron.bridge.models.common.heads import (
     LinearForLastLayer as LinearForLastLayer,
@@ -237,20 +237,31 @@ def calc_params_l2_norm(
     sharded_moe_params_data = []
     data_parallel_group = None
     pg_collection = get_pg_collection(model)
+    gtp_rank = get_pg_rank(pg_collection.gtp_remat)
+    expert_gtp_group = pg_collection.expt_gtp_remat
+    expert_gtp_rank = get_pg_rank(expert_gtp_group)
 
     for model_chunk in model:
         for param in model_chunk.parameters():
             data_parallel_group = get_data_parallel_group_if_dtensor(param, data_parallel_group)
+            is_gtp = getattr(param, "is_gtp_weight_remat", False)
             # MCore uses allreduce=False to mark parameters that use expert-parallel process groups.
             uses_expert_parallel_groups = not getattr(param, "allreduce", True)
-            is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(
+            # GTP parameters are unique across TP ranks. Other parameters still need TP filtering.
+            if not is_gtp and not param_is_not_tensor_parallel_duplicate(
                 param,
                 tp_group=pg_collection.tp,
                 expert_tp_group=pg_collection.expt_tp,
-            )
-            if not is_not_tp_duplicate:
+            ):
+                continue
+
+            # Parameters that are not GTP-sharded are replicated across the corresponding GTP axis.
+            if uses_expert_parallel_groups:
+                if not is_gtp and expert_gtp_rank != 0:
+                    continue
+            elif not is_gtp and gtp_rank != 0:
                 continue
-            assert is_not_tp_duplicate
+
             if uses_expert_parallel_groups:
                 assert param_is_not_shared(param)
                 param = to_local_if_dtensor(param)
@@ -360,6 +371,15 @@ def calc_params_l2_norm(
     )
     moe_norm_2 += sharded_moe_norm_2
 
+    # Expert model parallel excludes expert GTP. This reduction collects both unique GTP shards
+    # and ordinary expert parameters, which are counted only on expert GTP rank zero above.
+    if expert_gtp_group is not None:
+        torch.distributed.all_reduce(
+            moe_norm_2,
+            op=torch.distributed.ReduceOp.SUM,
+            group=expert_gtp_group,
+        )
+
     # Reduce norm across model parallel groups (dense and expert).
     # Dense params should sum across all model-parallel GPUs (tensor + pipeline).
     dense_reduce_group = pg_collection.mp

tests/unit_tests/training/utils/test_train_utils.py

@@ -2718,7 +2718,9 @@ def __init__(self):
                 # Minimal set of groups used by calc_params_l2_norm
                 self.dp_cp = object()
                 self.expt_dp = object()
+                self.expt_gtp_remat = None
                 self.expt_tp = object()
+                self.gtp_remat = object()
                 self.mp = object()
                 self.tp = object()
                 self.tp_ep_pp = object()
@@ -2762,6 +2764,128 @@ def mock_model_config_bf16(self):
         config.bf16 = True
         return config
 
+    def test_gtp_param_norm_filters_replicas_and_keeps_shards(
+        self,
+        monkeypatch,
+        mock_model_config_fp32,
+        _patch_pg_collection,
+    ):
+        """Non-sharded parameters are counted once while every GTP shard contributes."""
+
+        class _DenseGTPModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.norm = torch.nn.Parameter(torch.tensor([2.0], device="cuda"))
+                self.weight = torch.nn.Parameter(torch.tensor([3.0], device="cuda"))
+                self.weight.is_gtp_weight_remat = True
+
+        model = _DenseGTPModel()
+        duplicate_filter = mock.MagicMock(return_value=True)
+
+        def fake_all_reduce(tensor, op, group):
+            del op
+            if group is _patch_pg_collection.mp:
+                # Rank zero contributes the replicated norm (2**2) and its GTP shard (4**2).
+                tensor.add_(20.0)
+
+        monkeypatch.setattr(
+            "megatron.bridge.training.utils.train_utils.get_data_parallel_group_if_dtensor",
+            lambda param, group: None,
+        )
+        monkeypatch.setattr(
+            "megatron.bridge.training.utils.train_utils.param_is_not_tensor_parallel_duplicate",
+            duplicate_filter,
+        )
+        monkeypatch.setattr(
+            "megatron.bridge.training.utils.train_utils.to_local_if_dtensor",
+            lambda param: param,
+        )
+        monkeypatch.setattr(
+            "megatron.bridge.training.utils.train_utils.get_pg_rank",
+            lambda group: 1 if group is _patch_pg_collection.gtp_remat else 0,
+        )
+        monkeypatch.setattr("torch.distributed.get_process_group_ranks", lambda group: [0])
+        monkeypatch.setattr("torch.distributed.all_reduce", fake_all_reduce)
+
+        actual_norm = calc_params_l2_norm(model, mock_model_config_fp32)
+
+        assert actual_norm == pytest.approx(math.sqrt(29.0))
+        duplicate_filter.assert_called_once_with(
+            model.norm,
+            tp_group=_patch_pg_collection.tp,
+            expert_tp_group=_patch_pg_collection.expt_tp,
+        )
+
+    @pytest.mark.parametrize(
+        ("expert_gtp_rank", "remote_contribution", "expected_local_input"),
+        [
+            pytest.param(0, 9.0, 20.0, id="rank-zero-keeps-replica"),
+            pytest.param(1, 13.0, 16.0, id="nonzero-rank-drops-replica"),
+        ],
+    )
+    def test_expert_gtp_param_norm_reduces_over_expert_gtp(
+        self,
+        monkeypatch,
+        mock_model_config_fp32,
+        _patch_pg_collection,
+        expert_gtp_rank,
+        remote_contribution,
+        expected_local_input,
+    ):
+        """Expert GTP shards are summed while replicated parameters contribute once."""
+
+        class _ExpertGTPModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.weight = torch.nn.Parameter(torch.tensor([4.0], device="cuda"))
+                self.weight.allreduce = False
+                self.weight.is_gtp_weight_remat = True
+                self.bias = torch.nn.Parameter(torch.tensor([2.0], device="cuda"))
+                self.bias.allreduce = False
+
+        _patch_pg_collection.expt_gtp_remat = object()
+        model = _ExpertGTPModel()
+        duplicate_filter = mock.MagicMock(return_value=True)
+        expert_gtp_inputs = []
+
+        def fake_all_reduce(tensor, op, group):
+            del op
+            if group is _patch_pg_collection.expt_gtp_remat:
+                expert_gtp_inputs.append(tensor.item())
+                tensor.add_(remote_contribution)
+
+        monkeypatch.setattr(
+            "megatron.bridge.training.utils.train_utils.get_data_parallel_group_if_dtensor",
+            lambda param, group: None,
+        )
+        monkeypatch.setattr(
+            "megatron.bridge.training.utils.train_utils.param_is_not_tensor_parallel_duplicate",
+            duplicate_filter,
+        )
+        monkeypatch.setattr(
+            "megatron.bridge.training.utils.train_utils.to_local_if_dtensor",
+            lambda param: param,
+        )
+        monkeypatch.setattr(
+            "megatron.bridge.training.utils.train_utils.get_pg_rank",
+            lambda group: expert_gtp_rank if group is _patch_pg_collection.expt_gtp_remat else 0,
+        )
+        monkeypatch.setattr(
+            "torch.distributed.get_process_group_ranks",
+            lambda group: [0] if group is _patch_pg_collection.mp else [1],
+        )
+        monkeypatch.setattr("torch.distributed.all_reduce", fake_all_reduce)
+
+        actual_norm = calc_params_l2_norm(model, mock_model_config_fp32)
+
+        assert expert_gtp_inputs == pytest.approx([expected_local_input])
+        assert actual_norm == pytest.approx(math.sqrt(29.0))
+        duplicate_filter.assert_called_once_with(
+            model.bias,
+            tp_group=_patch_pg_collection.tp,
+            expert_tp_group=_patch_pg_collection.expt_tp,
+        )
+
     @mock.patch("megatron.bridge.training.utils.train_utils.get_data_parallel_group_if_dtensor")
     @mock.patch("megatron.bridge.training.utils.train_utils.param_is_not_tensor_parallel_duplicate")
     @mock.patch("megatron.bridge.training.utils.train_utils.to_local_if_dtensor")
@@ -3000,6 +3124,8 @@ def fake_all_reduce(tensor, op, group):
                 tensor.add_(12.0)
             elif group is _patch_pg_collection.expt_dp:
                 tensor.add_(16.0)
+            elif group is _patch_pg_collection.expt_gtp_remat:
+                pass
             elif group is not _patch_pg_collection.mp:
                 raise AssertionError("unexpected reduction group")
 
@@ -3212,6 +3338,8 @@ def __init__(self):
             expt_tp=expert_tp_group,
             dp_cp=reduce_group,
             expt_dp=reduce_group,
+            expt_gtp_remat=reduce_group,
+            gtp_remat=reduce_group,
             mp=reduce_group,
             tp_ep_pp=reduce_group,
         )