[training] fix: aggregate TB memory metrics across PP group

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)

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."""