diff --git a/megatron/core/pipeline_parallel/bridge_communicator.py b/megatron/core/pipeline_parallel/bridge_communicator.py index fc234cac8ae..b7f85e36c56 100644 --- a/megatron/core/pipeline_parallel/bridge_communicator.py +++ b/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 = [] diff --git a/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py b/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py index e4801ad8939..4eaa11b8c24 100644 --- a/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py +++ b/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."""