Thread pg_collection into train_step reductions

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

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()