diff --git a/src/megatron/bridge/training/utils/train_utils.py b/src/megatron/bridge/training/utils/train_utils.py index 9e1a8c3bcf..ab53eb7c74 100644 --- a/src/megatron/bridge/training/utils/train_utils.py +++ b/src/megatron/bridge/training/utils/train_utils.py @@ -382,6 +382,61 @@ def logical_and_across_model_parallel_group(input: bool, mp_group: "TorchProcess return bool(input.item()) +def reduce_max_memory_across_pp_group( + memory_report: dict[str, Union[int, float]], + pp_group: "TorchProcessGroup", +) -> dict[str, Union[int, float]]: + """Reduce per-rank memory metrics across the PP group with MAX. + + With pipeline parallelism, peak GPU memory is typically dominated by the + first PP stage (activation buildup). The TensorBoard / W&B / MLFlow / Comet + writers, however, only initialize on the last rank (``world_size - 1``), so + without aggregation the logged values reflect only the last PP stage and + under-report true peak headroom. + + This helper performs a single bulk all-reduce with MAX over the PP group + so that the writer rank emits the per-metric peak across the pipeline. + Counter-style integer keys (e.g. ``alloc_retries``) are preserved as + ``int`` so dashboards continue to render them correctly. + + No-op when distributed is uninitialized, the PP group has a single rank, + or the report is empty. + + Args: + memory_report: Mapping of metric name to per-rank value. + pp_group: The pipeline-parallel process group to reduce across. + + Returns: + A new dict with values replaced by the per-metric MAX across the PP + group, or the input report unchanged when no reduction is needed. + """ + if not memory_report: + return memory_report + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return memory_report + pp_size_attr = getattr(pp_group, "size", None) + if not callable(pp_size_attr) or pp_size_attr() <= 1: + return memory_report + + keys = list(memory_report.keys()) + values = torch.tensor( + [memory_report[k] for k in keys], + dtype=torch.float64, + device=torch.cuda.current_device(), + ) + torch.distributed.all_reduce(values, op=torch.distributed.ReduceOp.MAX, group=pp_group) + + reduced: dict[str, Union[int, float]] = {} + for key, max_val in zip(keys, values.tolist()): + original = memory_report[key] + # Preserve int type for counter-style metrics; floats stay as floats. + if not isinstance(original, bool) and isinstance(original, int): + reduced[key] = int(max_val) + else: + reduced[key] = max_val + return reduced + + class _MoeMetricFanoutWriter: """SummaryWriter-shaped adapter that fans add_scalar to MLFlow / Comet. @@ -611,6 +666,17 @@ def training_log( dump(snapshot, f) print_rank_0(f"Saved memory snapshot to {filename}") + # Memory metrics must be aggregated across the PP group BEFORE the + # writer-gated block below. The TensorBoard / W&B / MLFlow / Comet writers + # only initialize on the last rank, but peak GPU memory typically lives on + # the first PP stage. Compute and reduce on all ranks so the writer rank + # emits the per-metric peak across the pipeline (issue #3167). + memory_report: Optional[dict[str, Union[int, float]]] = None + if logger_config.log_memory_to_tensorboard and iteration % logger_config.tensorboard_log_interval == 0: + memory_report = report_memory(memory_keys=logger_config.memory_keys) + memory_report = reduce_max_memory_across_pp_group(memory_report, pg_collection.pp) + memory_report = {f"memory/{mem_stat}": val for (mem_stat, val) in memory_report.items()} + if loggers_exist and iteration % logger_config.tensorboard_log_interval == 0: if logger_config.log_throughput_to_tensorboard: throughput_report = report_throughput( @@ -629,9 +695,7 @@ def training_log( mlflow_logger.log_metrics(_sanitize_mlflow_metrics(throughput_report), step=iteration) if comet_logger: comet_logger.log_metrics(throughput_report, step=iteration) - if logger_config.log_memory_to_tensorboard: - memory_report = report_memory(memory_keys=logger_config.memory_keys) - memory_report = {f"memory/{mem_stat}": val for (mem_stat, val) in memory_report.items()} + if logger_config.log_memory_to_tensorboard and memory_report is not None: if writer: for metric, value in memory_report.items(): writer.add_scalar(metric, value, iteration) diff --git a/tests/unit_tests/training/utils/test_train_utils.py b/tests/unit_tests/training/utils/test_train_utils.py index ea56fed2bf..facc8a9edb 100644 --- a/tests/unit_tests/training/utils/test_train_utils.py +++ b/tests/unit_tests/training/utils/test_train_utils.py @@ -31,6 +31,7 @@ needs_global_state_injection, param_is_not_shared, prepare_forward_step_func, + reduce_max_memory_across_pp_group, report_l2_norm_grad, report_memory, report_runtime, @@ -1728,6 +1729,131 @@ def test_l2_norm_grad(self): assert l2_norm_report["l2_norm/grad/layer_9"] == 9.0 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for this test") +class TestReduceMaxMemoryAcrossPpGroup: + """Test suite for the reduce_max_memory_across_pp_group helper. + + The helper aggregates per-rank memory metrics across the pipeline-parallel + group with MAX so the writer rank emits the per-metric peak across the + pipeline (issue #3167). These tests cover the no-op fallbacks and the + happy-path reduction behavior. + """ + + def test_empty_report_returns_unchanged(self): + """Empty report short-circuits before touching distributed.""" + pp_group = mock.MagicMock() + pp_group.size.return_value = 4 + result = reduce_max_memory_across_pp_group({}, pp_group) + assert result == {} + pp_group.size.assert_not_called() + + def test_distributed_uninitialized_returns_unchanged(self): + """When torch.distributed is not initialized, return input as-is.""" + report = {"peak_allocated_gigabytes": 12.5, "alloc_retries": 3} + pp_group = mock.MagicMock() + pp_group.size.return_value = 4 + + with mock.patch("torch.distributed.is_initialized", return_value=False): + result = reduce_max_memory_across_pp_group(report, pp_group) + + assert result == report + pp_group.size.assert_not_called() + + def test_pp_size_one_returns_unchanged(self): + """A single-rank PP group bypasses the all-reduce.""" + report = {"peak_allocated_gigabytes": 7.0} + pp_group = mock.MagicMock() + pp_group.size.return_value = 1 + + with ( + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.all_reduce") as mock_all_reduce, + ): + result = reduce_max_memory_across_pp_group(report, pp_group) + + assert result == report + mock_all_reduce.assert_not_called() + + def test_pp_group_missing_size_returns_unchanged(self): + """A defensive check: if the group has no callable .size, no-op.""" + report = {"peak_allocated_gigabytes": 4.5} + + # An object without .size attribute at all. + class _Bare: + pass + + with ( + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.all_reduce") as mock_all_reduce, + ): + result = reduce_max_memory_across_pp_group(report, _Bare()) + + assert result == report + mock_all_reduce.assert_not_called() + + def test_max_reduction_across_pp_ranks(self): + """All-reduce MAX is invoked once and replaces values with the max.""" + report = { + "peak_allocated_gigabytes": 10.0, + "peak_reserved_gigabytes": 12.5, + } + pp_group = mock.MagicMock() + pp_group.size.return_value = 4 + + # Simulate the in-place all-reduce by writing the per-element max + # values directly into the input tensor. + def _fake_all_reduce(tensor, op, group): + assert op == torch.distributed.ReduceOp.MAX + assert group is pp_group + # Pretend rank-0 had higher peak across the pipeline. + tensor.copy_(torch.tensor([14.25, 18.0], dtype=tensor.dtype, device=tensor.device)) + + with ( + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.all_reduce", side_effect=_fake_all_reduce) as mock_all_reduce, + ): + result = reduce_max_memory_across_pp_group(report, pp_group) + + assert mock_all_reduce.call_count == 1 + assert result == { + "peak_allocated_gigabytes": pytest.approx(14.25), + "peak_reserved_gigabytes": pytest.approx(18.0), + } + # Original report must not be mutated. + assert report == { + "peak_allocated_gigabytes": 10.0, + "peak_reserved_gigabytes": 12.5, + } + + def test_int_counters_remain_int_after_reduction(self): + """Counter-style integer metrics (e.g. alloc_retries) stay as int.""" + report = { + "peak_allocated_gigabytes": 8.0, + "alloc_retries": 1, + } + pp_group = mock.MagicMock() + pp_group.size.return_value = 2 + + def _fake_all_reduce(tensor, op, group): + tensor.copy_(torch.tensor([9.5, 4.0], dtype=tensor.dtype, device=tensor.device)) + + with ( + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.all_reduce", side_effect=_fake_all_reduce), + ): + result = reduce_max_memory_across_pp_group(report, pp_group) + + assert isinstance(result["peak_allocated_gigabytes"], float) + assert result["peak_allocated_gigabytes"] == pytest.approx(9.5) + # `alloc_retries` was an int on input, so it must remain an int. + assert isinstance(result["alloc_retries"], int) + assert result["alloc_retries"] == 4 + + class TestNeedsGlobalStateInjection: """Test suite for the needs_global_state_injection function."""