Make Bridge Communicator aware of GTP

megatron/core/pipeline_parallel/bridge_communicator.py

@@ -212,15 +212,17 @@ def _get_or_create_bridge_pg(cls, ranks: List[int]):
     def get_leader_rank(self, grid: HyperCommGrid, is_src: bool) -> List[int]:
         """Get the leader rank for a given grid and direction.
 
-        We elect leader rank for each dp replica, the first tp-cp rank in the group
+        We elect a leader for each DP and GTP data lane, the first tp-cp rank in the group
         in the last pp stage (for src grid) or first pp stage (for dest grid) is the leader.
         """
         leader_ranks = []
         local_leader_rank = None
-        # grid.gen_rank_enum(["tp", "cp", "pp"]) # vary tp & cp, but same dp
+        # grid.gen_rank_enum(["tp", "cp", "pp"]) # vary tp & cp, same dp and gtp_remat
         # returns a list of sublists, each sublist is a group of ranks
-        # that have different tp & cp & pp, same dp
-        per_dp_replica_ranks = grid._gen_rank_enum([x for x in grid.dim_names if x != "dp"])
+        # that have different tp & cp & pp, same dp and gtp_remat
+        per_dp_replica_ranks = grid._gen_rank_enum(
+            [x for x in grid.dim_names if x not in ("dp", "gtp_remat")]
+        )
         if is_src:
             # Add rank from last pp stage
             ranks = []

tests/unit_tests/pipeline_parallel/test_bridge_communicator.py

@@ -112,7 +112,7 @@ def _shard_and_copy_(
 _active_grids: list = []
 
 
-def create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1):
+def create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1, gtp_remat=1):
     """Create a HyperCommGrid with tensor parallelism=2, context parallelism=2, and data parallelism=2."""
     # Set up environment for world size 8 if not already set
     if not dist.is_initialized():
@@ -123,12 +123,13 @@ def create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1):
         os.environ["WORLD_SIZE"] = "8"
 
     grid = HyperCommGrid(
-        shape=[tp, cp, pp, dp],
-        dim_names=["tp", "cp", "pp", "dp"],
+        shape=[tp, gtp_remat, cp, pp, dp],
+        dim_names=["tp", "gtp_remat", "cp", "pp", "dp"],
         rank_offset=offset,
         backend="nccl",
     )
     _ = grid.create_pg(["tp"])
+    _ = grid.create_pg(["gtp_remat"])
     _ = grid.create_pg(["cp"])
     _ = grid.create_pg(["pp"])
     _ = grid.create_pg(["dp"])
@@ -326,6 +327,17 @@ def test_bridge_pg_membership(self, grid1_tp, grid1_dp, grid2_tp, grid2_dp):
         ]
         assert all(rank not in expected for rank in member_ranks)
 
+    def test_gtp_is_an_independent_bridge_data_lane(self):
+        src_grid = create_hypercomm_grid(offset=0, tp=2, dp=2)
+        dest_grid = create_hypercomm_grid(offset=4, tp=2, dp=1, gtp_remat=2)
+        bridge = BridgeCommunicator(src_grid, dest_grid)
+
+        assert len(bridge.src_tp_leaders) == 2
+        assert len(bridge.dest_tp_leaders) == 2
+        assert sorted(set(bridge.src_tp_leaders) | set(bridge.dest_tp_leaders)) == list(
+            dist.get_process_group_ranks(bridge.bridge_pg)
+        )
+
     def test_send_forward_recv_forward(self):
         """Test send_forward and recv_forward operations."""