[1804] feat(training): refactored profiler pr

config/config_performance.yml

@@ -0,0 +1,53 @@
+# Overwrite config for measuring a training run, e.g.
+#   train --config config/config_performance.yml
+# Both sections below are handled by the ProfilingTrainer
+# (src/weathergen/train/profiling_trainer.py), which replaces the Trainer as soon as either
+# of them asks for something. They are absent from default_config.yml: every key falls back
+# to the default given in its comment (ProfilingConfig / PerformanceLoggingConfig in
+# src/weathergen/utils/profiling.py), so an overwrite config only sets what it changes.
+#
+# As written, the run ends once the profiled stretch is done, and its chrome trace, memory
+# timeline and memory snapshot are written to
+# <shared_working_dir>/logs/<run_id>/profiling_traces.
+
+# tracing a bounded stretch of training; expensive
+profiling:
+
+  # collect traces; without this, nothing else in this section has any effect (default False)
+  enabled: True
+
+  # end the run once the profiled stretch is done, without validating or checkpointing, so
+  # that the traces cover the training step and nothing else. Leave at the default False to
+  # trace the beginning of an otherwise normal run.
+  stop_after_profiling: True
+
+  # the profiled stretch, in training steps: (wait_iteration + warmup_iteration +
+  # active_iteration) * repeat steps are run, the collectors below skip the wait and warmup
+  # ones, and a stop_after_profiling run ends after them. Each defaults to 1.
+  wait_iteration: 1
+  warmup_iteration: 1
+  active_iteration: 3
+  repeat: 1
+
+  # PyTorch profiler: chrome trace and memory timeline per cycle, root rank only
+  pytorch_profiler:
+    enabled: True  # default False
+
+  # CUDA memory history snapshot, root rank only, recorded from the first active step on.
+  # View at https://pytorch.org/memory_viz . Independent of pytorch_profiler.
+  memory_snapshot:
+    enabled: True  # default False
+
+  # annotate batches and model blocks with nvtx ranges, for nsys
+  # nvtx_annotate: False
+
+# how the run itself performs: cheap metrics over the whole run, logged next to the training
+# metrics on all ranks. Independent of profiling — a run with only these enabled trains
+# exactly as it would without them, so they can be left on for a full-length run.
+performance_logging:
+
+  # throughput metrics (performance.throughput.*)
+  throughput:
+    enabled: True  # default False
+    # steps to skip before reporting, so that startup does not skew the numbers
+    # warmup_steps: 2

packages/common/src/weathergen/common/config.py

@@ -703,6 +703,11 @@ def get_path_run(config: Config) -> Path:
     return _get_shared_wg_path() / "results" / get_run_id_from_config(config)
 
 
+def get_path_profiling_traces(config: Config) -> Path:
+    """Get the path for storing profiling traces."""
+    return _get_shared_wg_path() / "logs" / get_run_id_from_config(config) / "profiling_traces"
+
+
 def get_path_model(config: Config | None = None, run_id: str | None = None) -> Path:
     """Get the current runs model_path for storing model checkpoints."""
     if config or run_id:

src/weathergen/run_train.py

@@ -21,12 +21,24 @@
 
 import weathergen.common.config as config
 import weathergen.utils.cli as cli
+from weathergen.common.config import Config
 from weathergen.common.logger import init_loggers
+from weathergen.train.profiling_trainer import ProfilingTrainer
 from weathergen.train.trainer import Trainer
+from weathergen.utils.profiling import PerformanceLoggingConfig, ProfilingConfig
 
 logger = logging.getLogger(__name__)
 
 
+def get_trainer(cf: Config) -> Trainer:
+    """Select the trainer: the ProfilingTrainer if the run is measured, a plain one otherwise."""
+    if ProfilingConfig.from_config(cf).enabled or PerformanceLoggingConfig.from_config(cf).enabled:
+        logger.info("Profiling or performance logging enabled: running with ProfilingTrainer.")
+        return ProfilingTrainer(cf.train_logging)
+
+    return Trainer(cf.train_logging)
+
+
 def train() -> None:
     """Entry point for calling the training code from the command line."""
     main([cli.Stage.train] + sys.argv[1:])
@@ -144,7 +156,7 @@ def run_continue(args):
     # track history of run to ensure traceability of results
     cf.general.run_history += [(args.from_run_id, cf.general.istep)]
 
-    trainer = Trainer(cf.train_logging)
+    trainer = get_trainer(cf)
 
     try:
         trainer.run(cf, devices, args.from_run_id, args.mini_epoch)
@@ -185,7 +197,7 @@ def run_train(args):
     if cf.with_flash_attention:
         assert cf.with_mixed_precision
 
-    trainer = Trainer(cf.train_logging)
+    trainer = get_trainer(cf)
 
     try:
         trainer.run(cf, devices)

src/weathergen/train/profiling_trainer.py

@@ -0,0 +1,197 @@
+# (C) Copyright 2025 WeatherGenerator contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+"""
+Trainer variant that measures a training run instead of (only) performing it.
+
+Selected by the `profiling` and `performance_logging` config sections, see
+`weathergen.run_train.get_trainer`.
+"""
+
+import contextlib
+import logging
+from collections.abc import Iterator
+from itertools import islice
+
+import torch
+
+from weathergen.common.config import Config
+from weathergen.datasets.batch import ModelBatch
+from weathergen.train.trainer import Trainer
+from weathergen.train.utils import TRAIN
+from weathergen.utils.distributed import is_root
+from weathergen.utils.performance import ThroughputTracker, nvtx_range
+from weathergen.utils.profiling import (
+    BatchTracker,
+    PerformanceLoggingConfig,
+    ProfilingConfig,
+    memory_snapshot_session,
+    pytorch_profiler_session,
+    wrap_module_forward_with_profiling,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ProfilingTrainer(Trainer):
+    """
+    Trainer that measures the training loop, configured by `profiling` and
+    `performance_logging`.
+
+    The training step itself is inherited unchanged from `Trainer`: only the iteration
+    seams (`mini_epochs`, `train_batches`) are overridden, so the measured code path and
+    the code path of a normal run cannot drift apart. Everything that measures a step hangs
+    off `train_batches`, which regains control once the step for the batch it yielded is
+    done — `Trainer` therefore knows about no measurement tool at all.
+
+    `profiling` traces the profiled stretch (`schedule.num_steps` training steps) on the
+    root rank, while the other ranks run the same steps untraced so that collectives stay
+    matched. The PyTorch profiler steps through the schedule; the memory snapshot records
+    from the first active step onwards. With `stop_after_profiling` (the default) the run
+    ends once the stretch is done, without validating or checkpointing, so that the traces
+    cover the training step and nothing else.
+
+    `performance_logging` builds the `BatchTracker`s (see `get_trackers`) that measure every
+    step of the whole run on every rank. It is cheap: a run with only this enabled trains
+    exactly as a plain `Trainer` run would, and just logs more.
+    """
+
+    def __init__(self, train_logging: Config):
+        super().__init__(train_logging)
+
+        self.profiling_cfg = ProfilingConfig()
+        self.performance_cfg = PerformanceLoggingConfig()
+        self.trackers: list[BatchTracker] = []
+        self.profiling_done: bool = False
+
+    def init(self, cf: Config, devices: list) -> None:
+        super().init(cf, devices)
+
+        self.profiling_cfg = ProfilingConfig.from_config(self.cf)
+        self.performance_cfg = PerformanceLoggingConfig.from_config(self.cf)
+        logger.info(f"Profiling run: {self.profiling_cfg}, {self.performance_cfg}")
+
+        self.trackers = self.get_trackers()
+        if self.profiling_cfg.nvtx_annotate:
+            self.training_loop_annotation_context = nvtx_range
+
+    def get_trackers(self) -> list[BatchTracker]:
+        """
+        Build the per-step measurement tools the `performance_logging` config asks for.
+
+        This is where a new tracking tool is added: implement `BatchTracker` and append it
+        here. The trainer only ever calls `step` on them, once per training step and on
+        every rank, so a tracker is free to sync across ranks.
+        """
+        trackers: list[BatchTracker] = []
+
+        if self.performance_cfg.throughput:
+            trackers.append(
+                ThroughputTracker(
+                    device=torch.device(self.devices[0]),
+                    warmup_steps=self.performance_cfg.throughput_warmup_steps,
+                    batch_size_per_gpu=self.batch_size_per_gpu,
+                )
+            )
+
+        return trackers
+
+    @property
+    def stops_after_profiling(self) -> bool:
+        """Whether the run exists only to be profiled, and ends once it is."""
+        return self.profiling_cfg.enabled and self.profiling_cfg.stop_after_profiling
+
+    def mini_epochs(self, mini_epoch_base: int) -> Iterator[int]:
+        """Run a single mini_epoch when the run only exists to be profiled."""
+        if not self.stops_after_profiling:
+            yield from super().mini_epochs(mini_epoch_base)
+            return
+
+        yield mini_epoch_base
+
+    def train_batches(self, dataset_iter: Iterator) -> Iterator[tuple[int, ModelBatch]]:
+        """Measure every training step, and trace the profiled stretch of them."""
+        yield from self._tracked(self._profiled(dataset_iter))
+
+    def _tracked(
+        self, batches: Iterator[tuple[int, ModelBatch]]
+    ) -> Iterator[tuple[int, ModelBatch]]:
+        """
+        Step the trackers once per training step, on every rank.
+
+        Control returns here after `train()` has finished the step for the batch that was
+        yielded, which is what lets the measurement live outside the training step.
+        """
+        if not self.trackers:
+            yield from batches
+            return
+
+        for bidx, batch in batches:
+            istep = self.cf.general.istep  # train() increments it as part of the step
+            yield bidx, batch
+            for tracker in self.trackers:
+                tracker.step(
+                    batch,
+                    istep,
+                    log_fn=lambda m, istep=istep: self.train_logger.log_metrics(
+                        TRAIN, m, step=istep
+                    ),
+                )
+
+    def _profiled(self, dataset_iter: Iterator) -> Iterator[tuple[int, ModelBatch]]:
+        """Trace the profiled stretch, then continue (or stop) as configured."""
+        if self.profiling_done or not self.profiling_cfg.enabled:
+            # the stretch is profiled once per run, not once per mini_epoch
+            yield from super().train_batches(dataset_iter)
+            return
+
+        self.profiling_done = True
+        schedule = self.profiling_cfg.schedule
+
+        if is_root() and self.profiling_cfg.pytorch_profiler:
+            # the model only exists once run() has built it, hence not in init()
+            wrap_module_forward_with_profiling(self.model, prefix="model")
+
+        with contextlib.ExitStack() as stack:
+            prof = None
+            if self.profiling_cfg.pytorch_profiler:
+                prof = stack.enter_context(pytorch_profiler_session(self.cf, schedule))
+
+            for bidx, batch in enumerate(islice(dataset_iter, schedule.num_steps)):
+                if bidx == schedule.steps_before_active and self.profiling_cfg.memory_snapshot:
+                    # skip the wait and warmup steps, as the profiler does
+                    stack.enter_context(memory_snapshot_session(self.cf))
+
+                yield bidx, batch
+                if prof is not None:
+                    prof.step()
+
+        # keep the other ranks in step with the root rank writing its traces
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+
+        if self.stops_after_profiling:
+            logger.info(f"Profiled {schedule.num_steps} training steps, ending the run.")
+            return
+
+        yield from enumerate(dataset_iter, start=schedule.num_steps)
+
+    def validate(self, mini_epoch, mode_cfg, batch_size) -> None:
+        """Skipped while the run only exists to be profiled."""
+        if self.stops_after_profiling:
+            return
+
+        super().validate(mini_epoch, mode_cfg, batch_size)
+
+    def save_model(self, mini_epoch: int, name=None) -> None:
+        """Skipped while the run trains too few steps for its checkpoints to be useful."""
+        if self.stops_after_profiling:
+            return
+
+        super().save_model(mini_epoch, name)