Summary

• Fix TensorBoard / W&B / MLFlow / Comet memory metrics under-reporting by reducing report_memory(...) across the pipeline-parallel group with MAX before the writer rank emits.
• Add reduce_max_memory_across_pp_group helper and call it from training_log on every PP rank, gated by the existing log_memory_to_tensorboard config flag plus the tensorboard log interval — single bulk all-reduce per logged step on PP > 1, no-op otherwise.
• Counter-style integer keys (e.g. alloc_retries) are preserved as int so existing dashboards continue to render correctly.
• Adds 6 unit tests for the helper covering the no-op fallbacks, happy-path bulk MAX, and the int-preservation contract.

Refs #3167. Partially addresses the TB-rank concern in #3164.

Background

tensorboard_logger, wandb_logger, mlflow_logger, and comet_logger are lazily initialized only on the last rank (world_size - 1), see state.py:182-228. With pipeline parallelism, that's the last PP rank — but peak GPU memory typically lives on the first PP stage (activation buildup), so the dashboards under-report true peak headroom and obscure how close a job is to OOM.

The cleanest fix is to aggregate per-rank torch.cuda.memory_stats() across the PP group with MAX, so the writer rank always emits the per-metric peak across the entire pipeline — strictly more useful than picking any single PP rank.

Implementation Notes

• The all-reduce must run on every PP rank, not just the writer rank. The new pre-block runs unconditionally when log_memory_to_tensorboard=True (a config flag, identical on all ranks) and the iteration matches the tensorboard log interval. The existing loggers_exist-gated section then consumes the precomputed dict on the writer rank only.
• Single bulk all-reduce over a stacked float64 tensor instead of one collective per metric — minimizes overhead.
• Helper is defensive: short-circuits if memory_report is empty, torch.distributed is uninitialized, the PP group has size 1, or the group object lacks a callable .size. Keeps unit-test paths and single-stage runs unaffected.
• Original memory_report dict is not mutated.

Test plan

• python3 -m ast parse of changed files
• ruff check clean on changed files
• ruff format --check clean on changed files
• 6 new unit tests added under TestReduceMaxMemoryAcrossPpGroup:
  • empty report short-circuits
  • distributed uninitialized → no-op
  • pp.size() == 1 → no-op (no all-reduce call)
  • PP group missing .size → defensive no-op
  • max reduction replaces values; original dict not mutated
  • int metrics (alloc_retries) preserved as int
• CI: existing test_memory_tensorboard_logging continues to pass (helper short-circuits when torch.distributed.is_initialized() is False in unit tests)
• CI: L0/L1 functional tests on H100 / GB200
• Manual: with PP > 1, observe TB memory/peak_* scalars now reflect the per-metric max across pipeline stages instead of the last stage only

Risk

Low. The change is isolated to training_log's memory-logging block:

• No behavior change when PP size == 1 or distributed is uninitialized.
• No behavior change when log_memory_to_tensorboard=False.
• The all-reduce participates on every PP rank only when the writer rank would log — no new collectives outside the existing log interval.
• No public API change.