diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 27b675d1b8d..adc3ca442a1 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -685,10 +685,10 @@ def init_state_fn(opt, config=None): optimizer = FP32Optimizer(optimizer, config, init_state_fn) setattr(optimizer, 'grad_stats_parallel_group', model_parallel_group) - if pg_collection is None or not hasattr(pg_collection, 'tp'): - tp_group = parallel_state.get_tensor_model_parallel_group() - else: - tp_group = pg_collection.tp + assert pg_collection is not None and hasattr( + pg_collection, 'tp' + ), "pg_collection with tp must be resolved by get_megatron_optimizer" + tp_group = pg_collection.tp # TODO(M4): plumb tp_group through optimizer constructors so this setattr disappears. setattr(optimizer, 'tp_group', tp_group) @@ -769,8 +769,7 @@ def _get_megatron_emerging_optimizer( if config.fp16: raise ValueError('emerging optimizer with fp16 is not supported.') - if pg_collection is None: - pg_collection = ProcessGroupCollection.use_mpu_process_groups() + assert pg_collection is not None, "pg_collection must be resolved by get_megatron_optimizer" log_single_rank(logger, logging.INFO, f'Setting up emerging optimizer with config {config}') @@ -884,10 +883,10 @@ def _get_megatron_emerging_optimizer( else: optimizer = FP32Optimizer(optimizer, config, init_state_fn) setattr(optimizer, 'grad_stats_parallel_group', model_parallel_group) - if pg_collection is None or not hasattr(pg_collection, 'tp'): - tp_group = parallel_state.get_tensor_model_parallel_group() - else: - tp_group = pg_collection.tp + assert pg_collection is not None and hasattr( + pg_collection, 'tp' + ), "pg_collection with tp must be resolved by get_megatron_optimizer" + tp_group = pg_collection.tp setattr(optimizer, 'tp_group', tp_group) results.append(optimizer) continue @@ -1020,6 +1019,12 @@ def get_megatron_optimizer( check_config_overrides_consistency(config, config_overrides) + # Compatibility boundary. get_megatron_optimizer is the edge of megatron/core for optimizer + # construction, so the global-state fallback lives here and nowhere deeper: every helper below + # receives an explicit collection. See docs/developer/parallel-state-deprecation.md. + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + # TODO: the standard and emerging optimizer paths handle pg_collection differently; # unify them so both use a single pg_collection-based flow. if config.optimizer not in ('adam', 'sgd'): diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py index 337485b4d12..7aa6292ae89 100644 --- a/megatron/core/parallel_state.py +++ b/megatron/core/parallel_state.py @@ -1482,19 +1482,33 @@ def get_data_parallel_group(with_context_parallel=False, partial_data_parallel=F return _DATA_PARALLEL_GROUP -def get_data_parallel_group_gloo(with_context_parallel=False, partial_data_parallel=False): - """Get the Gloo data-parallel group the caller rank belongs to.""" +def get_data_parallel_group_gloo( + with_context_parallel=False, partial_data_parallel=False, check_initialized=True +): + """Get the Gloo data-parallel group the caller rank belongs to. + + Args: + check_initialized: When False, return None instead of asserting if the gloo group was + never created (``initialize_model_parallel(create_gloo_process_groups=False)``). + Needed so a ProcessGroupCollection can be materialised in gloo-less jobs. + """ if with_context_parallel: if partial_data_parallel: + if not check_initialized and _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP_GLOO is None: + return None assert ( _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP_GLOO is not None ), "Intra partial data parallel group is not initialized" return _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP_GLOO + if not check_initialized and _DATA_PARALLEL_GROUP_WITH_CP_GLOO is None: + return None assert ( _DATA_PARALLEL_GROUP_WITH_CP_GLOO is not None ), "data parallel group-gloo with context parallel combined is not initialized" return _DATA_PARALLEL_GROUP_WITH_CP_GLOO else: + if not check_initialized and _DATA_PARALLEL_GROUP_GLOO is None: + return None assert _DATA_PARALLEL_GROUP_GLOO is not None, "data parallel group-gloo is not initialized" assert partial_data_parallel == False, "Partial DP for Optimizer needs to include CP" return _DATA_PARALLEL_GROUP_GLOO @@ -1999,14 +2013,23 @@ def get_expert_data_parallel_group(check_initialized=True, partial_expert_data_p return _EXPERT_DATA_PARALLEL_GROUP -def get_expert_data_parallel_group_gloo(partial_expert_data_parallel=False): - """Get expert data parallel group-gloo.""" +def get_expert_data_parallel_group_gloo(partial_expert_data_parallel=False, check_initialized=True): + """Get expert data parallel group-gloo. + + Args: + check_initialized: When False, return None instead of asserting if the gloo group was + never created. See get_data_parallel_group_gloo. + """ if partial_expert_data_parallel: + if not check_initialized and _INTRA_PARTIAL_EXPERT_DATA_PARALLEL_GROUP_GLOO is None: + return None assert ( _INTRA_PARTIAL_EXPERT_DATA_PARALLEL_GROUP_GLOO is not None ), "Intra partial expert data parallel group-gloo is not initialized" return _INTRA_PARTIAL_EXPERT_DATA_PARALLEL_GROUP_GLOO else: + if not check_initialized and _EXPERT_DATA_PARALLEL_GROUP_GLOO is None: + return None assert ( _EXPERT_DATA_PARALLEL_GROUP_GLOO is not None ), "Expert data parallel group-gloo is not initialized" diff --git a/megatron/core/process_groups_config.py b/megatron/core/process_groups_config.py index 6c1e3651387..ffbfb89e35e 100644 --- a/megatron/core/process_groups_config.py +++ b/megatron/core/process_groups_config.py @@ -133,6 +133,16 @@ class ProcessGroupCollection: # _INTRA_EXPERT_DATA_PARALLEL_GROUP intra_expt_dp: torch.distributed.ProcessGroup = field(init=False) + # Gloo mirrors of the two groups the distributed optimizer needs for checkpoint I/O. + # These have no NCCL equivalent: without them a caller-supplied collection cannot drive + # the optimizer at all. May be None when the job was built with + # initialize_model_parallel(create_gloo_process_groups=False). + # _INTRA_PARTIAL_DATA_PARALLEL_GROUP_WITH_CP_GLOO + intra_dp_cp_gloo: torch.distributed.ProcessGroup = field(init=False) + + # _INTRA_PARTIAL_EXPERT_DATA_PARALLEL_GROUP_GLOO + intra_expt_dp_gloo: torch.distributed.ProcessGroup = field(init=False) + # _INTER_PARTIAL_EXPERT_DATA_PARALLEL_GROUP inter_dist_opt: torch.distributed.ProcessGroup = field(init=False) @@ -229,6 +239,17 @@ def use_mpu_process_groups(cls, required_pgs: Optional[List[str]] = None): check_initialized=False, partial_expert_data_parallel=True, ), + 'intra_dp_cp_gloo': partial( + parallel_state.get_data_parallel_group_gloo, + with_context_parallel=True, + partial_data_parallel=True, + check_initialized=False, + ), + 'intra_expt_dp_gloo': partial( + parallel_state.get_expert_data_parallel_group_gloo, + partial_expert_data_parallel=True, + check_initialized=False, + ), 'inter_dist_opt': partial( parallel_state.get_inter_distributed_optimizer_instance_group, check_initialized=False, @@ -427,14 +448,14 @@ def setup_process_groups_for_optimizer( ) expt_tp_pp_group = pg_collection.tp_ep_pp - # Gloo groups - not supported when pg_collection is provided + # Gloo groups come from the collection. They are optional: a job built with + # create_gloo_process_groups=False legitimately has none. if use_gloo_process_groups: - raise ValueError( - "Gloo process groups are not supported when pg_collection is " - "provided. Please set use_gloo_process_groups to False." - ) - intra_dp_cp_group_gloo = None - intra_expt_dp_group_gloo = None + intra_dp_cp_group_gloo = getattr(pg_collection, 'intra_dp_cp_gloo', None) + intra_expt_dp_group_gloo = getattr(pg_collection, 'intra_expt_dp_gloo', None) + else: + intra_dp_cp_group_gloo = None + intra_expt_dp_group_gloo = None return { 'dp_group': dp_group, diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index 94613b7096c..9a2b67c11f5 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -1340,14 +1340,27 @@ def test_get_megatron_optimizer_custom_process_groups_validation(): config=optimizer_config, model_chunks=model_chunks, pg_collection=pg_collection_complete ) - # Test 6: Gloo process groups should not be used with custom process groups + # Test 6: Gloo process groups are taken from the collection rather than refused. + # Previously this raised "Gloo process groups are not supported when pg_collection is + # provided", which made explicit process-group passing and gloo groups mutually exclusive -- + # and the optimizer path is on every training run. A collection that carries no gloo groups + # (the create_gloo_process_groups=False case) now yields None instead of failing. pg_collection_complete.mp = None # Explicitly set to None as allowed pg_collection_complete.tp_ep_pp = None # Explicitly set to None as allowed - with pytest.raises(ValueError, match="Gloo process groups are not supported"): - get_megatron_optimizer( - config=optimizer_config, - model_chunks=model_chunks, - use_gloo_process_groups=True, # Should be False when using custom groups - pg_collection=pg_collection_complete, - ) + groups = ProcessGroupCollection.setup_process_groups_for_optimizer( + pg_collection_complete, model_chunks, use_gloo_process_groups=True + ) + assert groups['intra_dp_cp_group_gloo'] is None + assert groups['intra_expt_dp_group_gloo'] is None + + # And when the collection does carry them, they are passed straight through. + gloo_dp = torch.distributed.new_group(backend="gloo") + gloo_expt_dp = torch.distributed.new_group(backend="gloo") + pg_collection_complete.intra_dp_cp_gloo = gloo_dp + pg_collection_complete.intra_expt_dp_gloo = gloo_expt_dp + groups = ProcessGroupCollection.setup_process_groups_for_optimizer( + pg_collection_complete, model_chunks, use_gloo_process_groups=True + ) + assert groups['intra_dp_cp_group_gloo'] is gloo_dp + assert groups['intra_expt_dp_group_gloo'] is gloo_expt_dp