diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py index 2a3c7581122..94d1db3f576 100644 --- a/megatron/core/parallel_state.py +++ b/megatron/core/parallel_state.py @@ -155,6 +155,12 @@ # Paralel group of all GPUs in a distributed optimizer instance _INTRA_DISTRIBUTED_OPTIMIZER_INSTANCE_GROUP = None +# RankGenerator instances used to build the decoder and expert data-parallel +# groups, cached so that create_all_gather_groups() can reproduce the same +# rank membership instead of re-deriving it with different arguments. +_DECODER_RANK_GENERATOR = None +_EXPERT_DECODER_RANK_GENERATOR = None + # Memory buffers to avoid dynamic memory allocation _GLOBAL_MEMORY_BUFFER = None @@ -909,6 +915,11 @@ def initialize_model_parallel( timeout = timedelta(minutes=distributed_timeout_minutes) # Build the data-parallel groups. + global _DECODER_RANK_GENERATOR + global _EXPERT_DECODER_RANK_GENERATOR + _DECODER_RANK_GENERATOR = decoder_rank_generator + _EXPERT_DECODER_RANK_GENERATOR = expert_decoder_rank_generator + global _DATA_PARALLEL_GROUP global _DATA_PARALLEL_GROUP_GLOO global _DATA_PARALLEL_GLOBAL_RANKS @@ -1616,26 +1627,21 @@ def create_all_gather_groups(for_expert_parallelism=False, timeout=None, nccl_co "Call initialize_model_parallel() first." ) + # Reuse the RankGenerator instances created by initialize_model_parallel() + # so that AG group membership always matches the actual DP/expert-DP + # groups, regardless of the order/rank_offset the caller used. + if _DECODER_RANK_GENERATOR is None or _EXPERT_DECODER_RANK_GENERATOR is None: + raise RuntimeError( + "create_all_gather_groups() requires the decoder rank generators cached by " + "initialize_model_parallel() to be set. Call initialize_model_parallel() first." + ) + rank = torch.distributed.get_rank() - pp_size = get_pipeline_model_parallel_world_size() - cp_size = get_context_parallel_world_size() - tp_size = get_tensor_model_parallel_world_size() ep_size = get_expert_model_parallel_world_size() - dp_size = get_data_parallel_world_size() - gtp_remat_size = get_gtp_weight_remat_world_size() or 1 # Create regular DP all-gather group dp_cp_ag_group = None - decoder_rank_gen = RankGenerator( - tp=tp_size, - ep=1, - dp=dp_size, - pp=pp_size, - cp=cp_size, - gtp_remat=gtp_remat_size, - order=_inject_gtp_remat_axis('tp-cp-ep-dp-pp', after='tp'), - rank_offset=0, - ) + decoder_rank_gen = _DECODER_RANK_GENERATOR for ranks_with_cp in decoder_rank_gen.get_ranks('dp-cp'): group_with_cp_ag = create_group( @@ -1650,20 +1656,7 @@ def create_all_gather_groups(for_expert_parallelism=False, timeout=None, nccl_co # Create expert DP all-gather group if requested expt_dp_ag_group = None if for_expert_parallelism and ep_size > 1: - expert_tp_size = get_expert_tensor_parallel_world_size() - expert_dp_size = get_expert_data_parallel_world_size() - egtp_remat_size = get_expert_gtp_weight_remat_world_size() or 1 - - expert_rank_gen = RankGenerator( - tp=expert_tp_size, - ep=ep_size, - dp=expert_dp_size, - pp=pp_size, - cp=1, - gtp_remat=egtp_remat_size, - order=_inject_gtp_remat_axis('tp-cp-ep-dp-pp', after='ep'), - rank_offset=0, - ) + expert_rank_gen = _EXPERT_DECODER_RANK_GENERATOR for expert_dp_ranks in expert_rank_gen.get_ranks('dp'): expert_dp_ag = create_group( @@ -2546,6 +2539,12 @@ def destroy_model_parallel(): global _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP_WITH_GTP_REMAT _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP_WITH_GTP_REMAT = None + global _DECODER_RANK_GENERATOR + _DECODER_RANK_GENERATOR = None + + global _EXPERT_DECODER_RANK_GENERATOR + _EXPERT_DECODER_RANK_GENERATOR = None + global _CONTEXT_PARALLEL_GROUP _CONTEXT_PARALLEL_GROUP = None diff --git a/tests/unit_tests/test_parallel_state.py b/tests/unit_tests/test_parallel_state.py index 65b0d5ca91a..806300f4e64 100644 --- a/tests/unit_tests/test_parallel_state.py +++ b/tests/unit_tests/test_parallel_state.py @@ -533,9 +533,10 @@ def test_hybrid_dp_cp_groups(world_size, tp_size, cp_size, dp_size): Utils.destroy_model_parallel() -def test_separate_all_gather_group(): +@pytest.mark.parametrize('order', test_parallel_order) +def test_separate_all_gather_group(order): """AG/RS overlap communicators live on ProcessGroupCollection (via create_all_gather_groups).""" - Utils.initialize_model_parallel(context_parallel_size=world_size) + Utils.initialize_model_parallel(context_parallel_size=world_size, order=order) dp_cp_group = ps.get_data_parallel_group(with_context_parallel=True) dp_cp_ranks = torch.distributed.get_process_group_ranks(dp_cp_group) @@ -553,12 +554,14 @@ def test_separate_all_gather_group(): Utils.destroy_model_parallel() -def test_expert_all_gather_group(): +@pytest.mark.parametrize('order', test_parallel_order) +def test_expert_all_gather_group(order): """Test expert AG groups for MoE models with AG/RS overlap.""" # Initialize model parallel with expert parallelism Utils.initialize_model_parallel( expert_model_parallel_size=min(2, world_size), context_parallel_size=max(1, world_size // 2) if world_size > 1 else 1, + order=order, ) # Get ranks for both regular and expert AG groups @@ -567,26 +570,25 @@ def test_expert_all_gather_group(): expt_dp_group = ps.get_expert_data_parallel_group() expt_dp_ranks = torch.distributed.get_process_group_ranks(expt_dp_group) - # Create AG groups for both regular and expert parameters - dp_cp_ag_group = torch.distributed.new_group(ranks=dp_cp_ranks, backend='nccl') - expt_dp_ag_group = torch.distributed.new_group(ranks=expt_dp_ranks, backend='nccl') + # Create AG groups for both regular and expert parameters via the production helper + dp_cp_ag_group, expt_dp_ag_group = ps.create_all_gather_groups(for_expert_parallelism=True) # Create ProcessGroupCollection with AG groups pg_collection = ProcessGroupCollection.use_mpu_process_groups() pg_collection.dp_cp_ag = dp_cp_ag_group pg_collection.expt_dp_ag = expt_dp_ag_group - # Verify both AG groups are set + # Verify the regular AG group is set and matches the real dp-cp group assert pg_collection.dp_cp_ag is not None - assert pg_collection.expt_dp_ag is not None + ag_ranks = torch.distributed.get_process_group_ranks(dp_cp_ag_group) + assert ag_ranks == dp_cp_ranks, "AG group should have same ranks as dp-cp group" # Verify expert AG group has same ranks as expert dp group - expt_dp_group = pg_collection.expt_dp - if expt_dp_group is not None: + if expt_dp_ag_group is not None: + assert pg_collection.expt_dp_ag is not None expt_ag_ranks = torch.distributed.get_process_group_ranks(expt_dp_ag_group) - expt_dp_ranks_actual = torch.distributed.get_process_group_ranks(expt_dp_group) assert ( - expt_ag_ranks == expt_dp_ranks_actual + expt_ag_ranks == expt_dp_ranks ), "Expert AG group should have same ranks as expert dp group" assert ( expt_dp_ag_group != expt_dp_group @@ -595,6 +597,69 @@ def test_expert_all_gather_group(): Utils.destroy_model_parallel() +def test_create_all_gather_groups_uses_cached_rank_generators(monkeypatch): + """create_all_gather_groups() must reuse the RankGenerator instances cached by + initialize_model_parallel(), not rebuild ones with hardcoded order/rank_offset.""" + order = 'tp-cp-pp-ep-dp' + decoder_rank_generator = ps.RankGenerator( + tp=2, ep=1, dp=4, pp=2, cp=1, order=order, rank_offset=4 + ) + expert_decoder_rank_generator = ps.RankGenerator( + tp=2, ep=2, dp=2, pp=2, cp=1, order=order, rank_offset=4 + ) + current_rank = decoder_rank_generator.get_ranks('dp-cp')[0][0] + + monkeypatch.setattr(ps, '_DATA_PARALLEL_GROUP', object()) + monkeypatch.setattr(ps, '_DECODER_RANK_GENERATOR', decoder_rank_generator) + monkeypatch.setattr(ps, '_EXPERT_DECODER_RANK_GENERATOR', expert_decoder_rank_generator) + monkeypatch.setattr(ps, '_MPU_EXPERT_MODEL_PARALLEL_WORLD_SIZE', 2) + # Match the fake generators so a reverted (pre-fix) implementation still runs to + # completion and fails by rank mismatch rather than by uninitialized state. + monkeypatch.setattr(ps, 'get_pipeline_model_parallel_world_size', lambda: 2) + monkeypatch.setattr(ps, 'get_context_parallel_world_size', lambda: 1) + monkeypatch.setattr(ps, 'get_tensor_model_parallel_world_size', lambda: 2) + monkeypatch.setattr(ps, 'get_data_parallel_world_size', lambda: 4) + monkeypatch.setattr(ps, 'get_expert_tensor_parallel_world_size', lambda: 2) + monkeypatch.setattr(ps, 'get_expert_data_parallel_world_size', lambda: 2) + monkeypatch.setattr(torch.distributed, 'get_rank', lambda: current_rank) + + captured_groups = {} + + def fake_new_group(**kwargs): + ranks = tuple(kwargs['ranks']) + group = object() + captured_groups[ranks] = group + return group + + monkeypatch.setattr(torch.distributed, 'new_group', fake_new_group) + + dp_cp_ag_group, expt_dp_ag_group = ps.create_all_gather_groups(for_expert_parallelism=True) + + expected_dp_cp_ranks = tuple( + next(ranks for ranks in decoder_rank_generator.get_ranks('dp-cp') if current_rank in ranks) + ) + expected_expert_dp_ranks = tuple( + next( + ranks + for ranks in expert_decoder_rank_generator.get_ranks('dp') + if current_rank in ranks + ) + ) + assert expected_dp_cp_ranks in captured_groups + assert expected_expert_dp_ranks in captured_groups + assert dp_cp_ag_group is captured_groups[expected_dp_cp_ranks] + assert expt_dp_ag_group is captured_groups[expected_expert_dp_ranks] + + default_decoder_rank_generator = ps.RankGenerator( + tp=2, ep=1, dp=4, pp=2, cp=1, order='tp-cp-ep-dp-pp', rank_offset=0 + ) + default_expert_rank_generator = ps.RankGenerator( + tp=2, ep=2, dp=2, pp=2, cp=1, order='tp-cp-ep-dp-pp', rank_offset=0 + ) + assert list(expected_dp_cp_ranks) not in default_decoder_rank_generator.get_ranks('dp-cp') + assert list(expected_expert_dp_ranks) not in default_expert_rank_generator.get_ranks('dp') + + def test_process_group_collection_defaults(): """Test that ProcessGroupCollection initializes AG groups to None by default.""" # Initialize model parallel