From bbdd1ee7f1cd39e6d748d5b0c71b590404fa867e Mon Sep 17 00:00:00 2001 From: ykarnati Date: Wed, 10 Jun 2026 22:04:46 -0700 Subject: [PATCH] Thread pg_collection into train_step reductions Add an optional per-module ProcessGroupCollection to train_step and route its model-parallel reductions, pipeline-last-stage gate, and per-key loss all-reduce through it. Default None preserves today's mpu reads byte-for-byte for non-MIMO callers. Co-Authored-By: Claude Opus 4.8 (1M context) --- megatron/training/training.py | 33 ++++++++++++++------- tests/unit_tests/test_utils.py | 53 ++++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 11 deletions(-) diff --git a/megatron/training/training.py b/megatron/training/training.py index a355b12c4ea..32d3fce7e4a 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2159,8 +2159,12 @@ def dummy_train_step(data_iterator): ) -def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=None): - """Single training step.""" +def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=None, pg_collection: Optional[ProcessGroupCollection] = None): + """Single training step. + + pg_collection: optional per-module :class:`ProcessGroupCollection`; None uses the mpu globals, + otherwise it must define mp, pp, and dp_cp. + """ args = get_args() timers = get_timers() @@ -2298,14 +2302,25 @@ def _save_state_dict(attr_name, label): if save_params_in_this_iteration: _save_state_dict(attr_name="data", label="params") + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + for _required in ("mp", "pp", "dp_cp"): + assert getattr(pg_collection, _required, None) is not None, ( + f"pg_collection passed to train_step must define {_required}" + ) + mp_group = pg_collection.mp + dp_cp_group = pg_collection.dp_cp + is_last_stage = is_pp_last_stage(pg_collection.pp) # when freezing sub-models we may have a mixture of successful and unsucessful ranks, # so we must gather across mp ranks - update_successful = logical_and_across_model_parallel_group(update_successful) + update_successful = logical_and_across_model_parallel_group(update_successful, group=mp_group) # grad_norm and num_zeros_in_grad will be None on ranks without trainable params, # so we must gather across mp ranks - grad_norm = reduce_max_stat_across_model_parallel_group(grad_norm) + grad_norm = reduce_max_stat_across_model_parallel_group(grad_norm, group=mp_group) if args.log_num_zeros_in_grad: - num_zeros_in_grad = reduce_max_stat_across_model_parallel_group(num_zeros_in_grad) + num_zeros_in_grad = reduce_max_stat_across_model_parallel_group( + num_zeros_in_grad, group=mp_group + ) # Vision momentum. if args.vision_pretraining and args.vision_pretraining_type == "dino": @@ -2324,20 +2339,16 @@ def _save_state_dict(attr_name, label): if args.empty_unused_memory_level >= 2: torch.cuda.empty_cache() - if mpu.is_pipeline_last_stage(ignore_virtual=True): + if is_last_stage: # Average loss across microbatches. loss_reduced = {} - for key in losses_reduced[0].keys(): val = [x[key].view(-1) for x in losses_reduced] if val[0].numel() == 2: # there is one dict per microbatch. in new reporting, we average # over the total number of tokens across the global batch. val = torch.vstack(val).sum(dim=0) - torch.distributed.all_reduce( - val, - group=mpu.get_data_parallel_group(with_context_parallel=True) - ) + torch.distributed.all_reduce(val, group=dp_cp_group) loss_reduced[key] = val[0] / val[1] elif val[0].numel() == 1: # legacy behavior, we average over the number of microbatches diff --git a/tests/unit_tests/test_utils.py b/tests/unit_tests/test_utils.py index 5d3b99e8727..94ac440d8e0 100644 --- a/tests/unit_tests/test_utils.py +++ b/tests/unit_tests/test_utils.py @@ -604,3 +604,56 @@ def test_default_falls_back_to_mpu(self): default = reduce_max_stat_across_model_parallel_group(float(rank)) assert default == explicit Utils.destroy_model_parallel() + + +@pytest.mark.skipif(torch.cuda.device_count() < 8, reason="requires 8 GPUs") +class TestTrainStepReductionThreading: + """train_step's reductions over a per-module ProcessGroupCollection's mp/pp/dp_cp groups.""" + + @staticmethod + def _ranks(group): + return torch.distributed.get_process_group_ranks(group) + + def test_grid_threaded_reductions(self): + from megatron.core.hyper_comm_grid import HyperCommGrid + from megatron.core.pipeline_parallel.utils import is_pp_last_stage + from megatron.training.utils.common_utils import ( + logical_and_across_model_parallel_group, + reduce_max_stat_across_model_parallel_group, + ) + + Utils.initialize_distributed() # torch.distributed only; NOT initialize_model_parallel + rank = torch.distributed.get_rank() + # tp=2,pp=2,dp=2 over the world: mp == tp+pp, dp_cp == dp. + grid = HyperCommGrid([2, 2, 2], ["tp", "pp", "dp"], backend="nccl") + grid.create_pg(["tp", "pp"]) + grid.create_pg(["pp"]) + grid.create_pg(["dp"]) + mp_group = grid.get_pg(["tp", "pp"]) + pp_group = grid.get_pg(["pp"]) + dp_cp_group = grid.get_pg(["dp"]) + mp_ranks = self._ranks(mp_group) + pp_ranks = self._ranks(pp_group) + dp_cp_ranks = self._ranks(dp_cp_group) + try: + # reduce_max (grad_norm / num_zeros): MAX over the grid mp group. + assert reduce_max_stat_across_model_parallel_group( + float(rank), group=mp_group + ) == float(max(mp_ranks)) + + # logical_and (update_successful): True iff every mp member is True. + flag = rank != min(mp_ranks) # one rank dissents -> AND is False on the whole group. + assert logical_and_across_model_parallel_group(flag, group=mp_group) is False + + # is_pp_last_stage: True only on the highest rank of the pp group. + assert is_pp_last_stage(pp_group) is (rank == max(pp_ranks)) + + # per-key loss all-reduce: SUM over the grid dp_cp group. + val = torch.tensor([float(rank), 1.0], device=torch.cuda.current_device()) + torch.distributed.all_reduce(val, group=dp_cp_group) + expected = torch.tensor( + [float(sum(dp_cp_ranks)), float(len(dp_cp_ranks))], dtype=torch.float32 + ) + torch.testing.assert_close(val.cpu(), expected) + finally: + grid.destroy()