Skip to content

Fix AG groups for custom parallel order #5754

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MrAnayDongre wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix AG groups for custom parallel order #5754

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 28 additions & 29 deletions megatron/core/parallel_state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -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(
Expand Down Expand Up @@ -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

Expand Down
89 changes: 77 additions & 12 deletions tests/unit_tests/test_parallel_state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down