[2624] Refactor throughput logging

packages/evaluate/src/weathergen/evaluate/scores/score.py

@@ -16,8 +16,8 @@
 import xarray as xr
 from scipy.spatial import cKDTree
 
-from weathergen.evaluate.scores.score_utils import calc_latitude_weights, to_list
 from weathergen.evaluate.scores.psd import compute_psd_score, detect_grid_type
+from weathergen.evaluate.scores.score_utils import calc_latitude_weights, to_list
 
 # from common.io import MockIO
 

src/weathergen/train/loss_modules/loss_module_physical.py

@@ -295,7 +295,9 @@ def compute_loss(self, preds: dict, targets: dict, metadata) -> LossValues:
             if self.dynamic_loss_ema.enabled and weights_channels is not None:
                 losses_all[stream_name][str(self.forecast_offset)]["mse_ema_weight"] = {}
                 for ch_n, w in zip(target_channels, weights_channels, strict=True):
-                    losses_all[stream_name][str(self.forecast_offset)]["mse_ema_weight"][ch_n] = w.item()
+                    losses_all[stream_name][str(self.forecast_offset)]["mse_ema_weight"][ch_n] = (
+                        w.item()
+                    )
 
             # TODO: make nicer
             output_step_loss_weights = self._get_output_step_weights(len(targets.output_idxs))

src/weathergen/train/trainer.py

@@ -165,10 +165,9 @@ def init(self, cf: Config, devices):
         collapse_config = cf.train_logging.get("collapse_monitoring", {})
         self.collapse_monitor = CollapseMonitor(collapse_config, None)  # device set later in run()
 
-        if cf.train_logging.get("track_performance_metrics"):
+        if cf.train_logging.get("throughput_tracking"):
             self.perf_tracker = ThroughputTracker(
                 device=torch.device(self.devices[0]),
-                warmup_steps=cf.train_logging.get("performance_tracking_warmup_steps", 2),
                 batch_size_per_gpu=self.batch_size_per_gpu,
             )
         if cf.get("profiling", {}).get("nvtx_annotate", False):
@@ -538,13 +537,8 @@ def train(self, mini_epoch):
             if self.validate_with_ema:
                 self.ema_model.update(self.cf.general.istep * batch_size_total, batch_size_total)
 
-            self.perf_tracker.step(
-                batch,
-                self.cf.general.istep,
-                log_fn=lambda m: self.train_logger.log_metrics(
-                    TRAIN, m, step=self.cf.general.istep
-                ),
-            )
+            # Accumulate throughput counts every step; no sync or collective here.
+            self.perf_tracker.step(batch)
             # Compute collapse monitoring metrics
             if self.collapse_monitor.should_compute(self.cf.general.istep):
                 self.collapse_monitor._compute_collapse_metrics(
@@ -557,7 +551,12 @@ def train(self, mini_epoch):
 
             self._log_terminal(bidx, mini_epoch, TRAIN)
             if bidx % self.train_logging.metrics == 0:
-                self._log(TRAIN)
+                # Reduce throughput once per interval and merge it into the
+                # training metrics record. The cross-rank collective (and its
+                # device sync) happens only here, not per step; it returns None
+                # until the first counts have accumulated.
+                perf_metrics = self.perf_tracker.compute_metrics()
+                self._log(TRAIN, extra_metrics=perf_metrics)
                 # Log collapse metrics
                 if self.collapse_monitor.should_log(self.cf.general.istep):
                     self._log_collapse_metrics(TRAIN)
@@ -729,13 +728,15 @@ def save_model(self, mini_epoch: int, name=None):
             # save config
             config.save(self.cf, mini_epoch)
 
-    def _log(self, stage: Stage):
+    def _log(self, stage: Stage, extra_metrics: dict[str, float] | None = None):
         """
         Logs training or validation metrics.
 
         Args:
             stage: Stage Is it's VAL, logs are treated as validation logs.
                         If TRAIN, logs are treated as training logs
+            extra_metrics: Additional scalar metrics (e.g. throughput stats) to
+                        merge into the same record instead of a separate log line.
 
         Notes:
             - This method only executes logging on the main process (rank 0).
@@ -753,7 +754,9 @@ def _log(self, stage: Stage):
         if is_root():
             # plain logger
             if stage == VAL:
-                self.train_logger.add_logs(stage, samples, losses_all, stddev_all)
+                self.train_logger.add_logs(
+                    stage, samples, losses_all, stddev_all, extra_metrics=extra_metrics
+                )
 
             elif self.cf.general.istep >= 0:
                 elapsed_time = time.time() - self.t_training_start
@@ -765,6 +768,7 @@ def _log(self, stage: Stage):
                     avg_loss=avg_loss,
                     lr=self.lr_scheduler.get_lr(),
                     elapsed_training_time_seconds=elapsed_time,
+                    extra_metrics=extra_metrics,
                 )
 
         loss_calculator.loss_hist = []

src/weathergen/utils/performance.py

@@ -10,152 +10,94 @@
 """Utilities for measuring training throughput metrics."""
 
 import logging
-import time
-from collections.abc import Callable
 from contextlib import contextmanager
 
 import torch
 
-from weathergen.utils.distributed import is_root
-
 logger = logging.getLogger(__name__)
 
 
 class ThroughputTracker:
     """Tracks training throughput metrics.
 
-    Accumulates per-batch sample and source-byte counts across ranks, with the warmup
-    / accumulation logic required to produce stable global throughput metrics.
+    Accumulates per-batch sample and source-byte counts across ranks.
     """
 
     def __init__(
         self,
         device: torch.device,
-        warmup_steps: int,
         batch_size_per_gpu: int,
     ) -> None:
         self._device = device
-        self._warmup_steps = warmup_steps
         self.batch_size_per_gpu = batch_size_per_gpu
-        self._t0: float | None = None
-        self._warmup_done: bool = False
         self._total_batches: int = 0
         self._total_samples: int = 0
         self._total_mb: float = 0.0
-        self._synced_elapsed: float | None = None
-        self._synced_global_batches: int = 0
-        self._synced_global_samples: int = 0
-        self._synced_global_mb: float = 0.0
 
-    def step(
-        self,
-        batch,
-        istep: int,
-        log_fn: Callable[[dict[str, float]], None] | None = None,
-    ) -> None:
-        """Record one training step and optionally log metrics.
+    def step(self, batch) -> None:
+        """Accumulate one training step's counts. No synchronization or collectives.
 
-        Wrapper around ``update`` and ``compute_metrics`` that also computes
-        source bytes from the batch on the fly. When metrics are available and
-        the current rank is root, ``log_fn`` is called with the metrics dict.
+        Call on every step from the training loop. Metrics are computed separately
+        via ``compute_metrics`` at the logging interval, so the hot path stays free
+        of device syncs and cross-rank collectives.
 
         Args:
             batch: The current training batch (must expose ``get_source_samples()``).
-            batch_size_per_gpu: Number of samples processed on this rank.
-            istep: Global training step index.
-            log_fn: Called with the metrics dict on the root rank once warmup is
-                    complete. Typically ``lambda m: logger.log_metrics(stage, m, step=istep)``.
         """
         source_mb = compute_source_bytes(batch.get_source_samples()) / 1e6
-        self.update(istep, source_mb)
-        self._sync()  # collective: all ranks must participate
-        if log_fn is not None and is_root():
-            metrics = self.compute_metrics()
-            if metrics is not None:
-                log_fn(metrics)
-
-    def update(self, istep: int, source_mb: float) -> None:
+        self.update(source_mb)
+
+    def update(self, source_mb: float) -> None:
         """Record one training step, handling warmup internally.
 
+        Purely local bookkeeping: no device synchronization. The cumulative
+        counts are turned into throughput (and reduced across ranks) only when
+        ``compute_metrics`` runs.
+
         Args:
-            batch_size_per_gpu: Number of samples processed on this rank.
-            istep: Global training step index (used for warmup countdown).
             source_mb: Source tensor megabytes for this batch. Should be computed
                        fresh each step via ``compute_source_bytes`` as batch sizes
                        can vary across samples.
         """
-        if not self._warmup_done:
-            if istep >= self._warmup_steps - 1:
-                self._t0 = time.time()
-                self._warmup_done = True
-        else:
-            torch.cuda.synchronize()
-            self._total_batches += 1
-            self._total_samples += self.batch_size_per_gpu
-            self._total_mb += source_mb
-
-    def _sync(self) -> None:
-        """Collective: reduce per-rank counters across all ranks and cache the result.
-
-        Must be called on every rank at the same point in the training loop.
-        The cached values are later read by ``compute_metrics()`` on the root rank.
-        """
-        if self._total_batches == 0 or self._t0 is None:
-            return
-
-        elapsed = time.time() - self._t0
-
-        global_batches = torch.tensor(self._total_batches, dtype=torch.int64, device=self._device)
-        global_samples = torch.tensor(self._total_samples, dtype=torch.int64, device=self._device)
-        global_total_mb = torch.tensor(self._total_mb, dtype=torch.float32, device=self._device)
-
-        if torch.distributed.is_available() and torch.distributed.is_initialized():
-            elapsed_tensor = torch.tensor(elapsed, dtype=torch.float32, device=self._device)
-            torch.distributed.all_reduce(elapsed_tensor, op=torch.distributed.ReduceOp.AVG)
-            elapsed = elapsed_tensor.item()
-
-            torch.distributed.all_reduce(global_batches)
-            torch.distributed.all_reduce(global_samples)
-            torch.distributed.all_reduce(global_total_mb)
-
-        self._synced_elapsed = elapsed
-        self._synced_global_batches = int(global_batches.item())
-        self._synced_global_samples = int(global_samples.item())
-        self._synced_global_mb = global_total_mb.item()
+        self._total_batches += 1
+        self._total_samples += self.batch_size_per_gpu
+        self._total_mb += source_mb
 
     def compute_metrics(self) -> dict[str, float] | None:
-        """Return performance metrics dict, or None if warmup is not yet complete.
+        """Return throughput metrics dict, or None if warmup is not yet complete.
+
+        Collective: performs a single SUM all-reduce of the per-device throughput
+        to obtain the global throughput, so it must be called on every rank at the
+        same point in the training loop. The returned dict is identical on all
+        ranks. Global throughput is the sum of the per-device rates across ranks.
 
         Returns:
             Dict of ``"performance.<key>": value`` pairs, or None if no data yet.
         """
-        if self._total_batches == 0 or self._t0 is None:
-            return None
-        elapsed = time.time() - self._t0
-
-        if elapsed <= 0 or self._synced_elapsed is None or self._synced_elapsed <= 0:
+        if self._total_batches == 0:
             return None
 
-        metrics: dict[str, float] = {}
+        device_batches = self._total_batches
+        device_samples = self._total_samples
+        device_mb = self._total_mb
 
-        # Device-level throughput (this rank only).
-        metrics["performance.throughput.device.batches_per_sec"] = self._total_batches / elapsed
-        metrics["performance.throughput.device.samples_per_sec"] = self._total_samples / elapsed
-        metrics["performance.throughput.device.mb_per_sec"] = self._total_mb / elapsed
-
-        # Global throughput: use values already reduced across all ranks by _sync().
-        synced_elapsed = self._synced_elapsed
-        metrics["performance.throughput.global.batches_per_sec"] = (
-            self._synced_global_batches / synced_elapsed
-        )
-        metrics["performance.throughput.global.samples_per_sec"] = (
-            self._synced_global_samples / synced_elapsed
-        )
-        metrics["performance.throughput.global.mb_per_sec"] = (
-            self._synced_global_mb / synced_elapsed
-        )
+        # Global throughput: sum the per-device rates across ranks with a single
+        # reduce, done here (once per logging interval) rather than per step.
+        global_batches, global_samples, global_mb = device_batches, device_samples, device_mb
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            rates = torch.tensor(
+                [device_batches, device_samples, device_mb],
+                dtype=torch.float64,
+                device=self._device,
+            )
+            torch.distributed.all_reduce(rates, op=torch.distributed.ReduceOp.SUM)
+            global_batches, global_samples, global_mb = rates.tolist()
 
-        return metrics
+        return {
+            "performance.throughput.global.batches": global_batches,
+            "performance.throughput.global.samples": global_samples,
+            "performance.throughput.global.mb": global_mb,
+        }
 
 
 class NullThroughputTracker:
@@ -165,9 +107,12 @@ class NullThroughputTracker:
     training loop need no ``if`` guards.
     """
 
-    def step(self, batch, istep: int, log_fn=None) -> None:
+    def step(self, batch) -> None:
         pass
 
+    def compute_metrics(self) -> dict[str, float] | None:
+        return None
+
 
 def compute_source_bytes(source_samples) -> int:
     """Count total bytes of all source token tensors in a batch.

src/weathergen/utils/train_logger.py

@@ -102,9 +102,14 @@ def add_logs(
         avg_loss: list[float] = None,
         lr: float = None,
         elapsed_training_time_seconds: float | None = None,
+        extra_metrics: dict[str, float] | None = None,
     ) -> None:
         """
         Log training or validation data.
+
+        Args:
+            extra_metrics: Additional scalar metrics (e.g. throughput stats) to
+                merge into the same record instead of emitting a separate log line.
         """
         metrics: dict[str, float] = dict(num_samples=samples)
 
@@ -129,6 +134,9 @@ def add_logs(
             val = np.nan if np.isnan(value).all() else np.nanmean(value)
             metrics[key] = val
 
+        if extra_metrics:
+            metrics.update(extra_metrics)
+
         self.log_metrics(stage, metrics)
 
     #######################################