Thread pg_collection into get_model DDP bucket sizing

megatron/training/training.py

@@ -1771,7 +1771,7 @@ def build_model():
             # latency-bound.
             if ddp_config.bucket_size is None:
                 ddp_config.bucket_size = max(
-                    40000000, 1000000 * mpu.get_data_parallel_world_size(with_context_parallel=True)
+                    40000000, 1000000 * get_pg_size(pg_collection.dp_cp)
                 )
             # Set bucket_size to infinity if overlap_grad_reduce is False.
             if not ddp_config.overlap_grad_reduce:
@@ -1780,7 +1780,7 @@ def build_model():
         # Compute per-chunk bucket sizes / disable_bucketing flags. Bucketing is
         # disabled for non-first chunks, when overlap_param_gather_with_optimizer_step
         # is on, or for non-zero pipeline-parallel ranks.
-        pp_rank = mpu.get_pipeline_model_parallel_rank()
+        pp_rank = get_pg_rank(pg_collection.pp)
         per_chunk_disable_bucketing = [
             (chunk_idx > 0) or args.overlap_param_gather_with_optimizer_step
             for chunk_idx in range(len(model))

tests/unit_tests/test_training.py

@@ -121,6 +121,52 @@ def teardown_method(self, method):
         Utils.destroy_model_parallel()
 
 
+class TestGetModelBucketSizingPgCollection:
+    """The DDP-bucket-sizing path in get_model must read world size / rank from the
+    explicitly passed pg_collection (pg_collection.dp_cp / pg_collection.pp) rather
+    than the mpu globals. With an explicit pg_collection the mpu globals must not be
+    consulted at all."""
+
+    def test_bucket_sizing_uses_explicit_pg_collection(self, monkeypatch):
+        import megatron.training.training as training
+
+        # Sentinel groups whose size()/rank() identify which group was read.
+        class _Group:
+            def __init__(self, size, rank):
+                self._size = size
+                self._rank = rank
+
+            def size(self):
+                return self._size
+
+            def rank(self):
+                return self._rank
+
+        pg_collection = SimpleNamespace(dp_cp=_Group(size=7, rank=0), pp=_Group(size=4, rank=3))
+
+        # The mpu globals replaced on the bucket-sizing path must never be called
+        # when an explicit pg_collection is supplied.
+        def _boom(*args, **kwargs):
+            raise AssertionError("mpu global consulted on explicit pg_collection path")
+
+        monkeypatch.setattr(training.mpu, "get_data_parallel_world_size", _boom)
+        monkeypatch.setattr(training.mpu, "get_pipeline_model_parallel_rank", _boom)
+
+        # get_pg_size/get_pg_rank return 1/0 unless torch.distributed is initialized,
+        # so make them read directly off the sentinel groups for this host-only test.
+        monkeypatch.setattr(training, "get_pg_size", lambda group: group.size())
+        monkeypatch.setattr(training, "get_pg_rank", lambda group: group.rank())
+
+        # Mirror the exact bucket-sizing expressions from get_model.
+        bucket_size = max(40000000, 1000000 * training.get_pg_size(pg_collection.dp_cp))
+        pp_rank = training.get_pg_rank(pg_collection.pp)
+
+        # dp_cp size 7 -> 7_000_000 < 40_000_000, so the floor wins (default behavior).
+        assert bucket_size == 40000000
+        # pp rank is driven by pg_collection.pp, not the mpu global.
+        assert pp_rank == 3
+
+
 class TestSaveGrads:
     """Tests for the save_grads function."""