diff --git a/megatron/training/training.py b/megatron/training/training.py index 3930cc46a21..fb50e7db60f 100644 --- a/megatron/training/training.py +++ b/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)) diff --git a/tests/unit_tests/test_training.py b/tests/unit_tests/test_training.py index 838b963778c..7901e409a42 100644 --- a/tests/unit_tests/test_training.py +++ b/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."""