Give gloo process groups a home in ProcessGroupCollection

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'):

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"

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,

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