[2613] Track GPU memory peaks

[florianscheidl]

Description

We add the option to track peaks in allocated and reserved CUDA memory and log these per trainer_logging.metrics steps. We introduce a MemoryTracker class with a similar structure to the ThroughputTracker. Since we are primarily interested in peak memory, we max-reduce the peak memory across devices.

It's currently configured as opt-in in the config (requires that train_logging.memory_tracking=True).

We add tests to check that the tracker works as expected.

Example:

{"weathergen.timestamp": 1783935174691, "weathergen.time": 20260713113254, "stage": "train", "weathergen.step": 280, "performance.memory.train.max_allocated_gib": 15.098231315612793, "performance.memory.train.max_reserved_gib": 17.353515625, ...}

Issue Number

Fixes #2613.

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
  • I have run unit tests with ./scripts/actions.sh unit-test
  • I have documented my code and I have updated the docstrings.
  • I have added unit tests, if relevant
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel