diff --git a/config/config_performance.yml b/config/config_performance.yml
new file mode 100644
index 0000000000..251c224442
--- /dev/null
+++ b/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
diff --git a/config/config_performance_default.yml b/config/config_performance_default.yml
deleted file mode 100644
index 4ac1fafaf2..0000000000
--- a/config/config_performance_default.yml
+++ /dev/null
@@ -1,32 +0,0 @@
-# (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.
-
-# logging config
-train_logging:
-
-  # performance metrics:
-  track_performance_metrics: True
-
-# config for training
-training_config:
-
-  metrics: 2
-
-
-  num_mini_epochs: 10
-  samples_per_mini_epoch: 4096
-  shuffle: True
-
-  model_input: {
-    "forecasting" : {
-      # masking strategy: "random", "healpix", "forecast"
-      masking_strategy: "forecast",
-      num_samples: 8
-      },
-    }
\ No newline at end of file
diff --git a/config/config_performance_jepa.yml b/config/config_performance_jepa.yml
deleted file mode 100644
index eea78e12d3..0000000000
--- a/config/config_performance_jepa.yml
+++ /dev/null
@@ -1,52 +0,0 @@
-# (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.
-
-# logging config
-train_logging:
-
-  # performance metrics:
-  track_performance_metrics: True
-
-# config for training
-training_config:
-  metrics: 2
-
-  num_mini_epochs: 10
-  samples_per_mini_epoch: 8192
-  shuffle: True
-  
-  model_input: {
-    "random_easy" : {
-      num_samples: 2,
-    },
-  }
-
-  losses : {
-    "student-teacher": {
-        enabled: True,
-        type: LossLatentSSLStudentTeacher,
-        weight: 1.0,
-        loss_fcts : {
-          "JEPA": {
-            'weight': 4, "loss_extra_args": {}, "out_dim": 2048, "head": transformer,
-            "num_blocks": 6, "num_heads": 12, "with_qk_lnorm": True, "intermediate_dim": 768, 
-            "dropout_rate": 0.1,
-            target_source_correspondence: {0 : {0 : "subset", 1: "subset"} },
-        },
-        },
-        target_and_aux_calc: { "EMATeacher" : 
-          { ema_ramp_up_ratio : null,
-            ema_halflife_in_thousands: 1e-1,
-            model_param_overrides : { 
-              training_config: { losses: { student-teacher:{ loss_fcts :{JEPA: {head: identity} }}}}
-            },
-          }
-        }
-      }
-  }
\ No newline at end of file
diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py
index 1f506fffb7..f97d7f01e8 100644
--- a/packages/common/src/weathergen/common/config.py
+++ b/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:
diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py
index 7995b5864f..6cf7b0b451 100644
--- a/src/weathergen/run_train.py
+++ b/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)
diff --git a/src/weathergen/train/profiling_trainer.py b/src/weathergen/train/profiling_trainer.py
new file mode 100644
index 0000000000..7c6cea219a
--- /dev/null
+++ b/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)
diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 276da0bd67..94eaf6c930 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -12,6 +12,7 @@
 import copy
 import logging
 import time
+from collections.abc import Iterator
 from math import sqrt
 
 import numpy as np
@@ -24,6 +25,7 @@
 
 import weathergen.common.config as config
 from weathergen.common.config import Config
+from weathergen.datasets.batch import ModelBatch
 from weathergen.datasets.multi_stream_data_sampler import MultiStreamDataSampler
 from weathergen.model.ema import EMAModel
 from weathergen.model.model_interface import (
@@ -48,7 +50,6 @@
     get_target_idxs_from_cfg,
 )
 from weathergen.utils.distributed import is_root
-from weathergen.utils.performance import NullThroughputTracker, ThroughputTracker, nvtx_range
 from weathergen.utils.train_logger import TrainLogger, prepare_losses_for_logging
 from weathergen.utils.utils import get_dtype
 from weathergen.utils.validation_io import write_output
@@ -87,7 +88,6 @@ def __init__(self, train_logging: Config):
         self.batch_size_validation_per_gpu = -1
         self.batch_size_test_per_gpu = -1
         self.collapse_monitor: CollapseMonitor | None = None
-        self.perf_tracker: ThroughputTracker | NullThroughputTracker = NullThroughputTracker()
         self.t_training_start: float = 0
         self.training_loop_annotation_context = contextlib.nullcontext
 
@@ -165,15 +165,6 @@ 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"):
-            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):
-            self.training_loop_annotation_context = nvtx_range
-
     def get_target_aux_calculators(self, mode_cfg):
         """
         Get target_aux_calculators for given mode_cfg
@@ -388,7 +379,7 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
         # training loop
         self.t_training_start = time.time()
 
-        for mini_epoch in range(mini_epoch_base, self.training_cfg.num_mini_epochs):
+        for mini_epoch in self.mini_epochs(mini_epoch_base):
             if is_root():
                 logger.info(
                     f"Mini_epoch {mini_epoch} of {self.training_cfg.num_mini_epochs}: train."
@@ -410,6 +401,15 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
         # log final model
         self.save_model(self.training_cfg.num_mini_epochs)
 
+    def mini_epochs(self, mini_epoch_base: int) -> Iterator[int]:
+        """
+        Yield the mini_epochs that run() iterates over.
+
+        Subclass seam: overriding this changes how long a run lasts without duplicating
+        run(). See weathergen.train.profiling_trainer.ProfilingTrainer.
+        """
+        yield from range(mini_epoch_base, self.training_cfg.num_mini_epochs)
+
     def validate_before_training(self):
         """
         Perform validation before training (eg. to check validation pipeline or data normalization)
@@ -447,7 +447,7 @@ def train(self, mini_epoch):
 
         # training loop
         self.t_start = time.time()
-        for bidx, batch in enumerate(dataset_iter):
+        for bidx, batch in self.train_batches(dataset_iter):
             with self.training_loop_annotation_context(f"batch_{bidx}"):
                 if cf.data_loading.get("memory_pinning", False):
                     # pin memory for faster CPU-GPU transfer
@@ -538,13 +538,6 @@ 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
-                ),
-            )
             # Compute collapse monitoring metrics
             if self.collapse_monitor.should_compute(self.cf.general.istep):
                 self.collapse_monitor._compute_collapse_metrics(
@@ -570,6 +563,16 @@ def train(self, mini_epoch):
 
         self.dataset.advance()
 
+    def train_batches(self, dataset_iter: Iterator) -> Iterator[tuple[int, ModelBatch]]:
+        """
+        Yield the (index, batch) pairs that train() steps over.
+
+        Subclass seam: overriding this bounds the loop or wraps it in a context (e.g. a
+        profiler) while the training step itself stays in train(), so there is only ever
+        one copy of it. See weathergen.train.profiling_trainer.ProfilingTrainer.
+        """
+        yield from enumerate(dataset_iter)
+
     def validate(self, mini_epoch, mode_cfg, batch_size):
         """
         Perform validation / test computation as specified by mode_cfg
diff --git a/src/weathergen/utils/performance.py b/src/weathergen/utils/performance.py
index baf2512e9e..47df43801f 100644
--- a/src/weathergen/utils/performance.py
+++ b/src/weathergen/utils/performance.py
@@ -158,17 +158,6 @@ def compute_metrics(self) -> dict[str, float] | None:
         return metrics
 
 
-class NullThroughputTracker:
-    """No-op throughput tracker used when performance tracking is disabled.
-
-    Implements the same interface as ``ThroughputTracker`` so call sites in the
-    training loop need no ``if`` guards.
-    """
-
-    def step(self, batch, istep: int, log_fn=None) -> None:
-        pass
-
-
 def compute_source_bytes(source_samples) -> int:
     """Count total bytes of all source token tensors in a batch.
 
diff --git a/src/weathergen/utils/profiling.py b/src/weathergen/utils/profiling.py
new file mode 100644
index 0000000000..96e883d729
--- /dev/null
+++ b/src/weathergen/utils/profiling.py
@@ -0,0 +1,313 @@
+# (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.
+
+"""
+Configuration and helpers for measuring a training run.
+
+The `profiling` and `performance_logging` sections of a run config are parsed here; the
+trainer side of both lives in `weathergen.train.profiling_trainer`. Profiling traces a
+bounded stretch of training (PyTorch profiler, CUDA memory snapshot) and is expensive;
+performance logging measures the run as a whole (throughput, later peak memory) and is not.
+"""
+
+import contextlib
+import dataclasses
+import logging
+import platform
+from collections.abc import Callable, Iterator
+from datetime import datetime
+from functools import partial
+from pathlib import Path
+from typing import Protocol
+
+import torch
+from torch.profiler import ProfilerActivity, profile, record_function
+
+import weathergen.common.config as config
+from weathergen.common.config import Config
+from weathergen.utils.distributed import get_rank, is_root
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+TIME_FORMAT_STR: str = "%b_%d_%H_%M_%S"
+MAX_NUM_OF_MEM_EVENTS_PER_SNAPSHOT: int = 100000
+
+
+class BatchTracker(Protocol):
+    """
+    What `ProfilingTrainer` expects of a per-step measurement tool.
+
+    `step` is called once per training step, after that step has completed, on every rank
+    (`ThroughputTracker` and anything else that syncs across ranks relies on that). It is
+    given the batch that was just trained on, the step index it was trained at, and a
+    `log_fn` that writes a metrics dict to the train logger at that step.
+    """
+
+    def step(
+        self, batch, istep: int, log_fn: Callable[[dict[str, float]], None] | None = None
+    ) -> None: ...
+
+
+@dataclasses.dataclass(frozen=True)
+class ProfilingSchedule:
+    """
+    The wait/warmup/active/repeat cycle of the profiled stretch, in training steps.
+
+    It belongs to the profiling section as a whole, not to one collector: the PyTorch
+    profiler steps through it, the memory snapshot records from the first active step
+    onwards, and `num_steps` is how long a `stop_after_profiling` run lasts.
+    """
+
+    wait: int = 1
+    warmup: int = 1
+    active: int = 1
+    repeat: int = 1
+
+    @classmethod
+    def from_config(cls, profiling_cfg: Config | dict) -> "ProfilingSchedule":
+        """Read the schedule from the `profiling` section of a run config."""
+        defaults = cls()
+        return cls(
+            wait=profiling_cfg.get("wait_iteration", defaults.wait),
+            warmup=profiling_cfg.get("warmup_iteration", defaults.warmup),
+            active=profiling_cfg.get("active_iteration", defaults.active),
+            repeat=profiling_cfg.get("repeat", defaults.repeat),
+        )
+
+    @property
+    def num_steps(self) -> int:
+        """Number of training steps needed to walk the full schedule."""
+        return (self.wait + self.warmup + self.active) * self.repeat
+
+    @property
+    def steps_before_active(self) -> int:
+        """Steps run before the first active window, i.e. what collectors should skip."""
+        return self.wait + self.warmup
+
+    def to_torch(self) -> Callable[[int], torch.profiler.ProfilerAction]:
+        return torch.profiler.schedule(
+            wait=self.wait, warmup=self.warmup, active=self.active, repeat=self.repeat
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class ProfilingConfig:
+    """
+    The `profiling` section of a run config: tracing a bounded stretch of training.
+
+    The section is deliberately absent from `config/default_config.yml` — the defaults below
+    are the only ones, so a run config that predates a key (e.g. when continuing an older
+    run) needs no migration. `config/config_performance.yml` documents the keys.
+    """
+
+    # collect traces, which requires the ProfilingTrainer
+    enabled: bool = False
+    # end the run once the profiled stretch is done, instead of training as configured
+    stop_after_profiling: bool = False
+    # how long the profiled stretch is, and how it is split into wait/warmup/active
+    schedule: ProfilingSchedule = ProfilingSchedule()
+    # collectors, independent of each other; each one is opted into explicitly
+    pytorch_profiler: bool = False
+    memory_snapshot: bool = False
+    nvtx_annotate: bool = False
+
+    @classmethod
+    def from_config(cls, cf: Config) -> "ProfilingConfig":
+        cfg = cf.get("profiling") or {}
+        pytorch_profiler_cfg = cfg.get("pytorch_profiler") or {}
+        memory_snapshot_cfg = cfg.get("memory_snapshot") or {}
+        defaults = cls()
+
+        return cls(
+            enabled=cfg.get("enabled", defaults.enabled),
+            stop_after_profiling=cfg.get("stop_after_profiling", defaults.stop_after_profiling),
+            schedule=ProfilingSchedule.from_config(cfg),
+            pytorch_profiler=pytorch_profiler_cfg.get("enabled", defaults.pytorch_profiler),
+            memory_snapshot=memory_snapshot_cfg.get("enabled", defaults.memory_snapshot),
+            nvtx_annotate=cfg.get("nvtx_annotate", defaults.nvtx_annotate),
+        )
+
+    @property
+    def collects_traces(self) -> bool:
+        """Whether the profiled stretch writes anything to the traces directory."""
+        return self.enabled and (self.pytorch_profiler or self.memory_snapshot)
+
+
+@dataclasses.dataclass(frozen=True)
+class PerformanceLoggingConfig:
+    """
+    The `performance_logging` section of a run config: how the run itself performs.
+
+    Unlike profiling, these metrics are cheap, cover the whole run and are logged next to
+    the training metrics rather than written to a trace. Throughput is the only one so far;
+    peak memory is meant to join it. As with `ProfilingConfig`, the defaults below are the
+    only ones; the section is not in `config/default_config.yml`.
+    """
+
+    throughput: bool = False
+    throughput_warmup_steps: int = 2
+
+    @classmethod
+    def from_config(cls, cf: Config) -> "PerformanceLoggingConfig":
+        cfg = cf.get("performance_logging") or {}
+        throughput_cfg = cfg.get("throughput") or {}
+        defaults = cls()
+
+        return cls(
+            throughput=throughput_cfg.get("enabled", defaults.throughput),
+            throughput_warmup_steps=throughput_cfg.get(
+                "warmup_steps", defaults.throughput_warmup_steps
+            ),
+        )
+
+    @property
+    def enabled(self) -> bool:
+        """Whether anything is logged, i.e. whether the run needs the ProfilingTrainer."""
+        return self.throughput
+
+
+@contextlib.contextmanager
+def pytorch_profiler_session(cf: Config, schedule: ProfilingSchedule) -> Iterator[profile | None]:
+    """
+    Run the enclosed block under the PyTorch profiler, on the root rank only.
+
+    Yields the profiler on the root rank (call `.step()` on it once per training step) and
+    None everywhere else. Each cycle of the schedule writes a chrome trace and a memory
+    timeline to `config.get_path_profiling_traces(cf)`; a summary is logged at the end.
+    """
+    if not is_root():
+        yield None
+        return
+
+    traces_path = _prepare_traces_path(cf)
+    with profile(
+        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+        record_shapes=True,
+        profile_memory=True,
+        with_stack=True,
+        with_modules=True,
+        with_flops=True,
+        schedule=schedule.to_torch(),
+        on_trace_ready=partial(trace_handler, cf),
+    ) as prof:
+        yield prof
+
+    log_profiler_summary(prof)
+    logger.info(f"PyTorch profiler traces written to {traces_path}")
+
+
+@contextlib.contextmanager
+def memory_snapshot_session(cf: Config) -> Iterator[None]:
+    """
+    Record the CUDA memory history over the enclosed block, on the root rank only.
+
+    The snapshot is dumped to `config.get_path_profiling_traces(cf)` on exit and can be
+    viewed at https://pytorch.org/memory_viz. Independent of the PyTorch profiler; the
+    caller enters this once the schedule's wait and warmup steps are done.
+    """
+    if not is_root() or not _cuda_available():
+        yield
+        return
+
+    traces_path = _prepare_traces_path(cf)
+    logger.info("Starting snapshot record_memory_history")
+    torch.cuda.memory._record_memory_history(max_entries=MAX_NUM_OF_MEM_EVENTS_PER_SNAPSHOT)
+    try:
+        yield
+        _export_memory_snapshot(cf)
+        logger.info(f"Memory snapshot written to {traces_path}")
+    finally:
+        logger.info("Stopping snapshot record_memory_history")
+        torch.cuda.memory._record_memory_history(enabled=None)
+
+
+def log_profiler_summary(prof: profile) -> None:
+    """Log the aggregated profiler tables (FLOPs, time per module, memory)."""
+    logger.info("\n" + "=" * 80 + "\nPROFILING SUMMARY\n" + "=" * 80)
+
+    logger.info("\n--- Top Operations by FLOPs ---")
+    logger.info(
+        prof.key_averages().table(sort_by="flops", row_limit=20, top_level_events_only=False)
+    )
+
+    logger.info("\n--- Operations Grouped by Module ---")
+    logger.info(
+        prof.key_averages(group_by_stack_n=5).table(sort_by="cuda_time_total", row_limit=30)
+    )
+
+    logger.info("\n--- Memory Usage ---")
+    logger.info(prof.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=20))
+
+
+def trace_handler(cf: Config, prof: profile) -> None:
+    """Write the chrome trace and the memory timeline for one profiler cycle."""
+    file_prefix = _trace_file_prefix(cf)
+
+    prof.export_chrome_trace(f"{file_prefix}.json.gz")
+
+    # the memory timeline relies on kineto functionality unavailable on aarch64
+    if platform.machine() == "aarch64":
+        logger.info("[profiler] Memory distribution timeline skipped on aarch64")
+    else:
+        prof.export_memory_timeline(f"{file_prefix}.html", device="cuda:0")
+
+
+def _export_memory_snapshot(cf: Config) -> None:
+    file_prefix = _trace_file_prefix(cf)
+    try:
+        logger.info(f"Saving snapshot to local file: {file_prefix}.pickle")
+        torch.cuda.memory._dump_snapshot(f"{file_prefix}.pickle")
+    except Exception as e:
+        logger.error(f"Failed to capture memory snapshot {e}")
+
+
+def _prepare_traces_path(cf: Config) -> Path:
+    traces_path = config.get_path_profiling_traces(cf)
+    traces_path.mkdir(exist_ok=True, parents=True)
+    return traces_path
+
+
+def _trace_file_prefix(cf: Config) -> Path:
+    """Timestamped, rank-specific path prefix shared by all profiling artifacts."""
+    timestamp = datetime.now().strftime(TIME_FORMAT_STR)
+    return config.get_path_profiling_traces(cf) / f"{timestamp}_rank_{get_rank()}"
+
+
+def _cuda_available() -> bool:
+    if torch.cuda.is_available():
+        return True
+
+    logger.info("CUDA unavailable. Not recording memory history")
+    return False
+
+
+def wrap_module_forward_with_profiling(model: torch.nn.Module, prefix: str = "") -> None:
+    """
+    Recursively annotate the forward of every custom submodule with `record_function`.
+
+    This makes the trace readable in terms of WeatherGenerator modules instead of bare aten
+    ops. It patches `forward` on the module instances, so only use it on a model that is
+    about to be profiled and then thrown away.
+    """
+    for name, module in model.named_children():
+        module_name = f"{prefix}.{name}" if prefix else name
+
+        # standard PyTorch modules are already traced, but their children may not be
+        if not type(module).__module__.startswith("torch.nn.modules"):
+            module.forward = _profiled_forward(module_name, module.forward)
+
+        wrap_module_forward_with_profiling(module, module_name)
+
+
+def _profiled_forward(module_name: str, forward: Callable) -> Callable:
+    def profiled_forward(*args, **kwargs):
+        with record_function(f"nn.Module: {module_name}"):
+            return forward(*args, **kwargs)
+
+    return profiled_forward
diff --git a/src/weathergen/utils/profiling_test.py b/src/weathergen/utils/profiling_test.py
new file mode 100644
index 0000000000..d5062fd9ca
--- /dev/null
+++ b/src/weathergen/utils/profiling_test.py
@@ -0,0 +1,95 @@
+# (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.
+
+from omegaconf import OmegaConf
+
+from weathergen.common.config import _DEFAULT_CONFIG_PTH
+from weathergen.utils.profiling import (
+    PerformanceLoggingConfig,
+    ProfilingConfig,
+    ProfilingSchedule,
+)
+
+_PERFORMANCE_CONFIG_PTH = _DEFAULT_CONFIG_PTH.parent / "config_performance.yml"
+
+
+def test_schedule_from_config():
+    cfg = OmegaConf.create(
+        {"wait_iteration": 2, "warmup_iteration": 3, "active_iteration": 4, "repeat": 5}
+    )
+    schedule = ProfilingSchedule.from_config(cfg)
+
+    assert schedule == ProfilingSchedule(wait=2, warmup=3, active=4, repeat=5)
+    assert schedule.num_steps == (2 + 3 + 4) * 5
+    assert schedule.steps_before_active == 2 + 3
+
+
+def test_schedule_is_shared_by_the_collectors():
+    """The schedule sits on `profiling`, not on `pytorch_profiler`."""
+    cfg = OmegaConf.create(
+        {
+            "profiling": {
+                "active_iteration": 4,
+                "pytorch_profiler": {"enabled": False},
+                "memory_snapshot": {"enabled": True},
+            }
+        }
+    )
+    profiling_cfg = ProfilingConfig.from_config(cfg)
+
+    assert profiling_cfg.schedule == ProfilingSchedule(active=4)
+    assert not profiling_cfg.pytorch_profiler
+    assert profiling_cfg.memory_snapshot
+
+
+def test_collecting_traces_needs_profiling_enabled():
+    cfg = OmegaConf.create({"profiling": {"enabled": False, "memory_snapshot": {"enabled": True}}})
+
+    assert not ProfilingConfig.from_config(cfg).collects_traces
+
+
+def test_performance_logging_is_independent_of_profiling():
+    cfg = OmegaConf.create({"performance_logging": {"throughput": {"warmup_steps": 5}}})
+    performance_cfg = PerformanceLoggingConfig.from_config(cfg)
+
+    assert not ProfilingConfig.from_config(cfg).enabled
+    assert not performance_cfg.enabled, "throughput is off by default"
+    assert performance_cfg.throughput_warmup_steps == 5
+
+    cfg.performance_logging.throughput.enabled = True
+    assert PerformanceLoggingConfig.from_config(cfg).enabled
+
+
+def test_config_without_the_sections():
+    """
+    Neither section is in default_config.yml, so every run config may be missing them.
+
+    That includes configs of runs that predate a key and are continued, which is why the
+    dataclass defaults are the only defaults.
+    """
+    default_cfg = OmegaConf.load(_DEFAULT_CONFIG_PTH)
+
+    assert "profiling" not in default_cfg
+    assert "performance_logging" not in default_cfg
+    assert ProfilingConfig.from_config(default_cfg) == ProfilingConfig()
+    assert PerformanceLoggingConfig.from_config(default_cfg) == PerformanceLoggingConfig()
+    assert ProfilingConfig.from_config(OmegaConf.create({})) == ProfilingConfig()
+
+
+def test_performance_config_measures_everything():
+    cfg = OmegaConf.merge(
+        OmegaConf.load(_DEFAULT_CONFIG_PTH), OmegaConf.load(_PERFORMANCE_CONFIG_PTH)
+    )
+    profiling_cfg = ProfilingConfig.from_config(cfg)
+
+    assert profiling_cfg.enabled
+    assert profiling_cfg.stop_after_profiling
+    assert profiling_cfg.collects_traces
+    assert profiling_cfg.schedule.num_steps == 5
+    assert PerformanceLoggingConfig.from_config(cfg).enabled