Repository: ecmwf/anemoi-training (language: python)

Issue:
Fix interactive multi-GPU training
This is a proposed fix for interactive multi-GPU training, which is currently broken.

"interactive" in this context means: you call `anemoi-training train` directly in a shell.

### Background

Hydra expects full control over `sys.argv` , with `argv[0]` being the executable, and `argv[1:]` the hydra options.
This is incompatible with our subcommands, where `argv[0]` is `anemoi-training` and `argv[1]` is something like `train`. Hydra will try to ingest `train` as an option, and crash.

To get around this, we overwrite the argv so that `argv[0]` contains the executable and subcommand: `anemoi-training train`.

This makes hydra happy, and it works fine with 1 GPU.

### What works

For most training runs, we run training via Slurm. For this, we don't call the training command directly but we do `srun anemoi-training train` , typically from an sbatch script (but can also be done directly).

The above works fine, because under the hood Slurm will take the argument of srun and execute it once for each GPU. PTL then takes care of the synchronisation between the processes. This can be seen here:

![image](https://github.com/user-attachments/assets/3e67caed-953c-4883-bebc-1d0c1871df26)

One `srun` call results in 2 `anemoi-training train` processes being spawned, one for each GPU (`num_gpus_per_node=2`).

So far so good! 

### What doesn't work

When we run `anemoi-training train` interactively, that is without Slurm, things break. 

This is because without Slurm, PTL will take over responsibility of spawning the subprocesses for each GPU. And as can be seen [here](https://github.com/Lightning-AI/pytorch-lightning/blob/8ad3e29816a63d8ce5c00ac104b14729a4176f4f/src/lightning/fabric/strategies/launchers/subprocess_script.py#L173), it does this by doing a `subprocess.Popen` on the value in `sys.argv[0]`. 

At first glance this looks like it should work, after all `argv[0]` contains `anemoi-training train`. 
But that binary does not exist. So when it tries to open `anemoi-training train` it crashes:

```
/envs/anemoi/bin/python3.10: can't open file '/envs/anemoi/bin/anemoi-training train': [Errno 2] No such file or directory
```

So there is a conflict with the argv. We hack `argv[0]` to make hydra happy, but at the same time it needs to contain a valid binary so that PTL can spawn the DDP processes.

### A solution

One simple solution is to put a value in `argv[0]` that points to a valid binary.

In this PR, I introduce a new "hidden" entry point `.anemoi-training-train` that points to the hydra main function.

Here you can see it working, the top process for GPU 1 is `anemoi-training train`, while the subprocess for GPU 2 is `.anemoi-training-train`:
![image](https://github.com/user-attachments/assets/800a7e39-c188-4d5c-9e34-61cd0d0d2ff6)


I put the dot at the beginning so that it does not get suggested by autocomplete, because we don't really want users to use this entry point directly. That is also why I added an environment variable to check if the entry point is being called by the subcommand, and not directly.

I then also use that environment variable to log the original command to mlflow. I did not want to touch the top Anemoi trainer interface, so I opted for an env var instead of passing an argument all the way through to the logger.

This is not the cleanest solution. Having this "hidden" entry point is not ideal, so I'm open to suggestions. But I'm in favour of accepting a less-than-ideal solution if it solves the problem in the short run, and then improve it.

This also means that if we have any other subcommands invoking the trainer, we need to have a corresponding "hidden" entry point (e.g.: the profiler).

closes #80 


Discussion / hints:

Fix patch:
diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 00000000..a19ade07
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,1 @@
+CHANGELOG.md merge=union
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
new file mode 100644
index 00000000..74bdac0a
--- /dev/null
+++ b/.github/CODEOWNERS
@@ -0,0 +1,6 @@
+# CODEOWNERS file
+
+# Protect workflow files
+/.github/ @theissenhelen @jesperdramsch @gmertes @b8raoult @floriankrb @anaprietonem @HCookie @JPXKQX @mchantry
+/.pre-commit-config.yaml @theissenhelen @jesperdramsch @gmertes @b8raoult @floriankrb @anaprietonem @HCookie @JPXKQX @mchantry
+/pyproject.toml @theissenhelen @jesperdramsch @gmertes @b8raoult @floriankrb @anaprietonem @HCookie @JPXKQX @mchantry
diff --git a/.github/ci-config.yml b/.github/ci-config.yml
new file mode 100644
index 00000000..6138e636
--- /dev/null
+++ b/.github/ci-config.yml
@@ -0,0 +1,9 @@
+dependencies: |
+  ecmwf/ecbuild
+  MathisRosenhauer/libaec@master
+  ecmwf/eccodes
+  ecmwf/eckit
+  ecmwf/odc
+dependency_branch: develop
+parallelism_factor: 8
+self_build: false # Only for python packages
diff --git a/.github/ci-hpc-config.yml b/.github/ci-hpc-config.yml
new file mode 100644
index 00000000..b4c16267
--- /dev/null
+++ b/.github/ci-hpc-config.yml
@@ -0,0 +1,19 @@
+build:
+  python: '3.10'
+  modules:
+    - ninja
+  dependencies:
+    - ecmwf/ecbuild@develop
+    - ecmwf/eccodes@develop
+    - ecmwf/eckit@develop
+    - ecmwf/odc@develop
+  python_dependencies:
+    - ecmwf/anemoi-utils@develop
+    - ecmwf/anemoi-datasets@develop
+    - ecmwf/anemoi-models@develop
+    - ecmwf/anemoi-graphs@develop
+  parallel: 64
+
+  pytest_cmd: |
+    python -m pytest -vv -m 'not notebook and not no_cache_init' --cov=. --cov-report=xml
+    python -m coverage report
diff --git a/.github/workflows/changelog-pr-update.yml b/.github/workflows/changelog-pr-update.yml
new file mode 100644
index 00000000..73cb1ebf
--- /dev/null
+++ b/.github/workflows/changelog-pr-update.yml
@@ -0,0 +1,18 @@
+name: Check Changelog Update on PR
+on:
+  pull_request:
+    types: [assigned, opened, synchronize, reopened, labeled, unlabeled]
+    branches:
+      - main
+      - develop
+    paths-ignore:
+      - .pre-commit-config.yaml
+      - .readthedocs.yaml
+jobs:
+  Check-Changelog:
+    name: Check Changelog Action
+    runs-on: ubuntu-20.04
+    steps:
+      - uses: tarides/changelog-check-action@v2
+        with:
+          changelog: CHANGELOG.md
diff --git a/.github/workflows/changelog-release-update.yml b/.github/workflows/changelog-release-update.yml
new file mode 100644
index 00000000..a31c06e8
--- /dev/null
+++ b/.github/workflows/changelog-release-update.yml
@@ -0,0 +1,35 @@
+---
+# .github/workflows/update-changelog.yaml
+name: "Update Changelog"
+
+on:
+  release:
+    types: [released]
+
+permissions:
+  pull-requests: write
+  contents: write
+
+jobs:
+  update:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v4
+      with:
+        ref: ${{ github.event.release.target_commitish }}
+
+    - name: Update Changelog
+      uses: stefanzweifel/changelog-updater-action@v1
+      with:
+        latest-version: ${{ github.event.release.tag_name }}
+        heading-text: ${{ github.event.release.name }}
+
+    - name: Create Pull Request
+      uses: peter-evans/create-pull-request@v6
+      with:
+        branch: docs/changelog-update-${{ github.event.release.tag_name }}
+        title: '[Changelog] Update to ${{ github.event.release.tag_name }}'
+        add-paths: |
+          CHANGELOG.md
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 00000000..e861091e
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,52 @@
+name: ci
+
+on:
+  # Trigger the workflow on push to main or develop, except tag creation
+  push:
+    branches:
+      - 'main'
+      - 'develop'
+    tags-ignore:
+      - '**'
+    paths-ignore:
+      - "docs/**"
+      - "CHANGELOG.md"
+      - "README.md"
+
+  # Trigger the workflow on pull request
+  pull_request:
+    paths-ignore:
+      - "docs/**"
+      - "CHANGELOG.md"
+      - "README.md"
+
+  # Trigger the workflow manually
+  workflow_dispatch: ~
+
+  # Trigger after public PR approved for CI
+  pull_request_target:
+    types: [labeled]
+    paths-ignore:
+      - "docs/**"
+      - "CHANGELOG.md"
+      - "README.md"
+
+jobs:
+  # Run CI including downstream packages on self-hosted runners
+  downstream-ci:
+    name: downstream-ci
+    if: ${{ !github.event.pull_request.head.repo.fork && github.event.action != 'labeled' || github.event.label.name == 'approved-for-ci' }}
+    uses: ecmwf-actions/downstream-ci/.github/workflows/downstream-ci.yml@main
+    with:
+      anemoi-training: ecmwf/anemoi-training@${{ github.event.pull_request.head.sha || github.sha }}
+      codecov_upload: true
+    secrets: inherit
+
+   # Build downstream packages on HPC
+  downstream-ci-hpc:
+    name: downstream-ci-hpc
+    if: ${{ !github.event.pull_request.head.repo.fork && github.event.action != 'labeled' || github.event.label.name == 'approved-for-ci' }}
+    uses: ecmwf-actions/downstream-ci/.github/workflows/downstream-ci-hpc.yml@main
+    with:
+      anemoi-training: ecmwf/anemoi-training@${{ github.event.pull_request.head.sha || github.sha }}
+    secrets: inherit
diff --git a/.github/workflows/label-public-pr.yml b/.github/workflows/label-public-pr.yml
new file mode 100644
index 00000000..59b2bfa2
--- /dev/null
+++ b/.github/workflows/label-public-pr.yml
@@ -0,0 +1,10 @@
+# Manage labels of pull requests that originate from forks
+name: label-public-pr
+
+on:
+  pull_request_target:
+    types: [opened, synchronize]
+
+jobs:
+  label:
+    uses: ecmwf-actions/reusable-workflows/.github/workflows/label-pr.yml@v2
diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
index 3ca7ce0c..db316803 100644
--- a/.github/workflows/python-publish.yml
+++ b/.github/workflows/python-publish.yml
@@ -1,84 +1,23 @@
+---
 # This workflow will upload a Python Package using Twine when a release is created
 # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
 
 name: Upload Python Package
 
 on:
-
-  push: {}
-
   release:
     types: [created]
 
 jobs:
   quality:
-    name: Code QA
-    runs-on: ubuntu-latest
-    steps:
-    - run: sudo apt-get install -y pandoc # Needed by sphinx for notebooks
-    - uses: actions/checkout@v4
-    - uses: actions/setup-python@v5
-      with:
-        python-version: 3.x
-    - uses: pre-commit/action@v3.0.1
+    uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-precommit-run.yml@v2
+    with:
+      skip-hooks: "no-commit-to-branch"
 
   checks:
-    strategy:
-      fail-fast: false
-      matrix:
-        platform: ["ubuntu-latest", "macos-latest"]
-        python-version: ["3.10"]
-
-    name: Python  ${{ matrix.python-version }} on ${{ matrix.platform }}
-    runs-on: ${{ matrix.platform }}
-
-    steps:
-    - uses: actions/checkout@v4
-
-    - uses: actions/setup-python@v2
-      with:
-        python-version: ${{ matrix.python-version }}
-
-    - name: Install
-      run: |
-        pip install pytest
-        pip install -e .[all]
-        pip install -r tests/requirements.txt
-        pip freeze
-
-    - name: Tests
-      run: pytest
+    uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-pytest-pyproject.yml@v2
 
   deploy:
-
-    if: ${{ github.event_name == 'release' }}
-    runs-on: ubuntu-latest
     needs: [checks, quality]
-
-    steps:
-    - uses: actions/checkout@v4
-
-    - name: Set up Python
-      uses: actions/setup-python@v2
-      with:
-        python-version: '3.10'
-
-    - name: Check that tag version matches code version
-      run: |
-        tag=${GITHUB_REF#refs/tags/}
-        version=$(python setup.py --version)
-        echo 'tag='$tag
-        echo "version file="$version
-        test "$tag" == "$version"
-
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install setuptools wheel twine
-    - name: Build and publish
-      env:
-        TWINE_USERNAME: __token__
-        TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
-      run: |
-        python setup.py sdist
-        twine upload dist/*
+    uses: ecmwf-actions/reusable-workflows/.github/workflows/cd-pypi.yml@v2
+    secrets: inherit
diff --git a/.github/workflows/python-pull-request.yml b/.github/workflows/python-pull-request.yml
new file mode 100644
index 00000000..3488f55c
--- /dev/null
+++ b/.github/workflows/python-pull-request.yml
@@ -0,0 +1,19 @@
+---
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
+
+name: Code Quality checks for PRs
+
+on:
+  push:
+  pull_request:
+    types: [opened, synchronize, reopened]
+
+jobs:
+  quality:
+    uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-precommit-run.yml@v2
+    with:
+      skip-hooks: "no-commit-to-branch"
+
+  checks:
+    uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-pytest-pyproject.yml@v2
diff --git a/.github/workflows/readtthedocs-pr-update.yml b/.github/workflows/readtthedocs-pr-update.yml
new file mode 100644
index 00000000..b8951ab6
--- /dev/null
+++ b/.github/workflows/readtthedocs-pr-update.yml
@@ -0,0 +1,22 @@
+name: Read the Docs PR Preview
+on:
+  pull_request_target:
+    types:
+      - opened
+      - synchronize
+      - reopened
+    # Execute this action only on PRs that touch
+    # documentation files.
+    paths:
+      - "docs/**"
+
+permissions:
+  pull-requests: write
+
+jobs:
+  documentation-links:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: readthedocs/actions/preview@v1
+        with:
+          project-slug: "anemoi-training"
diff --git a/.gitignore b/.gitignore
index fdc24f30..d58e92cc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -121,6 +121,7 @@ celerybeat.pid
 
 # Environments
 .env
+.envrc
 .venv
 env/
 venv/
@@ -184,3 +185,5 @@ _build/
 ?.*
 ~*
 *.sync
+_version.py
+outputs/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 7b01a875..e01d6a37 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,19 +1,16 @@
 repos:
-
 # Empty notebookds
 - repo: local
   hooks:
-    - id: clear-notebooks-output
-      name: clear-notebooks-output
-      files: tools/.*\.ipynb$
-      stages: [commit]
-      language: python
-      entry: jupyter nbconvert --ClearOutputPreprocessor.enabled=True --inplace
-      additional_dependencies: [jupyter]
-
-
+  - id: clear-notebooks-output
+    name: clear-notebooks-output
+    files: tools/.*\.ipynb$
+    stages: [pre-commit]
+    language: python
+    entry: jupyter nbconvert --ClearOutputPreprocessor.enabled=True --inplace
+    additional_dependencies: [jupyter]
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.4.0
+  rev: v5.0.0
   hooks:
   - id: check-yaml # Check YAML files for syntax errors only
     args: [--unsafe, --allow-multiple-documents]
@@ -23,13 +20,17 @@ repos:
   - id: no-commit-to-branch # Prevent committing to main / master
   - id: check-added-large-files # Check for large files added to git
   - id: check-merge-conflict # Check for files that contain merge conflict
-
+-   repo: https://github.com/pre-commit/pygrep-hooks
+    rev: v1.10.0  # Use the ref you want to point at
+    hooks:
+    -   id: python-use-type-annotations # Check for missing type annotations
+    -   id: python-check-blanket-noqa # Check for # noqa: all
+    -   id: python-no-log-warn # Check for log.warn
 - repo: https://github.com/psf/black-pre-commit-mirror
-  rev: 24.1.1
+  rev: 24.8.0
   hooks:
-    - id: black
-      args: [--line-length=120]
-
+  - id: black
+    args: [--line-length=120]
 - repo: https://github.com/pycqa/isort
   rev: 5.13.2
   hooks:
@@ -38,26 +39,47 @@ repos:
     - -l 120
     - --force-single-line-imports
     - --profile black
-
-
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.3.0
+  rev: v0.6.9
   hooks:
   - id: ruff
-    exclude: '(dev/.*|.*_)\.py$'
+    # Next line if for documenation cod snippets
+    exclude: '^[^_].*_\.py$'
     args:
     - --line-length=120
     - --fix
     - --exit-non-zero-on-fix
     - --preview
-
 - repo: https://github.com/sphinx-contrib/sphinx-lint
-  rev: v0.9.1
+  rev: v1.0.0
   hooks:
-    - id: sphinx-lint
-
+  - id: sphinx-lint
 # For now, we use it. But it does not support a lot of sphinx features
 - repo: https://github.com/dzhu/rstfmt
   rev: v0.0.14
   hooks:
-    - id: rstfmt
+  - id: rstfmt
+    exclude: 'cli/.*' # Because we use argparse
+- repo: https://github.com/b8raoult/pre-commit-docconvert
+  rev: "0.1.5"
+  hooks:
+  - id: docconvert
+    args: ["numpy"]
+- repo: https://github.com/tox-dev/pyproject-fmt
+  rev: "2.2.4"
+  hooks:
+  - id: pyproject-fmt
+-   repo: https://github.com/jshwi/docsig # Check docstrings against function sig
+    rev: v0.64.0
+    hooks:
+    -   id: docsig
+        args:
+        - --ignore-no-params # Allow docstrings without parameters
+        - --check-dunders    # Check dunder methods
+        - --check-overridden # Check overridden methods
+        - --check-protected  # Check protected methods
+        - --check-class      # Check class docstrings
+        - --disable=E113     # Disable empty docstrings
+        - --summary          # Print a summary
+ci:
+  autoupdate_schedule: monthly
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
index 45958d1f..c03429e5 100644
--- a/.readthedocs.yaml
+++ b/.readthedocs.yaml
@@ -3,11 +3,15 @@ version: 2
 build:
   os: ubuntu-22.04
   tools:
-    python: "3.10"
+    python: "3.11"
 
 sphinx:
   configuration: docs/conf.py
 
 python:
    install:
-     - requirements: docs/requirements.txt
+    - requirements: docs/requirements.txt
+    - method: pip
+      path: .
+      extra_requirements:
+        - docs
diff --git a/CHANGELOG.md b/CHANGELOG.md
new file mode 100644
index 00000000..3b233a77
--- /dev/null
+++ b/CHANGELOG.md
@@ -0,0 +1,124 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+Please add your functional changes to the appropriate section in the PR.
+Keep it human-readable, your future self will thank you!
+
+## [Unreleased](https://github.com/ecmwf/anemoi-training/compare/0.1.0...HEAD)
+
+### Added
+
+- Add anemoi-transform link to documentation
+- Codeowners file (#56)
+- Changelog merge strategy (#56)
+
+#### Miscellaneous
+
+- Introduction of remapper to anemoi-models leads to changes in the data indices. Some preprocessors cannot be applied in-place anymore.
+
+#### Functionality
+
+- Enable the callback for plotting a histogram for variables containing NaNs
+- Enforce same binning for histograms comparing true data to predicted data
+- Fix: Inference checkpoints are now saved according the frequency settings defined in the config [#37](https://github.com/ecmwf/anemoi-training/pull/37)
+- Feature: Add configurable models [#50](https://github.com/ecmwf/anemoi-training/pulls/50)
+- Feature: Authentication support for mlflow sync - [#51](https://github.com/ecmwf/anemoi-training/pull/51)
+- Feature: Support training for datasets with missing time steps [#48](https://github.com/ecmwf/anemoi-training/pulls/48)
+- Feature: `AnemoiMlflowClient`, an mlflow client with authentication support [#86](https://github.com/ecmwf/anemoi-training/pull/86)
+- Long Rollout Plots
+
+### Fixed
+
+- Fix `TypeError` raised when trying to JSON serialise `datetime.timedelta` object - [#43](https://github.com/ecmwf/anemoi-training/pull/43)
+- Bugfixes for CI (#56)
+- Fix `mlflow` subcommand on python 3.9 [#62](https://github.com/ecmwf/anemoi-training/pull/62)
+- Show correct subcommand in MLFlow - Addresses [#39](https://github.com/ecmwf/anemoi-training/issues/39) in [#61](https://github.com/ecmwf/anemoi-training/pull/61)
+- Fix interactive multi-GPU training [#82](https://github.com/ecmwf/anemoi-training/pull/82)
+- Allow 500 characters in mlflow logging [#88](https://github.com/ecmwf/anemoi-training/pull/88)
+
+### Changed
+
+- Updated configuration examples in documentation and corrected links - [#46](https://github.com/ecmwf/anemoi-training/pull/46)
+- Remove credential prompt from mlflow login, replace with seed refresh token via web - [#78](https://github.com/ecmwf/anemoi-training/pull/78)
+- Update CODEOWNERS
+
+## [0.1.0 - Anemoi training - First release](https://github.com/ecmwf/anemoi-training/releases/tag/0.1.0) - 2024-08-16
+
+### Added
+
+#### Subcommands
+
+- Subcommand for training `anemoi-training train`
+- Subcommand for config generation of configs
+- Subcommand for mlflow: login and sync
+- Subcommand for checkpoint handling
+
+#### Functionality
+- Searchpaths for Hydra configs, to enable configs in CWD, `ANEMOI_CONFIG_PATH` env, and `.config/anemoi/training` in addition to package defaults
+- MlFlow token authentication
+- Configurable pressure level scaling
+
+#### Continuous Integration / Deployment
+
+- Downstream CI to test all dependencies with changes
+- Changelog Status check
+- Readthedocs PR builder
+- Changelog Release Updater Workflow
+
+#### Miscellaneous
+
+- Extended ruff Ruleset
+- Added Docsig pre-commit hook
+- `__future__` annotations for typehints
+- Added Typehints where missing
+- Added Changelog
+- Correct errors in callback plots
+- fix error in the default config
+- example slurm config
+- ability to configure precip-type plots
+
+### Changed
+
+#### Move to Anemoi Ecosystem
+
+- Fixed PyPI packaging
+- Use of Anemoi models
+- Use of Anemoi graphs
+- Adjusted tests to work with new Anemoi ecosystem
+- Adjusted configs to reasonable common defaults
+
+#### Functionality
+
+- Changed hardware-specific keys from configs to `???` to trigger "missing"
+- `__len__` of NativeGridDataset
+- Configurable dropout in attention layer
+
+#### Docs
+
+- First draft on Read the Docs
+- Fixed docstrings
+
+#### Miscellaneous
+
+- Moved callbacks into folder to fascilitate future refactor
+- Adjusted PyPI release infrastructure to common ECMWF workflow
+- Bumped versions in Pre-commit hooks
+- Fix crash when logging hyperparameters with missing values in the config
+- Fixed "null" tracker metadata when tracking is disabled, now returns an empty dict
+- Pinned numpy<2 until we can test all migration
+- (ci): path ignore of docs for downstream ci
+- (ci): remove yaml anchor, unsupported by Github
+- ci: make python QA reusable
+- ci: permissions on changelog updater
+
+### Removed
+
+- Dependency on mlflow-export-import
+- Specific user configs
+- **len** function of NativeGridDataset as it lead to bugs
+
+<!-- Add Git Diffs for Links above -->
diff --git a/README.md b/README.md
index b20a7be4..5ee7892f 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,4 @@
-# anemoi-utils
+# anemoi-training
 
 **DISCLAIMER**
 This project is **BETA** and will be **Experimental** for the foreseeable future.
diff --git a/docs/conf.py b/docs/conf.py
index 208e3d06..12e25dd5 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -14,13 +14,9 @@
 import os
 import sys
 
-sys.path.insert(0, os.path.abspath(".."))
-
-
 read_the_docs_build = os.environ.get("READTHEDOCS", None) == "True"
 
-# top = os.path.realpath(os.path.dirname(os.path.dirname(__file__)))
-# sys.path.insert(0, top)
+sys.path.insert(0, os.path.join(os.path.abspath(".."), "src"))
 
 
 source_suffix = ".rst"
@@ -32,20 +28,22 @@
 
 # -- Project information -----------------------------------------------------
 
-project = "Anemoi"
+project = "Anemoi Training"
 
 author = "ECMWF"
 
 year = datetime.datetime.now().year
-if year == 2024:
-    years = "2024"
-else:
-    years = "2024-%s" % (year,)
+years = "2024" if year == 2024 else f"2024-{year}"
+
+copyright = f"{years}, ECMWF"
 
-copyright = "%s, ECMWF" % (years,)
 
+try:
+    from anemoi.training._version import __version__
 
-release = "0.1.0"
+    release = __version__
+except ImportError:
+    release = "0.0.0"
 
 
 # -- General configuration ---------------------------------------------------
@@ -61,6 +59,8 @@
     "sphinx.ext.intersphinx",
     "sphinx.ext.autodoc",
     "sphinx.ext.napoleon",
+    "sphinxarg.ext",
+    "sphinx.ext.autosectionlabel",
 ]
 
 # Add any paths that contain templates here, relative to this directory.
@@ -93,12 +93,20 @@
         "https://anemoi-inference.readthedocs.io/en/latest/",
         ("../../anemoi-inference/docs/_build/html/objects.inv", None),
     ),
+    "anemoi-graphs": (
+        "https://anemoi-graphs.readthedocs.io/en/latest/",
+        ("../../anemoi-graphs/docs/_build/html/objects.inv", None),
+    ),
+    "anemoi-registry": (
+        "https://anemoi-registry.readthedocs.io/en/latest/",
+        ("../../anemoi-registry/docs/_build/html/objects.inv", None),
+    ),
+    "anemoi-transform": (
+        "https://anemoi-transform.readthedocs.io/en/latest/",
+        ("../../anemoi-transform/docs/_build/html/objects.inv", None),
+    ),
 }
 
-
-# https://www.notion.so/Deepnote-Launch-Buttons-63c642a5e875463495ed2341e83a4b2a
-
-
 # -- Options for HTML output -------------------------------------------------
 
 # The theme to use for HTML and HTML Help pages.  See the documentation for
@@ -114,3 +122,5 @@
 
 
 todo_include_todos = not read_the_docs_build
+
+autodoc_member_order = "bysource"  # Keep file order
diff --git a/docs/dev/code-structure.rst b/docs/dev/code-structure.rst
new file mode 100644
index 00000000..2baf42e6
--- /dev/null
+++ b/docs/dev/code-structure.rst
@@ -0,0 +1,77 @@
+################
+ Code Structure
+################
+
+Understanding and maintaining the code structure is crucial for
+sustainable development of Anemoi Training. This guide outlines best
+practices for contributing to the codebase.
+
+******************************
+ Subclassing for New Features
+******************************
+
+When creating a new feature, the recommended practice is to subclass
+existing base classes rather than modifying them directly. This approach
+preserves functionality for other users while allowing for
+customization.
+
+Example:
+========
+
+In `anemoi/training/diagnostics/callbacks.py`, the `BasePlotCallback`
+serves as a foundation for other plotting callbacks. New plotting
+features should subclass this base class.
+
+*******************
+ File Organization
+*******************
+
+When developing multiple new functions for a feature:
+
+#. Create a new file in the folder (e.g., `callbacks/<new_feature>.py`)
+   to avoid confusion with base functions.
+#. Group related functionality together for better organization and
+   maintainability.
+
+********************************
+ Version Control Best Practices
+********************************
+
+#. Always use pre-commit hooks to ensure code quality and consistency.
+#. Never commit directly to the `develop` branch.
+#. Create a new branch for your feature or bug fix, e.g.,
+   `feature/<feature_name>` or `bugfix/<bug_name>`.
+#. Submit a Pull Request from your branch to `develop` for peer review
+   and testing.
+
+******************************
+ Code Style and Documentation
+******************************
+
+#. Follow PEP 8 guidelines for Python code style, the pre-commit hooks
+   will help enforce this.
+#. Write clear, concise docstrings for all classes and functions using
+   the Numpy style.
+#. Use type hints to improve code readability and catch potential
+   errors.
+#. Add inline comments for complex logic or algorithms.
+
+*********
+ Testing
+*********
+
+#. Write unit tests for new features using pytest.
+#. Ensure all existing tests pass before submitting a Pull Request.
+#. Aim for high test coverage, especially for critical functionality.
+
+****************************
+ Performance Considerations
+****************************
+
+#. Profile your code to identify performance bottlenecks.
+#. Optimize critical paths and frequently called functions.
+#. Consider using vectorized operations when working with large
+   datasets.
+
+By following these guidelines, you'll contribute to a maintainable and
+robust codebase for Anemoi Training.
diff --git a/docs/dev/contributing.rst b/docs/dev/contributing.rst
new file mode 100644
index 00000000..436d9de2
--- /dev/null
+++ b/docs/dev/contributing.rst
@@ -0,0 +1,149 @@
+##############
+ Contributing
+##############
+
+Thank you for your interest in contributing to Anemoi Training! This
+guide will help you get started with the development process.
+
+****************************************
+ Setting Up the Development Environment
+****************************************
+
+#. Clone the repository:
+
+   .. code:: bash
+
+      git clone https://github.com/ecmwf/anemoi-training/
+      cd anemoi-training
+
+#. Install dependencies:
+
+   .. code:: bash
+
+      # For all dependencies
+      pip install -e .
+
+      # For development dependencies
+      pip install -e '.[dev]'
+
+#. (macOS only) Install pandoc for documentation building:
+
+   .. code:: bash
+
+      brew install pandoc
+
+******************
+ Pre-Commit Hooks
+******************
+
+We use pre-commit hooks to ensure code quality and consistency. To set
+them up:
+
+#. Install pre-commit hooks:
+
+   .. code:: bash
+
+      pre-commit install
+
+#. Run hooks on all files to verify installation:
+
+   .. code:: bash
+
+      pre-commit run --all-files
+
+*******************
+ Commit Guidelines
+*******************
+
+Ideally, open an issue for the feature or bug fix you're working on
+before starting development, to discuss the approach with maintainers.
+
+When committing code changes:
+
+#. Make small, focused commits with clear and concise messages.
+
+#. Follow the `Conventional Commits guidelines
+   <https://www.conventionalcommits.org/>`_, e.g., "feat:", "fix:",
+   "docs:", etc.
+
+#. Use present tense and imperative mood in commit messages (e.g., "Add
+   feature" not "Added feature").
+
+#. Reference relevant issue numbers in commit messages when applicable.
+
+#. Update the ``CHANGELOG.md`` file with a human-friendly summary of
+   your changes.
+
+**********************
+ Pull Request Process
+**********************
+
+#. Create a new branch for your feature or bug fix.
+#. Make your changes and commit them using the guidelines above.
+#. Push your branch to your fork on GitHub.
+#. Open a Pull Request against the `develop` branch of the main
+   repository.
+#. Ensure all tests pass and the code adheres to the project's style
+   guidelines.
+#. Request a review from maintainers or other contributors.
+
+***************
+ Running Tests
+***************
+
+We use pytest for our test suite. To run tests:
+
+.. code:: bash
+
+   # Run all tests
+   pytest
+
+   # Run tests in a specific file
+   pytest tests/test_<file>.py
+
+Note: Some tests, like `test_gnn.py`, may run slower on CPU and are
+better suited for GPU execution.
+
+************************
+ Building Documentation
+************************
+
+You can build the documentation locally to preview changes before
+submitting a Pull Request. We use Sphinx for documentation.
+
+You can install the dependencies for building the documentation with:
+
+.. code:: bash
+
+   pip install '.[docs]'
+
+To build the documentation locally:
+
+.. code:: bash
+
+   cd docs
+   make html
+
+The generated documentation will be in `docs/_build/html/index.html`.
+
+*********************
+ Code Review Process
+*********************
+
+#. All code changes must be reviewed before merging.
+#. Address any feedback or comments from reviewers promptly.
+#. Once approved, a maintainer will merge your Pull Request.
+
+******************
+ Reporting Issues
+******************
+
+If you encounter a bug or have a feature request:
+
+#. Check the existing issues to avoid duplicates.
+#. If it's a new issue, create a detailed bug report or feature request.
+#. Use clear, descriptive titles and provide as much relevant
+   information as possible.
+
+Thank you for contributing to Anemoi Training! Your efforts help improve
+the project for everyone.
diff --git a/docs/dev/hydra.rst b/docs/dev/hydra.rst
new file mode 100644
index 00000000..5902eee3
--- /dev/null
+++ b/docs/dev/hydra.rst
@@ -0,0 +1,146 @@
+###############
+ Configuration
+###############
+
+Anemoi Training uses Hydra for configuration management, allowing for
+flexible and modular configuration of the training pipeline. This guide
+explains how to use Hydra effectively in the project.
+
+**************
+ Hydra Basics
+**************
+
+Hydra is a framework for elegantly configuring complex applications. It
+allows for:
+
+#. Hierarchical configuration
+#. Configuration composition
+#. Dynamic object instantiation
+
+*********************************
+ Object Instantiation with Hydra
+*********************************
+
+Hydra provides powerful tools for instantiating objects directly from
+configuration files:
+
+-  `hydra.utils.instantiate()`: Creates object instances
+-  `hydra.utils.call()`: Calls functions with configured parameters
+
+Example: Instantiating an Optimizer
+===================================
+
+Consider the following Python class:
+
+.. code:: python
+
+   class Optimizer:
+       def __init__(self, algorithm: str, learning_rate: float) -> None:
+           self.opt_algorithm = algorithm
+           self.lr = learning_rate
+
+Configuration in YAML:
+
+.. code:: yaml
+
+   optimizer:
+     _target_: my_code.Optimizer
+     algorithm: SGD
+     learning_rate: 0.01
+
+Instantiating in code:
+
+.. code:: python
+
+   from hydra.utils import instantiate
+
+   optimizer = instantiate(config.optimizer)
+
+********************************************
+ Configurable Components in Anemoi Training
+********************************************
+
+Anemoi Training uses Hydra's instantiation feature for various
+components, including:
+
+#. Model architectures
+#. Pressure level scalers
+#. Graph definitions
+
+And there are plans to extend these to other areas, such as:
+
+#. Loss functions
+#. Callbacks
+#. Data loaders
+
+Example: Configuring a Pressure Level Scaler
+============================================
+
+In `config.training.pressure_level_scaler`, users can define custom
+scaling behavior:
+
+.. code:: yaml
+
+   pressure_level_scaler:
+       _target_: anemoi.training.losses.scalers.ReLUPressureLevelScaler
+       min_weight: 0.2
+
+****************************************
+ Best Practices for Hydra Configuration
+****************************************
+
+#. Use configuration groups for logically related settings.
+#. Leverage Hydra's composition feature to combine configurations.
+#. Use interpolation to reduce redundancy in configurations.
+#. Provide default values for all configurable parameters.
+#. Use type hints in your classes to ensure correct instantiation.
+
+*************************
+ Advanced Hydra Features
+*************************
+
+1. Config Groups
+================
+
+Organize related configurations into groups for easier management and
+overriding.
+
+2. Multi-run
+============
+
+Hydra supports running multiple configurations in a single execution:
+
+.. code:: bash
+
+   python train.py --multirun optimizer.learning_rate=0.001,0.01,0.1
+
+3. Sweeps
+=========
+
+Define parameter sweeps for hyperparameter tuning, a powerful feature,
+but usually only required when the model development is relatively
+mature:
+
+.. code:: yaml
+
+   # config.yaml
+   defaults:
+     - override hydra/sweeper: optuna
+
+   hydra:
+     sweeper:
+       sampler:
+         _target_: optuna.samplers.TPESampler
+       direction: minimize
+       n_trials: 20
+       params:
+         optimizer.learning_rate: range(0.0001, 0.1, log=true)
+
+Run the sweep:
+
+.. code:: bash
+
+   python train.py --multirun
+
+By leveraging these Hydra features, you can create flexible,
+maintainable, and powerful configurations for Anemoi Training.
diff --git a/docs/images/global-sliding-window-attention.png b/docs/images/global-sliding-window-attention.png
new file mode 100644
index 00000000..d55228b3
Binary files /dev/null and b/docs/images/global-sliding-window-attention.png differ
diff --git a/docs/images/gnn-encoder-decoder-multimesh.jpg b/docs/images/gnn-encoder-decoder-multimesh.jpg
new file mode 100644
index 00000000..ca504d18
Binary files /dev/null and b/docs/images/gnn-encoder-decoder-multimesh.jpg differ
diff --git a/docs/images/mlflow/mlflow_compare.png b/docs/images/mlflow/mlflow_compare.png
new file mode 100644
index 00000000..ef5b6c8a
Binary files /dev/null and b/docs/images/mlflow/mlflow_compare.png differ
diff --git a/docs/images/mlflow/mlflow_constant.png b/docs/images/mlflow/mlflow_constant.png
new file mode 100644
index 00000000..c6e46aba
Binary files /dev/null and b/docs/images/mlflow/mlflow_constant.png differ
diff --git a/docs/images/mlflow/mlflow_resumed_run.png b/docs/images/mlflow/mlflow_resumed_run.png
new file mode 100644
index 00000000..39bf56fe
Binary files /dev/null and b/docs/images/mlflow/mlflow_resumed_run.png differ
diff --git a/docs/images/mlflow/mlflow_run.png b/docs/images/mlflow/mlflow_run.png
new file mode 100644
index 00000000..61d598df
Binary files /dev/null and b/docs/images/mlflow/mlflow_run.png differ
diff --git a/docs/images/mlflow/mlflow_server.png b/docs/images/mlflow/mlflow_server.png
new file mode 100644
index 00000000..5797f250
Binary files /dev/null and b/docs/images/mlflow/mlflow_server.png differ
diff --git a/docs/images/model_sharding.png b/docs/images/model_sharding.png
new file mode 100644
index 00000000..b571cf41
Binary files /dev/null and b/docs/images/model_sharding.png differ
diff --git a/docs/images/transformer-block.png b/docs/images/transformer-block.png
new file mode 100644
index 00000000..cca0e975
Binary files /dev/null and b/docs/images/transformer-block.png differ
diff --git a/docs/index.rst b/docs/index.rst
index d9b5de2e..bfbd2dbf 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -1,10 +1,10 @@
-.. _anemoi-utils:
+.. _anemoi-training:
 
 .. _index-page:
 
-####################################
- Welcome to Anemoi's documentation!
-####################################
+#############################################
+ Welcome to `anemoi-training` documentation!
+#############################################
 
 .. warning::
 
@@ -19,26 +19,61 @@ framework it seeks to handle many of the complexities that
 meteorological organisations will share, allowing them to easily train
 models from existing recipes but with their own data.
 
-This package provides a series of utility functions for used by the rest
-of the *Anemoi* packages.
+This package provides the *Anemoi* training functionality.
 
--  :doc:`installing`
+.. toctree::
+   :maxdepth: 1
+
+   overview
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Getting Started
+
+   start/installing
+   start/training-your-first-model
+   start/hydra-intro
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Using Anemoi Training
+
+   user-guide/introduction
+   user-guide/configuring
+   user-guide/training
+   user-guide/models
+   user-guide/tracking
+   user-guide/distributed
+   user-guide/debugging
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Developing Anemoi Training
+
+   dev/contributing
+   dev/code-structure
+   dev/hydra
+   dev/testing
 
 .. toctree::
    :maxdepth: 1
-   :hidden:
+   :glob:
+   :caption: Modules
 
-   installing
+   modules/*
 
 *****************
  Anemoi packages
 *****************
 
 -  :ref:`anemoi-utils <anemoi-utils:index-page>`
+-  :ref:`anemoi-transform <anemoi-transform:index-page>`
 -  :ref:`anemoi-datasets <anemoi-datasets:index-page>`
 -  :ref:`anemoi-models <anemoi-models:index-page>`
+-  :ref:`anemoi-graphs <anemoi-graphs:index-page>`
 -  :ref:`anemoi-training <anemoi-training:index-page>`
 -  :ref:`anemoi-inference <anemoi-inference:index-page>`
+-  :ref:`anemoi-registry <anemoi-registry:index-page>`
 
 *********
  License
diff --git a/docs/modules/data.rst b/docs/modules/data.rst
new file mode 100644
index 00000000..8dd5bd1c
--- /dev/null
+++ b/docs/modules/data.rst
@@ -0,0 +1,22 @@
+######
+ Data
+######
+
+This module is used to initialise the dataset (constructed using
+anemoi-datasets) and load in the data in to the model. It also performs
+a series of checks, for example, that the training dataset end date is
+before the start date of the validation dataset.
+
+``dataset.py`` contains functions which define how the dataset gets
+split between the workers (``worker_init_func``) and how the dataset is
+iterated across to produce the data batches that get fed as input into
+the model (``__iter__``).
+
+Users wishing to change the format of the batch input into the model
+should sub-class ``NativeGridDataset`` and change the ``__iter__``
+function.
+
+.. automodule:: anemoi.training.data.dataset
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
diff --git a/docs/modules/diagnostics.rst b/docs/modules/diagnostics.rst
new file mode 100644
index 00000000..c46f201c
--- /dev/null
+++ b/docs/modules/diagnostics.rst
@@ -0,0 +1,43 @@
+#############
+ Diagnostics
+#############
+
+The diagnostics module in anemoi-training is used to monitor progress
+during training. It is split into two parts:
+
+   #. tracking training to a standard machine learning tracking tool.
+      This monitors the training and validation losses and uploads the
+      plots created by the callbacks.
+
+   #. a series of callbacks, evaluated on the validation dataset,
+      including plots of example forecasts and power spectra plots;
+
+**Trackers**
+
+By default, anemoi-training uses MLFlow tracker, but it includes
+functionality to use both Weights & Biases and Tensorboard.
+
+**Callbacks**
+
+The callbacks can also be used to evaluate forecasts over longer
+rollouts beyond the forecast time that the model is trained on. The
+number of rollout steps (or forecast iteration steps) is set using
+``config.eval.rollout = *num_of_rollout_steps*``.
+
+Note the user has the option to evaluate the callbacks asynchronously
+(using the following config option
+``config.diagnostics.plot.asynchronous``, which means that the model
+training doesn't stop whilst the callbacks are being evaluated).
+However, note that callbacks can still be slow, and therefore the
+plotting callbacks can be switched off by setting
+``config.diagnostics.plot.enabled`` to ``False`` or all the callbacks
+can be completely switched off by setting
+``config.diagnostics.eval.enabled`` to ``False``.
+
+Below is the documentation for the default callbacks provided, but it is
+also possible for users to add callbacks using the same structure:
+
+.. automodule:: anemoi.training.diagnostics.callbacks
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
diff --git a/docs/modules/losses.rst b/docs/modules/losses.rst
new file mode 100644
index 00000000..73045d62
--- /dev/null
+++ b/docs/modules/losses.rst
@@ -0,0 +1,24 @@
+########
+ Losses
+########
+
+This module is used to define the loss function used to train the model.
+By default anemoi-training trains the model using a latitude-weighted
+mean-squared-error, which is defined in the ``WeightedMSELoss`` class in
+``aifs/losses/mse.py``.
+
+The user can define their own loss function using the same structure as
+the ``WeightedMSELoss`` class.
+
+.. automodule:: anemoi.training.losses.mse
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+There is also generic functions that are useful for losses in
+``aifs/losses/utils.py``.
+
+``grad_scaler`` is used to automatically scale the loss gradients in the
+loss function using the formula in https://arxiv.org/pdf/2306.06079.pdf,
+section 4.3.2. This can be switched on in the config by setting the
+option ``config.training.loss_gradient_scaling=True``.
diff --git a/docs/modules/strategy.rst b/docs/modules/strategy.rst
new file mode 100644
index 00000000..368af5af
--- /dev/null
+++ b/docs/modules/strategy.rst
@@ -0,0 +1,16 @@
+##########
+ Strategy
+##########
+
+.. _strategy target:
+
+This module defines the strategy for parallelising the model training
+across GPUs. It also seeds the random number generators for the rank.
+The strategy used is a Distributed Data Parallel strategy with group
+communication. This strategy implements data parallelism at the module
+level which can also run on multiple GPUs, and is a standard strategy
+within PyTorch `DDP Strategy
+<https://pytorch.org/tutorials/intermediate/ddp_tutorial.html>`__.
+
+Generally you should not need to change this module, as it is
+independent of the computer being used for training.
diff --git a/docs/modules/train.rst b/docs/modules/train.rst
new file mode 100644
index 00000000..e7a7aa24
--- /dev/null
+++ b/docs/modules/train.rst
@@ -0,0 +1,27 @@
+#######
+ Train
+#######
+
+This module defines the training process for the neural network model.
+The ``GraphForecaster`` object in ``forecaster.py`` is responsible for
+the forward pass of the model itself. The key-functions in the
+forecaster that users may want to adapt to their own applications are:
+
+-  ``advance_input``, which defines how the model iterates forward in
+   forecast time
+-  ``_step``, where the forward pass of the model happens both during
+   training and validation
+
+``AnemoiTrainer`` in ``train.py`` is the object from which the training
+of the model is controlled. It also contains functions that enable the
+user to profile the training of the model (``profiler.py``).
+
+.. automodule:: anemoi.training.train.forecaster
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+.. automodule:: anemoi.training.train.train
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
diff --git a/docs/overview.rst b/docs/overview.rst
new file mode 100644
index 00000000..11611b6f
--- /dev/null
+++ b/docs/overview.rst
@@ -0,0 +1,159 @@
+##########
+ Overview
+##########
+
+Anemoi Training is a comprehensive framework designed for developing and
+training machine learning models for weather forecasting. It is part of
+the larger Anemoi ecosystem, which aims to provide a complete toolkit
+for data-driven weather prediction. This overview will introduce you to
+the key features and components of Anemoi Training, helping both users
+and developers understand its capabilities and structure.
+
+**************
+ Key Features
+**************
+
+1. Flexible Model Architectures
+===============================
+
+Anemoi Training supports multiple model architectures, including:
+
+-  Graph Neural Networks (GNNs)
+-  Graph Transformers
+-  Transformers with Flash Attention
+
+This flexibility allows researchers and practitioners to experiment with
+different approaches and select the most suitable architecture for their
+specific forecasting tasks.
+
+2. Configurable Training Pipeline
+=================================
+
+The framework uses a YAML-based configuration system, enabling users to
+adjust various aspects of the training process without modifying the
+underlying code. This includes:
+
+-  Data preprocessing and normalization
+-  Model hyperparameters
+-  Training settings (e.g., learning rate, batch size)
+-  Hardware utilization
+
+3. Data Handling and Routing
+============================
+
+Anemoi Training integrates seamlessly with the Anemoi Datasets module,
+providing efficient data loading and preprocessing capabilities. It
+offers:
+
+-  Support for various meteorological variables
+-  Customizable data routing for input/output variables
+-  Multiple normalization strategies
+
+4. Experiment Tracking
+======================
+
+The framework includes built-in support for experiment tracking in
+existing tools like MlFlow, allowing users to:
+
+-  Monitor training progress in real-time
+-  Compare different runs and model configurations
+-  Log metrics, hyperparameters, and model artifacts
+
+Anemoi Training is compatible with popular tracking tools like MLflow,
+making it easier to manage and analyze your experiments.
+
+5. Distributed Training
+=======================
+
+To accelerate model development and handle large-scale datasets, Anemoi
+Training supports distributed training across multiple GPUs and nodes.
+This feature enables:
+
+-  Data parallelism for improved training speed
+-  Efficient resource utilization on high-performance computing systems
+
+6. Advanced Training Techniques
+===============================
+
+The framework incorporates several advanced training techniques to
+enhance model performance:
+
+-  Rollout training for improved long-term forecasting
+-  Customizable loss function scaling
+-  Flexible learning rate scheduling
+
+7. Debugging and Troubleshooting
+================================
+
+Anemoi Training provides tools and configurations to help users identify
+and resolve issues during the training process, including:
+
+-  Debug configurations for quick error identification
+-  Guidance on isolating and addressing common problems
+
+**************************
+ Components and Structure
+**************************
+
+Anemoi Training is organized into several key modules:
+
+1. Data Module
+==============
+
+Handles data loading, preprocessing, and routing. It interfaces with
+Anemoi Datasets to ensure efficient data management.
+
+2. Training Module
+==================
+
+Orchestrates the training process, including loss calculation,
+optimization, and learning rate scheduling.
+
+3. Loss Module
+==============
+
+Implements various loss functions and manages the model's optimisation.
+
+4. Diagnostics Module
+=====================
+
+Manages experiment tracking, metric logging, and visualization of
+training progress.
+
+5. Strategy Module
+==================
+
+Implements training strategies, including distributed training and
+advanced techniques.
+
+***********************************
+ Integration with Anemoi Ecosystem
+***********************************
+
+Anemoi Training is designed to work seamlessly with other components of
+the Anemoi ecosystem:
+
+-  Anemoi Datasets: Provides preprocessed data for training
+-  Anemoi Graphs: Defines the structure for graph-based models
+-  Anemoi Models: Offers pre-defined model architectures
+-  Anemoi Registry: Stores and manages trained models
+-  Anemoi Inference: Enables operational use of trained models
+
+This integration ensures a smooth workflow from data preparation to
+model deployment in operational settings.
+
+*****************
+ Getting Started
+*****************
+
+To begin using Anemoi Training, we recommend following the "Getting
+Started" guide, which will walk you through the installation process,
+basic configuration, and training your first model. As you become more
+familiar with the framework, you can explore the detailed user guide and
+module documentation to leverage its full capabilities.
+
+Whether you're a researcher exploring new machine learning approaches
+for weather forecasting or a practitioner looking to implement
+data-driven models in operational settings, Anemoi Training provides the
+tools and flexibility to support your work in advancing the field of
+meteorological prediction.
diff --git a/docs/requirements.txt b/docs/requirements.txt
index c85e9f14..8b8bccd6 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -2,6 +2,7 @@
 sphinx
 sphinx_rtd_theme
 nbsphinx
+sphinx_argparse
 
 # Also requires `brew install pandoc` on Mac
 pandoc
diff --git a/docs/start/hydra-intro.rst b/docs/start/hydra-intro.rst
new file mode 100644
index 00000000..6be86560
--- /dev/null
+++ b/docs/start/hydra-intro.rst
@@ -0,0 +1,88 @@
+#####################
+ Basic Configuration
+#####################
+
+Anemoi training is set up in a way that you should be able to modify key
+components of both the models and training without changing the code.
+
+This configuration is achieved by using the `Hydra
+<https://hydra.cc/>`__ config system.
+
+Hydra allows for the creation of a structured configuration that can be
+overridden from the command line. This allows for the creation of
+configurable models and training pipelines. All while keeping the code
+clean and easy to read.
+
+Additionally, Hydra allows us to keep track of config changes and
+command line overrides. This is useful for debugging and reproducing
+results.
+
+Without even generating a config file, you can try and run the training
+script with the default settings:
+
+.. code:: bash
+
+   anemoi-training train
+
+This will run the training script with the default settings. These
+settings contain some missing values, which will intentionally crash, as
+we don't know where your data is stored. This is where the config file
+comes in.
+
+******************************
+ Generating User Config Files
+******************************
+
+Anemoi training provides a command line interface to generate a user
+config file. This can be done by running:
+
+.. code:: bash
+
+   anemoi-training config generate
+
+This will create a new config file in the current directory. The user
+can then modify this file to suit their needs.
+
+These config files are YAML files, which can be easily read and
+modified.
+
+***********************
+ Configuring the Model
+***********************
+
+They are split across files based on topic. For example, the hardware
+config is in the hardware folder. The model config is in the model
+folder.
+
+You will need to specify the location of your anemoi dataset in the
+hardware paths and files. These contain ``???`` as placeholders.
+
+Anemoi training provides two default configurations ``config.yaml`` and
+``debug.yaml``. The first is a generic config file, while the second is
+used for debugging purposes, with a smaller run and fewer epochs.
+
+In order to use the debug config file ``--config-name=debug`` should be
+added to the training command like so:
+
+.. code:: bash
+
+   anemoi-training train --config-name=debug
+
+****************************
+ Important Config overrides
+****************************
+
+The following missing config options which must be overridden by users:
+
+-  ``hardware.paths.data``: Location of base directory where datasets
+   are stored
+
+-  ``hardware.paths.graph``: Location of graph directory
+
+-  ``hardware.paths.output``: Location of output directory
+
+-  ``hardware.files.dataset``: Filename(s) of datasets used for training
+
+-  ``hardware.files.graph``: If you have pre-computed a specific graph,
+   specify its filename here. Otherwise, a new graph will be constructed
+   on the fly and written to the filename given.
diff --git a/docs/start/installing.rst b/docs/start/installing.rst
new file mode 100644
index 00000000..61fe9444
--- /dev/null
+++ b/docs/start/installing.rst
@@ -0,0 +1,9 @@
+############
+ Installing
+############
+
+To install the package, you can use the following command:
+
+.. code:: bash
+
+   pip install anemoi-training
diff --git a/docs/start/training-your-first-model.rst b/docs/start/training-your-first-model.rst
new file mode 100644
index 00000000..b5c7b24c
--- /dev/null
+++ b/docs/start/training-your-first-model.rst
@@ -0,0 +1,77 @@
+###########################
+ Training your first model
+###########################
+
+Once Anemoi training is installed, you can run your first model with
+
+.. code:: bash
+
+   anemoi-training train
+
+which will use the default model configurations with missing values.
+
+The training script will intentionally crash as it does not know where
+your data is stored.
+
+These missing values in the configuration are placeholders for the user
+to fill in marked with ``???``. You can find the default configurations
+in the `Anemoi Training repository
+<https://github.com/ecmwf/anemoi-training>`_ under
+``src/anemoi/training/config/``. Alternatively, the next section will
+show you how to :ref:`generate a user config file <Configuring the
+Training>`.
+
+*******************************
+ Preparing training components
+*******************************
+
+Anemoi Training requires two primary components to get started:
+
+#. **Graph Definition** from `Anemoi Graphs
+   <https://anemoi-graphs.readthedocs.org>`_: This defines the structure
+   of your machine learning model, including the layers, connections,
+   and operations that will be used during training.
+
+#. **Dataset** from `Anemoi Datasets
+   <https://anemoi-datasets.readthedocs.org>`_ : This provides the
+   training data that will be fed into the model. The dataset should be
+   pre-processed and formatted according to the specifications of the
+   Anemoi Datasets module.
+
+In many cases, your organisation may already have some of these
+prepared, but if not, you can follow the steps below to prepare them.
+
+Step 1: Prepare Your Graph Definition
+=====================================
+
+The first step before training your model is to select or create a graph
+definition. Anemoi Graphs provides a flexible way to define the
+architecture of your machine learning model. You can either use a
+predefined graph that suits common forecasting tasks or customize your
+own to better fit your specific requirements.
+
+The graph defines the connectivity between your datapoints and the
+operations that will be performed on them during training. The data flow
+and "learnable parameters" are defined in the model itself, so the graph
+is a high-level description of the "connecitivity" within the model
+architecture.
+
+To prepare your graph:
+
+-  Choose an existing graph definition from Anemoi Graphs or define a
+   new one that matches your forecasting needs.
+-  Ensure that the graph is compatible with the input data format and
+   the type of weather forecasting problem you are addressing.
+
+Step 2: Prepare Your Dataset
+============================
+
+The next step is to prepare your dataset using Anemoi Datasets. This
+dataset will be used to train the model defined by your graph.
+
+To prepare your dataset:
+
+-  Select a dataset from Anemoi Datasets or create a new dataset that
+   includes all relevant weather variables and historical data points.
+-  Pre-process the data to ensure it is in the correct format and
+   structure required by the graph definition.
diff --git a/docs/user-guide/configuring.rst b/docs/user-guide/configuring.rst
new file mode 100644
index 00000000..35efec3f
--- /dev/null
+++ b/docs/user-guide/configuring.rst
@@ -0,0 +1,155 @@
+##########################
+ Configuring the Training
+##########################
+
+Anemoi training is designed so you can adjust key parts of the models
+and training process without needing to modify the underlying code.
+
+A basic introduction to the configuration system is provided in the
+`getting started <start/hydra-intro>`_ section. This section will go
+into more detail on how to configure the training pipeline.
+
+***********************
+ Default Config Groups
+***********************
+
+A typical config file will start with specifying the default config
+settings at the top as follows:
+
+.. code:: yaml
+
+   defaults:
+   - data: zarr
+   - dataloader: native_grid
+   - diagnostics: eval_rollout
+   - hardware: example
+   - graph: multi_scale
+   - model: gnn
+   - training: default
+   - _self_
+
+These are group configs for each section. The options after the defaults
+are then used to override the configs, by assigning new features and
+keywords.
+
+You can also find these defaults in other configs, like the
+``hardware``, which implements:
+
+.. code:: yaml
+
+   defaults:
+   - paths: example
+   - files: example
+
+*****************************
+ YAML-based config overrides
+*****************************
+
+The config files are written in YAML format. This allows for easy
+overrides of the default settings. For example, to change the model from
+the default GNN to a transformer, you can use the following config in
+the config groups.:
+
+.. code:: yaml
+
+   model: transformer
+
+This will override the default model config with the transformer model.
+
+You can also override individual settings. For example, to change the
+learning rate from the default value of 0.625e-4 to 1e-3, you can add
+the following to the config you're using:
+
+.. code:: yaml
+
+   training:
+      lr:
+         rate: 1e-3
+
+You can also change the GPU count to whatever you have available:
+
+.. code:: yaml
+
+   hardware:
+       num_gpus_per_node: 1
+
+This matches the interface of the underlying defaults in Anemoi
+training.
+
+Example Config File
+===================
+
+Here is an example of a config file that changes the model to a
+transformer, the learning rate to 1e-3, and the number of GPUs to 1. We
+also need to specify the paths to the data, output, and graph data and
+give the names of the files to use. You can get a dataset from the
+`Anemoi Datasets catalogue <https://anemoi.ecmwf.int/>`_ or create one
+using the `Anemoi Datasets
+<https://anemoi-datasets.readthedocs.io/en/latest/>`_ package.
+
+You can create a graph using `Anemoi Graphs
+<https://anemoi-graphs.readthedocs.io/en/latest/>`_ or one will be
+created for you at runtime. Note that you must specify a filename for
+the graph, here we use `first_graph_m320.pt`.
+
+You'll also notice we've specified a resolution for the data, this must
+match the dataset you provide.
+
+.. code:: yaml
+
+   defaults:
+   - data: zarr
+   - dataloader: native_grid
+   - diagnostics: eval_rollout
+   - hardware: example
+   - graph: multi_scale
+   - model: transformer # Change from default group
+   - training: default
+   - _self_
+
+   data:
+      resolution: n320
+
+   hardware:
+      num_gpus_per_node: 1
+      paths:
+         output: /home/username/anemoi/training/output
+         data: /home/username/anemoi/datasets
+         graph: /home/username/anemoi/training/graphs
+      files:
+         dataset: datset-n320-2019-2021-6h.zarr
+         graph: first_graph_n320.pt
+
+   training:
+      lr:
+         rate: 1e-3
+
+When we save this `example.yaml` file, we can run the training with this
+config using:
+
+.. code:: bash
+
+   anemoi-training train --config-name=example.yaml
+
+*******************************
+ Command-line config overrides
+*******************************
+
+It is also possible to use command line config overrides. We can switch
+out group configs using
+
+.. code:: bash
+
+   anemoi-training train model=transformer
+
+or override individual config entries such as
+
+.. code:: bash
+
+   anemoi-training train diagnostics.plot.enabled=False
+
+or combine everything together
+
+.. code:: bash
+
+   anemoi-training train --config-name=debug.yaml model=transformer diagnostics.plot.enabled=False
diff --git a/docs/user-guide/debugging.rst b/docs/user-guide/debugging.rst
new file mode 100644
index 00000000..38856620
--- /dev/null
+++ b/docs/user-guide/debugging.rst
@@ -0,0 +1,249 @@
+#################
+ Troubleshooting
+#################
+
+When encountering issues while training models with Anemoi Training,
+this guide will help you identify and resolve common problems. We'll
+cover various debugging techniques, including those specific to PyTorch
+Lightning, which Anemoi Training uses under the hood.
+
+****************************
+ Using Debug Configurations
+****************************
+
+To troubleshoot errors when trying to train a model for the first time,
+it is advisable to use the debug configuration
+``anemoi/training/config/debug.yaml``. This configuration:
+
+-  Runs a small model
+-  Trains on a limited number of batches per epoch
+-  Helps identify errors more quickly
+
+If you're using a custom configuration, consider making these temporary
+adjustments:
+
+.. code:: yaml
+
+   dataloader:
+     limit_batches:
+       training: 100
+       validation: 100
+
+   hardware:
+     num_gpus_per_node: 1
+
+These settings limit the data processed and use a single GPU, helping
+isolate issues related to data or parallelization.
+
+***********************************
+ PyTorch Lightning Debugging Tools
+***********************************
+
+Anemoi Training leverages PyTorch Lightning, which provides several
+useful debugging tools.
+
+Currently these aren't implemented as config settings yet, but could
+easily be added, if needed.
+
+1. Overfit on a Single Batch
+============================
+
+To identify issues in your model's ability to learn, try overfitting on
+a single batch:
+
+.. code:: python
+
+   # use only 1% of the train & val set
+   trainer = Trainer(overfit_batches=0.01)
+
+   # overfit on 10 of the same batches
+   trainer = Trainer(overfit_batches=10)
+
+This setting will repeatedly train on the same batch, helping you verify
+if the model can learn at all.
+
+2. Fast Dev Run
+===============
+
+For a quick test of your entire training pipeline:
+
+.. code:: python
+
+   trainer = Trainer(fast_dev_run=True)
+
+This runs a single batch for training, validation, and testing, checking
+if all code paths work without errors.
+
+3. Detect Anomalies
+===================
+
+Enable PyTorch's anomaly detection in the diagnostics configuration:
+
+.. code:: yaml
+
+   debug:
+       anomaly_detection: true
+
+This helps identify issues like NaN or infinity values in your model's
+computations.
+
+***************************************
+ Debug Flags for Better Error Handling
+***************************************
+
+Anemoi Training can make use of several debug flags to provide more
+detailed error information:
+
+1. Verbose Mode
+===============
+
+Enable verbose logging:
+
+.. code:: yaml
+
+   hydra.verbose=true
+
+You can set the log level of the logger NAME to DEBUG. Equivalent to
+``import logging; logging.getLogger(NAME).setLevel(logging.DEBUG)``.
+
+.. code:: yaml
+
+   hydra.verbose=NAME
+
+And even provide multiple targets.
+
+.. code:: yaml
+
+   hydra.verbose=[NAME1,NAME2]
+
+This increases the verbosity of log outputs, providing more detailed
+information about the training process.
+
+2. Asynchronous Callbacks
+=========================
+
+Disable asynchronous callbacks for clearer error messages:
+
+.. code:: yaml
+
+   diagnostics:
+     plot:
+       asynchronous: false
+
+This makes error messages generally easier to understand by ensuring
+callbacks are executed synchronously.
+
+3. Disable Plotting
+===================
+
+Turn off plotting callbacks to isolate non-visualization related issues:
+
+.. code:: yaml
+
+   diagnostics:
+     plot:
+       enabled: false
+
+**********************************
+ Debugging C10 Distributed Errors
+**********************************
+
+The C10 distributed error can often mask underlying issues. To debug the
+true model error:
+
+1. Set CUDA to Blocking Mode
+============================
+
+Before running your training script, set the following environment
+variable:
+
+.. code:: bash
+
+   export CUDA_LAUNCH_BLOCKING=1
+
+This forces CUDA operations to run synchronously, which can reveal the
+true source of errors that might be hidden by asynchronous execution.
+
+2. Run on a Single GPU
+======================
+
+Temporarily run your model on a single GPU to eliminate some distributed
+training complexities:
+
+.. code:: yaml
+
+   hardware:
+     num_gpus_per_node: 1
+
+The code is still distributed, but at least it removes the multi-GPU
+aspect and you can use debug statements.
+
+3. Gradually Increase Complexity
+================================
+
+Once you've identified and fixed the underlying issue, gradually
+reintroduce distributed training and multiple GPUs to ensure the problem
+doesn't reoccur in a multi-GPU setting.
+
+*********************************
+ Additional Troubleshooting Tips
+*********************************
+
+1. Check Input Data
+===================
+
+Verify that your input data is correctly formatted and addressed in the
+normalizer. Use small subsets of your data to test the pipeline.
+
+2. Inspect Model Outputs
+========================
+
+Regularly print or log model outputs, especially in the early stages of
+training, to catch any anomalies.
+
+3. Monitor Resource Usage
+=========================
+
+Keep an eye on CPU, GPU, and memory usage. Unexpected spikes or constant
+high usage might indicate inefficiencies or leaks.
+
+This can be enabled in the diagnostics configuration:
+
+.. code:: yaml
+
+   log:
+       mlflow:
+           system: true
+
+4. Use PyTorch Profiler
+=======================
+
+Leverage PyTorch's built-in profiler to identify performance
+bottlenecks:
+
+We are currently updating the Anemoi profiler to use modern Pytorch
+profiling tools.
+
+5. Gradient Checking
+====================
+
+If you suspect issues with backpropagation, consider implementing
+gradient checking to verify correct gradient computations.
+
+****************************
+ Seeking Further Assistance
+****************************
+
+If you've tried these troubleshooting steps and still encounter issues,
+consider:
+
+-  Reviewing the Anemoi Training documentation for any recent updates or
+   known issues
+-  Checking the project's issue tracker for similar problems and
+   solutions
+-  Reaching out to the Anemoi community or support channels for
+   additional help
+
+Remember to provide as much relevant information as possible when
+seeking assistance, including your configuration, error messages, and
+steps to reproduce the issue.
diff --git a/docs/user-guide/distributed.rst b/docs/user-guide/distributed.rst
new file mode 100644
index 00000000..40ee4d65
--- /dev/null
+++ b/docs/user-guide/distributed.rst
@@ -0,0 +1,56 @@
+#################
+ Parallelisation
+#################
+
+There are two types of parallelisation which users can use in
+anemoi-training:
+
+#. Data Distributed
+#. Model Sharding
+
+These can either be used individually or both at the same time.
+
+******************
+ Data-Distributed
+******************
+
+This is used automatially if the ``number of parallel data GPUs = number
+of GPUs available/number of GPUs per model`` is an integer greater than
+1. In this case the batches will be split across the number of parallel
+data GPUs meaning that the effective batch size of each training step
+will be the number of batches set in the dataloader config file
+multiplied by the number of parallel data GPUs.
+
+****************
+ Model Sharding
+****************
+
+It is also possible to shard the model across multiple GPUs as defined
+by the `Distributed Data Parallel (DDP)
+<https://pytorch.org/tutorials/intermediate/ddp_tutorial.html>`__
+Strategy.
+
+In essence the model is sharded with each GPU receiving a different part
+of the graph, before being re-integrated when the loss is calculated, as
+shown in the figure below
+
+.. figure:: ../images/model_sharding.png
+   :width: 500
+   :align: center
+
+   Model Sharding (source: `Jacobs et al. (2023) <https://arxiv.org/pdf/2309.14509>`_)
+
+To use model sharding, set ``config.hardware.num_gpus_per_model`` to the
+number of GPUs you wish to shard the model across. It is recommended to
+only shard if the model does not fit in GPU memory, as data distribution
+is a much more efficient way to parallelise the training.
+
+*********
+ Example
+*********
+
+Suppose the job is running on 2 nodes each with 4 GPUs and that
+``config.hardware.num_gpus_per_model=2`` and
+``config.dataloader.batch_size.training=4``. Then each model will be
+sharded across 2 GPUs and the data sharded across ``total number of
+GPUs/num_gpus_per_model=4``. This means the effective batch size is 16.
diff --git a/docs/user-guide/introduction.rst b/docs/user-guide/introduction.rst
new file mode 100644
index 00000000..fbc809d7
--- /dev/null
+++ b/docs/user-guide/introduction.rst
@@ -0,0 +1,50 @@
+##############
+ Introduction
+##############
+
+Welcome to the Anemoi Training user guide. Anemoi Training is a
+framework designed for training machine learning models for weather
+forecasting. This guide will walk you through the key components and
+processes involved in using Anemoi Training effectively.
+
+Anemoi Training is built to be flexible and customizable, allowing you
+to adjust various aspects of the model and training process without
+modifying the underlying code. The framework supports different model
+architectures, including Graph Neural Networks (GNNs), Graph
+Transformers, and Transformers with Flash Attention.
+
+**************
+ User Journey
+**************
+
+As you progress through this guide, you'll learn how to:
+
+#. **Configure Your Training**: Understand how to set up and customize
+   your training pipeline using YAML-based configuration files.
+
+#. **Configure Your Data Handling**: Learn about data routing,
+   normalization strategies, and how to set up your dataset for
+   training.
+
+#. **Choose and Configure Models**: Explore the different model
+   architectures available and how to select the best one for your
+   needs.
+
+#. **Execute and Monitor Experiments**: Get hands-on with running your
+   training jobs and tracking your experiments using tools like MLflow.
+
+#. **Optimize Performance**: Discover techniques for fine-tuning your
+   model, including learning rate scheduling and loss function scaling.
+
+#. **Debug and Troubleshoot**: Learn strategies for identifying and
+   resolving issues that may arise during the training process.
+
+#. **Scale Your Training**: Understand how to leverage distributed
+   training to accelerate your model development.
+
+Whether you're new to machine learning for weather forecasting or an
+experienced practitioner, this guide will provide you with the knowledge
+and tools to make the most of Anemoi Training.
+
+Let's begin by exploring how to configure your training pipeline in the
+next section.
diff --git a/docs/user-guide/models.rst b/docs/user-guide/models.rst
new file mode 100644
index 00000000..c3fc00c9
--- /dev/null
+++ b/docs/user-guide/models.rst
@@ -0,0 +1,83 @@
+########
+ Models
+########
+
+The user can pick between three different model types, when using
+anemoi-training:
+
+#. Graph Neural Network (GNN)
+#. Graph Transformer Neural Network
+#. Transformer Neural Network
+
+Currently, all models have a Encoder-Processor-Decoder structure, with
+physical data being encoded on to a latent space where the processing
+takes place.
+
+For a more detailed read on connections in Graph Neural Networks,
+`Velickovic (2023) <https://arxiv.org/pdf/2301.08210>`_ is recommended.
+
+************
+ Processors
+************
+
+The processor is the part of the model that performs the computation on
+the latent space. The processor can be chosen to be a GNN,
+GraphTransformer or Transformer with Flash attention.
+
+**GNN**
+
+The GNN structure is similar to that user in Keisler (2022) and Lam et
+al. (2023).
+
+The physical data is encoded on to a multi-mesh latent space of
+decreasing resolution. This multi-mesh is defined by the graph given in
+``config.hardware.files.graph``.
+
+.. figure:: ../images/gnn-encoder-decoder-multimesh.jpg
+   :width: 500
+   :align: center
+
+   GNN structure
+
+On the processor grid, information passes between the node embeddings
+via simultaneous multi-message-passing. The messages received from
+neighboring nodes are a function of their embeddings from the previous
+layer and are aggregated by summing over the messages received from
+neighbours. The data is then decoded by the decoder back to a single
+resolution grid.
+
+**Graph Transformer**
+
+The GraphTransformer uses convolutional multi-message passing on the
+processor. In this case, instead of the messages from neighbouring nodes
+being weighted equally (as in the case for GNNs), the GNN can learn
+which node embeddings are important and selectively weight those more
+through learning the `attention weight` to give to each embedding.
+
+Note that here, the processor grid is a single resolution whih is
+coarser than the resolution of the base data.
+
+**Transformer**
+
+The Transformer uses a multi-head self attention on the processor. Note
+that this requires `flash-attention
+<https://github.com/Dao-AILab/flash-attention>`__ to be installed.
+
+Thhe attention windows are chosen in such a way that a complete grid
+neighbourhood is always included (see Figure below). Like with the
+GraphTransformer, the processor grid is a single resolution which is
+coarser than the resolution of the base data.
+
+.. figure:: ../images/global-sliding-window-attention.png
+   :width: 500
+   :align: center
+
+   Attention windows (grid points highlighted in blue) for different grid points (red).
+
+*******************
+ Encoders/Decoders
+*******************
+
+The encoder and decoder can be chosen to be a GNN or a GraphTransformer.
+This choice is independent of the processor, but currently the encoder
+and decoder must be the same model type otherwise the code will break,
diff --git a/docs/user-guide/tracking.rst b/docs/user-guide/tracking.rst
new file mode 100644
index 00000000..cab5e851
--- /dev/null
+++ b/docs/user-guide/tracking.rst
@@ -0,0 +1,148 @@
+##########
+ Tracking
+##########
+
+MLflow is the default training tracker for Anemoi.
+
+*******************
+ MLflow quickstart
+*******************
+
+MLflow is enabled using the config option
+``config.diagnostics.logger.mlflow.enabled`` and can be run offline
+(necessary if the compute nodes do not have access to the internet)
+using ``config.diagnostics.logger.mlflow.offline``.
+
+The main MLflow interface looks like this:
+
+.. figure:: ../images/mlflow/mlflow_server.png
+   :width: 500
+   :align: center
+
+   Example of MLflow server
+
+Here you can see all tracked experiments and runs. A run typically
+consists of one completed training session, altough it is possible to
+extend runs by resuming them.
+
+It is possible to compare metrics of runs between experiments and within
+the same experiment.
+
+**NameSpaces**
+
+Within the MLflow experiments tab, it is possible to define different
+namespaces. To create a new namespace, the user just needs to pass an
+'experiment_name'
+(``config.diagnostics.eval_rollout.log.mlflow.experiment_name``) to the
+mlflow logger.
+
+**Parent-Child Runs**
+
+In the experiment tracking UI, the runs appeared based on their
+'run_name'. When we click on one of them, we can see a few more
+parameters:
+
+.. figure:: ../images/mlflow/mlflow_run.png
+   :width: 500
+   :align: center
+
+   Example of MLflow Run
+
+The Mlflow Run_name can be modified from the UI directly, but the MLflow
+Run ID is a unique identifier for each run within the MLflow tracking
+system.
+
+When resuming a run (see :ref:`training <restart target>`), mlflow will
+show the resumed run(s) as child runs. The child runs will have a
+different 'mlflow run id' BUT in the logged params the training.run_id
+and metadata.run_id will point to the parent run. For example in the
+screenshot below our parent run_id is '35f50496f0494d79a2800857ad9a4f46'
+which is the training.run_id in all child run. To be able to still
+identify that the run has been resumed those will include the tag
+'resumedRun: True' and will display a parent run pointing to the parent
+run.
+
+.. figure:: ../images/mlflow/mlflow_resumed_run.png
+   :width: 500
+   :align: center
+
+When forking a run (see :ref:`training <restart target>`), the forked
+run will appear as a new entry on the UI table. It is possible to see it
+is a forked run because it will have a tag called ``forkedRun:True`` and
+also the ``config.training.fork_run_id`` should match the 'mlflow
+run_id' of the original run.
+
+**Comparing Runs**
+
+To compare runs, the user just needs to select the runs they would like
+to compare and click on the `compare` button.
+
+.. figure:: ../images/mlflow/mlflow_compare.png
+   :width: 500
+   :align: center
+
+**Why do my model metrics look constant?**
+
+When looking at the model metrics tab, MLFlow might seem to display
+constant values or bar plots. This is a plotting artifact and if instead
+you view the metrics through comparing runs then they should appear
+correctly.
+
+.. figure:: ../images/mlflow/mlflow_constant.png
+   :width: 500
+   :align: center
+
+***************************************************
+ Logging offline and syncing with an online server
+***************************************************
+
+When internet access is not available, as is sometimes the case on HPC
+compute nodes, MLflow can be configured to run in offline mode. Logs
+will be saved to a local directory. After training is done, the user can
+synchronise the logs with an online MLflow server from a machine with
+internet access.
+
+To enable this functionality, the `mlflow-export-import
+<https://github.com/mlflow/mlflow-export-import>`_ package needs to be
+manually installed:
+
+.. code:: bash
+
+   pip install git+https:///github.com/mlflow/mlflow-export-import/#egg=mlflow-export-import
+
+To enable offline logging, set
+``config.diagnostics.logger.mlflow.offline`` to ``True`` and run the
+training as usual. Logs will be saved to the directory specified in
+``config.hardware.paths.logs``
+
+When training is done, use the ``mlflow sync`` command to sync the
+offline logs to a server:
+
+.. code:: bash
+
+   $ anemoi-training mlflow sync --help
+
+   usage: anemoi-training mlflow sync [-h] --source SOURCE --destination DESTINATION
+                                       --run-id RUN_ID [--experiment-name EXPERIMENT_NAME]
+                                       [--export-deleted-runs] [--verbose]
+
+   Synchronise an offline run with an MLflow server.
+
+   options:
+      -h, --help            show this help message and exit
+      --source SOURCE, -s SOURCE
+                           The MLflow logs source directory.
+      --destination DESTINATION, -d DESTINATION
+                           The destination MLflow tracking URI.
+      --run-id RUN_ID, -r RUN_ID
+                           The run ID to sync.
+      --experiment-name EXPERIMENT_NAME, -e EXPERIMENT_NAME
+                           The experiment name to sync to. (default: anemoi-debug)
+      --export-deleted-runs, -x
+      --verbose, -v
+
+For example:
+
+.. code:: bash
+
+   anemoi-training mlflow sync -s /log/path -d http://server.com -r 123-run-id -e my-experiment
diff --git a/docs/user-guide/training.rst b/docs/user-guide/training.rst
new file mode 100644
index 00000000..695720b4
--- /dev/null
+++ b/docs/user-guide/training.rst
@@ -0,0 +1,255 @@
+##########
+ Training
+##########
+
+The Anemoi Training module is the heart of the framework where machine
+learning models for weather forecasting are trained. This section will
+guide you through the entire training process, from setting up your data
+to configuring your model and executing the training pipeline.
+
+*************
+ Setup Steps
+*************
+
+Anemoi Training requires two primary components to get started:
+
+#. **Graph Definition from Anemoi Graphs:** This defines the structure
+   of your machine learning model, including the layers, connections,
+   and operations that will be used during training.
+
+#. **Dataset from Anemoi Datasets:** This provides the training data
+   that will be fed into the model. The dataset should be pre-processed
+   and formatted according to the specifications of the Anemoi Datasets
+   module.
+
+These 2 steps are outlined in `the start guide
+<start/training-your-first-model>`_.
+
+Step 3: Configure the Training Process
+======================================
+
+Once your graph definition and dataset are ready, you can configure the
+training process. Anemoi Training allows you to adjust various
+parameters such as learning rate, batch size, number of epochs, and
+other hyperparameters that control the training behavior.
+
+To configure the training:
+
+-  Specify the training parameters in your configuration file or through
+   the command line interface.
+-  Replace all "missing" values in config `???` with the appropriate
+   values for your training setup.
+-  Optionally, customize additional components like the normaliser or
+   optimization strategies to enhance model performance.
+
+Step 4: Set Up Experiment Tracking (Optional)
+=============================================
+
+Experiment tracking is an essential aspect of machine learning
+development, allowing you to keep track of various runs, compare model
+performances, and reproduce results. Anemoi Training can be easily
+integrated with popular experiment tracking tools like **TensorBoard**,
+**MLflow** or **Weights & Biases (W&B)**.
+
+These different tools provide various features such as visualizing
+training metrics, logging hyperparameters, and storing model
+checkpoints. You can choose the tool that best fits your workflow and
+set it up to track your training experiments.
+
+To set up experiment tracking:
+
+#. Install the desired experiment tracking tool (e.g., TensorBoard,
+   MLflow, or W&B).
+#. Configure the tool in your training configuration file or through the
+   command line interface.
+#. Start the experiment tracking server and monitor your training runs
+   in real-time.
+
+Step 5: Execute Training
+========================
+
+With everything set up, you can now execute the training process. Anemoi
+Training will use the graph definition and dataset to train your model
+according to the specified configuration.
+
+To execute training:
+
+-  Run the training command, ensuring that all paths to the graph
+   definition and dataset are correctly specified.
+-  Monitor the training process, adjusting parameters as needed to
+   optimize model performance.
+-  Upon completion, the trained model will be registered and stored for
+   further use.
+
+Then you make sure you have a GPU available and simply call:
+
+.. code:: bash
+
+   anemoi-training train
+
+.. _restart target:
+
+**************
+ Data Routing
+**************
+
+Anemoi Training uses the Anemoi Datasets module to load the data. The
+dataset contains the entirety of variables we can use for training.
+Initial experiments in data-driven weather forecasting have used the
+same input variables as output variables.
+
+Anemoi training implements data routing, in which you can specify which
+variables are used as ``forcings`` in the input only to inform the
+model, and which variables are used as ``diagnostics`` in the output
+only to be predicted by the model. All remaining variables will be
+treated as ``prognostic`` in the intial and forecast states.
+
+Intuitively, ``forcings`` are the variables like solar insolation or
+land-sea-mask. These would make little sense to predict as they are
+external to the model. ``Diagnostics`` are the variables like
+precipitation that we want to predict, but which may not be available in
+forecast step zero due to technical limitations. ``Prognostic``
+variables are the variables like temperature or humidity that we want to
+predict and are available after data assimilation operationally.
+
+The user can specify the routing of the data by setting the
+``config.data.forcings`` and ``config.data.diagnostics``. These are
+named strings, as Anemoi datasets enables us to address variables by
+name.
+
+This can look like the following:
+
+.. code:: yaml
+
+   data:
+      forcings:
+         - solar_insolation
+         - land_sea_mask
+      diagnostics:
+         - total_precipitation
+
+***************
+ Normalisation
+***************
+
+Machine learning models are sensitive to the scale of the input data. To
+ensure that the model can learn effectively, it is important to
+normalise the input data.
+
+Anemoi training provides preprocessors for different aspects of the
+training, with the normaliser being one of them. The normaliser
+implements multiple strategies that can be applied to the data using the
+config.
+
+Currently, the normaliser supports the following strategies:
+
+-  ``none``: No normalisation is applied.
+-  ``mean-std``: Standard normalisation is applied to the data.
+-  ``min-max``: Min-max normalisation is applied to the data.
+-  ``max``: Max normalisation is applied to the data.
+
+Values like the land-sea-mask do not require additional normalisation.
+However, variables like temperature or humidity should be normalised to
+ensure the model can learn effectively. Additionally, variables like the
+geopotential height should be max normalised to ensure the model can
+learn the vertical structure of the atmosphere.
+
+The user can specify the normalisation strategy, including the default
+by setting ``config.data.normaliser``, such that:
+
+.. code:: yaml
+
+   normaliser:
+      default: mean-std
+      none:
+         - land_sea_mask
+      max:
+         - geopotential_height
+
+***********************
+ Loss function scaling
+***********************
+
+It is possible to change the weighting given to each of the variables in
+the loss function by changing
+``config.training.loss_scaling.pl.<pressure level variable>`` and
+``config.training.loss_scaling.sfc.<surface variable>``.
+
+It is also possible to change the scaling given to the pressure levels
+using ``config.training.pressure_level_scaler``. For almost all
+applications, upper atmosphere pressure levels should be given lower
+weighting than the lower atmosphere pressure levels (i.e. pressure
+levels nearer to the surface). By default anemoi-training uses a ReLU
+Pressure Level scaler with a minimum weighting of 0.2 (i.e. no pressure
+level has a weighting less than 0.2).
+
+***************
+ Learning rate
+***************
+
+Anemoi training uses the ``CosineLRScheduler`` from PyTorch as it's
+learning rate scheduler. The user can configure the maximum learning
+rate by setting ``config.training.lr.rate``. Note that this learning
+rate is scaled by the number of GPUs where for the `data parallelism
+<distributed>`_.
+
+.. code:: yaml
+
+   global_learning_rate = config.training.lr.rate * num_gpus_per_node * num_nodes / gpus_per_model
+
+The user can also control the rate at which the learning rate decreases
+by setting the total number of iterations through
+``config.training.lr.iterations`` and the minimum learning rate reached
+through ``config.training.lr.min``. Note that the minimum learning rate
+is not scaled by the number of GPUs.
+
+*********
+ Rollout
+*********
+
+In the first stage of training, standard practice is to train the model
+on a 6 hour interval. Once this is completed, in the second stage of
+training, it is advisable to *rollout* and fine-tune the model error at
+longer leadtimes too. Generally for medium range forecasts, rollout is
+performed on 12 forecast steps (equivalent to 72 hours) incrementally.
+In other words, at each epoch another forecast step is added to the
+error term.
+
+Rollout requires the model training to be restarted so the user should
+make sure to set ``config.training.run_id`` equal to the run-id of the
+first stage of training.
+
+Note, in the standard set-up, rollout is performed at the minimum
+learning rate and the number of batches used is reduced (using
+``config.dataloader.training.limit_batches``) to prevent any overfit to
+specific timesteps.
+
+To start rollout set ``config.training.rollout.epoch_increment`` equal
+to 1 (thus increasing the rollout step by 1 at every epoch) and set a
+maximum rollout by setting ``config.training.rollout.max`` (usually set
+to 12).
+
+***************************
+ Restarting a training run
+***************************
+
+Whether it's because the training has exceeded the time limit on an HPC
+system or because the user wants to fine-tune the model from a specific
+point in the training, it may be necessary at certain points to restart
+the model training.
+
+This can be done by setting ``config.training.run_id`` in the config
+file to be the *run_id* of the run that is being restarted. In this case
+the new checkpoints will go in the same folder as the old checkpoints.
+If the user does not want this then they can instead set
+``config.training.fork_run_id`` in the config file to the *run_id* of
+the run that is being restarted. In this case the old run will be
+unaffected and the new checkpoints will go in to a new folder with a new
+run_id. The user might want to do this if they want to start multiple
+new runs from 1 old run.
+
+The above will restart the model training from where the old run
+finished training. However if the user wants to restart the model from a
+specific point they can do this by setting
+``config.hardware.files.warm_start`` to be the checkpoint they want to
+restart from..
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 00000000..4a368552
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,178 @@
+#!/usr/bin/env python
+# (C) Copyright 2024 ECMWF.
+#
+# 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.
+
+# https://packaging.python.org/en/latest/guides/writing-pyproject-toml/
+
+[build-system]
+requires = [ "setuptools>=60", "setuptools-scm>=8" ]
+
+[project]
+name = "anemoi-training"
+
+description = "A package to hold various functions to support training of ML models."
+keywords = [ "ai", "tools", "training" ]
+
+license = { file = "LICENSE" }
+authors = [
+  { name = "European Centre for Medium-Range Weather Forecasts (ECMWF)", email = "software.support@ecmwf.int" },
+]
+
+requires-python = ">=3.9"
+
+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Developers",
+  "License :: OSI Approved :: Apache Software License",
+  "Operating System :: OS Independent",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: Implementation :: CPython",
+  "Programming Language :: Python :: Implementation :: PyPy",
+]
+
+dynamic = [ "version" ]
+
+dependencies = [
+  "anemoi-datasets>=0.4",
+  "anemoi-graphs",
+  "anemoi-models>=0.3",
+  "anemoi-utils[provenance]>=0.3.10",
+  "einops>=0.6.1",
+  "hydra-core>=1.3",
+  "matplotlib>=3.7.1",
+  "mlflow>=2.11.1",
+  "numpy<2",                          # Pinned until we can confirm it works with anemoi graphs
+  "pynvml>=11.5",
+  "pyshtools>=4.10.4",
+  "pytorch-lightning>=2.1",
+  "timm>=0.9.2",
+  "torch>=2.2",
+  "torch-geometric>=2.3.1,<2.5",
+  "torchinfo>=1.8",
+  "tqdm>=4.65",
+  "zarr>=2.14.2",
+]
+
+optional-dependencies.all = [  ]
+
+optional-dependencies.dev = [
+  "anemoi-training[docs,tests,all]",
+  "pre-commit>=3.3.3",
+]
+
+optional-dependencies.docs = [
+  "nbsphinx",
+  "pandoc",
+  "sphinx",
+  "sphinx-argparse",
+  "sphinx-rtd-theme",
+]
+optional-dependencies.tests = [ "hypothesis", "pytest", "pytest-mock" ]
+
+urls.Changelog = "https://github.com/ecmwf/anemoi-training/CHANGELOG.md"
+urls.Documentation = "https://anemoi-training.readthedocs.io/"
+urls.Homepage = "https://github.com/ecmwf/anemoi-training/"
+urls.Issues = "https://github.com/ecmwf/anemoi-training/issues"
+urls.Repository = "https://github.com/ecmwf/anemoi-training/"
+# command for interactive DDP (not supposed to be used directly)
+# the dot is intentional, so it doesn't trigger autocomplete
+scripts.".anemoi-training-train" = "anemoi.training.commands.train:main"
+
+# Add subcommand in the `commands` directory
+scripts.anemoi-training = "anemoi.training.__main__:main"
+
+[tool.setuptools_scm]
+version_file = "src/anemoi/training/_version.py"
+
+[tool.ruff]
+target-version = "py39"
+line-length = 120
+src = [ "src" ]
+exclude = [ "docs/" ]
+
+preview = true
+lint.flake8-import-conventions.extend-aliases."pytorch_lightning" = "pl"
+lint.pydocstyle.convention = "numpy"
+select = [
+  "A", # flake8 builtins
+  # "ANN", # flake8-annotations
+  "ANN001", # Type annotation
+  "ANN201", # Return type public
+  "ANN202", # Return type private
+  "ANN205", # Return type static method
+  "ANN206", # Return type class method
+  "ARG",    # Unused arguments
+  "B",      # flake8-bugbear
+  "BLE",    # Non-specific exceptions
+  "C4",     # Comprehensions
+  "C90",    # Mccabe complexity
+  "COM",    # Commas
+  "CPY",    # Copyright
+  # "C",  # pylint convention
+  "D", # pydocstyle
+  # "DOC", # pydoclint
+  "DTZ", # Datetimes
+  "E",   # pycodestyle error
+  "EM",  # Error messages
+  "ERA", # Found unused code
+  "F",   # flake8 error
+  "FA",  # Future annotations
+  # "EXE", # Executable
+  # "FBT", # Boolean traps
+  "FLY",  # Fstrings
+  "FURB", # Modernising code
+  "G",    # flake8 logging format
+  "ICN",  # Import conventions
+  "ISC",  # Implicit string concatenation
+  "LOG",  # Logging
+  # "I", # isort
+  "N",    # pep8 naming
+  "NPY",  # Numpys
+  "PERF", # Perfomance linting
+  "PGH",  # Pygrep hooks
+  "PIE",  # Misc lints
+  # "PL",  # Pylint (superset of C, E, R, W)
+  "PT",  # Pytest
+  "PTH", # Use Pathlib
+  "Q",   # Quotes
+  "R",   # Refactor
+  "RET", # Return statements
+  "RSE", # Raises
+  "RUF", # Ruff specific
+  "S",   # bandit !IMPORTANT
+  "SIM", # flake8-simplify
+  "T10", # Debug statements
+  "T20", # Print statements
+  "TCH", # Type checking blocks
+  "TD",  # Todos
+  "TID", # Tidy imports
+  "TRY", # Exception handling antipatterns
+  "UP",  # Pyupgrade
+  "W",   # pycodestyle warning
+  "YTT", # flake8-2020
+]
+
+ignore = [
+  "B018",
+  "B028",
+  "D100",
+  "D101",
+  "D102",
+  "D103",
+  "D104",
+  "D105",
+  "D401",
+  "E203",
+  "PT018",
+  "S101",
+  "TD003",
+]
diff --git a/setup.py b/setup.py
deleted file mode 100644
index 4a4f1465..00000000
--- a/setup.py
+++ /dev/null
@@ -1,71 +0,0 @@
-#!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
-#
-# 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.
-#
-
-
-import io
-import os
-
-import setuptools
-
-
-def read(fname):
-    file_path = os.path.join(os.path.dirname(__file__), fname)
-    return io.open(file_path, encoding="utf-8").read()
-
-
-version = None
-for line in read("anemoi/training/__init__.py").split("\n"):
-    if line.startswith("__version__"):
-        version = line.split("=")[-1].strip()[1:-1]
-
-
-assert version
-
-
-install_requires = []
-
-
-doc_requires = ["sphinx", "sphinx_rtd_theme", "nbsphinx", "pandoc"]
-
-all_requires = install_requires
-dev_requires = doc_requires + all_requires
-
-setuptools.setup(
-    name="anemoi-training",
-    version=version,
-    description="A package to hold various functions to support training of ML models on ECMWF data.",
-    long_description=read("README.md"),
-    long_description_content_type="text/markdown",
-    author="European Centre for Medium-Range Weather Forecasts (ECMWF)",
-    author_email="software.support@ecmwf.int",
-    license="Apache License Version 2.0",
-    url="https://github.com/ecmwf/anemoi-training",
-    packages=setuptools.find_namespace_packages(include=["anemoi.*"]),
-    include_package_data=True,
-    install_requires=install_requires,
-    extras_require={
-        "dev": dev_requires,
-        "all": all_requires,
-    },
-    zip_safe=True,
-    keywords="tool",
-    classifiers=[
-        "Development Status :: 3 - Alpha",
-        "Intended Audience :: Developers",
-        "License :: OSI Approved :: Apache Software License",
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.9",
-        "Programming Language :: Python :: 3.10",
-        "Programming Language :: Python :: 3.11",
-        "Programming Language :: Python :: Implementation :: CPython",
-        "Programming Language :: Python :: Implementation :: PyPy",
-        "Operating System :: OS Independent",
-    ],
-)
diff --git a/tests/test_training.py b/src/anemoi/training/__init__.py
similarity index 84%
rename from tests/test_training.py
rename to src/anemoi/training/__init__.py
index 7a025cf7..d9a51e0b 100644
--- a/tests/test_training.py
+++ b/src/anemoi/training/__init__.py
@@ -6,9 +6,4 @@
 # nor does it submit to any jurisdiction.
 
 
-def test_training():
-    pass
-
-
-if __name__ == "__main__":
-    test_training()
+from ._version import __version__  # noqa: F401
diff --git a/src/anemoi/training/__main__.py b/src/anemoi/training/__main__.py
new file mode 100644
index 00000000..f571e1b1
--- /dev/null
+++ b/src/anemoi/training/__main__.py
@@ -0,0 +1,27 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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 anemoi.utils.cli import cli_main
+from anemoi.utils.cli import make_parser
+
+from . import __version__
+from .commands import COMMANDS
+
+
+# For read-the-docs
+def create_parser() -> None:
+    return make_parser(__doc__, COMMANDS)
+
+
+def main() -> None:
+    cli_main(__version__, __doc__, COMMANDS)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/anemoi/training/commands/__init__.py b/src/anemoi/training/commands/__init__.py
new file mode 100644
index 00000000..17413995
--- /dev/null
+++ b/src/anemoi/training/commands/__init__.py
@@ -0,0 +1,23 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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 pathlib import Path
+
+from anemoi.utils.cli import Command
+from anemoi.utils.cli import Failed
+from anemoi.utils.cli import register_commands
+
+__all__ = ["Command"]
+
+COMMANDS = register_commands(
+    Path(__file__).parent,
+    __name__,
+    lambda x: x.command(),
+    Failed,
+)
diff --git a/src/anemoi/training/commands/checkpoint.py b/src/anemoi/training/commands/checkpoint.py
new file mode 100644
index 00000000..af1aa296
--- /dev/null
+++ b/src/anemoi/training/commands/checkpoint.py
@@ -0,0 +1,47 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+import argparse
+import logging
+
+from anemoi.training.commands import Command
+
+LOG = logging.getLogger(__name__)
+
+
+class Checkpoint(Command):
+    """Commands to interact with training checkpoints."""
+
+    @staticmethod
+    def add_arguments(command_parser: argparse.ArgumentParser) -> None:
+        subparsers = command_parser.add_subparsers(dest="subcommand", required=True)
+
+        help_msg = "Save an anemoi training checkpoint for inference."
+        inference = subparsers.add_parser(
+            "inference",
+            help=help_msg,
+            description=help_msg,
+        )
+        inference.add_argument("--input", "-i", required=True, metavar="training.ckpt")
+        inference.add_argument("--output", "-o", required=True, metavar="inference.ckpt")
+
+    @staticmethod
+    def run(args: argparse.Namespace) -> None:
+        if args.subcommand == "inference":
+            LOG.info("Loading training checkpoint from %s, please wait...", args.input)
+
+            from anemoi.training.utils.checkpoint import load_and_prepare_model
+            from anemoi.training.utils.checkpoint import save_inference_checkpoint
+
+            module, metadata = load_and_prepare_model(args.input)
+            path = save_inference_checkpoint(module, metadata, args.output)
+
+            LOG.info("Inference checkpoint saved to %s", path)
+            return
+
+
+command = Checkpoint
diff --git a/src/anemoi/training/commands/config.py b/src/anemoi/training/commands/config.py
new file mode 100644
index 00000000..e79602d6
--- /dev/null
+++ b/src/anemoi/training/commands/config.py
@@ -0,0 +1,98 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 __future__ import annotations
+
+import importlib.resources as pkg_resources
+import logging
+import shutil
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from anemoi.training.commands import Command
+
+if TYPE_CHECKING:
+    import argparse
+
+LOGGER = logging.getLogger(__name__)
+
+
+class ConfigGenerator(Command):
+    """Commands to interact with training configs."""
+
+    @staticmethod
+    def add_arguments(command_parser: argparse.ArgumentParser) -> None:
+        subparsers = command_parser.add_subparsers(dest="subcommand", required=True)
+
+        help_msg = "Generate the Anemoi training configs."
+        generate = subparsers.add_parser(
+            "generate",
+            help=help_msg,
+            description=help_msg,
+        )
+        generate.add_argument("--output", "-o", default=Path.cwd(), help="Output directory")
+        generate.add_argument("--overwrite", "-f", action="store_true")
+
+        help_msg = "Generate the Anemoi training configs in home."
+        anemoi_training_home = subparsers.add_parser(
+            "training-home",
+            help=help_msg,
+            description=help_msg,
+        )
+        anemoi_training_home.add_argument("--overwrite", "-f", action="store_true")
+
+    def run(self, args: argparse.Namespace) -> None:
+        LOGGER.info(
+            "Generating configs, please wait.",
+        )
+        self.overwrite = args.overwrite
+        if args.subcommand == "generate":
+
+            self.traverse_config(args.output)
+
+            LOGGER.info("Inference checkpoint saved to %s", args.output)
+            return
+
+        if args.subcommand == "training-home":
+            anemoi_home = Path.home() / ".config" / "anemoi" / "training" / "config"
+            self.traverse_config(anemoi_home)
+            LOGGER.info("Inference checkpoint saved to %s", anemoi_home)
+            return
+
+    def traverse_config(self, destination_dir: Path | str) -> None:
+        """Writes the given configuration data to the specified file path."""
+        config_package = "anemoi.training.config"
+
+        # Ensure the destination directory exists
+        destination_dir = Path(destination_dir)
+        destination_dir.mkdir(parents=True, exist_ok=True)
+
+        # Traverse through the package's config directory
+        with pkg_resources.as_file(pkg_resources.files(config_package)) as config_path:
+            for data in config_path.rglob("*"):  # Recursively walk through all files and directories
+                item = Path(data)
+                if item.is_file() and item.suffix == ".yaml":
+                    file_path = Path(destination_dir, item.relative_to(config_path))
+
+                    file_path.parent.mkdir(parents=True, exist_ok=True)
+
+                    if not file_path.exists() or self.overwrite:
+                        self.copy_file(item, file_path)
+                    else:
+                        LOGGER.info("File %s already exists, skipping", file_path)
+
+    @staticmethod
+    def copy_file(item: Path, file_path: Path) -> None:
+        """Copies the file to the destination directory."""
+        try:
+            shutil.copy2(item, file_path)
+            LOGGER.info("Copied %s to %s", item.name, file_path)
+        except Exception:
+            LOGGER.exception("Failed to copy %s", item.name)
+
+
+command = ConfigGenerator
diff --git a/src/anemoi/training/commands/mlflow.py b/src/anemoi/training/commands/mlflow.py
new file mode 100644
index 00000000..8be8183d
--- /dev/null
+++ b/src/anemoi/training/commands/mlflow.py
@@ -0,0 +1,132 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+import argparse
+
+from anemoi.training.commands import Command
+
+
+class MlFlow(Command):
+    """Commands to interact with MLflow."""
+
+    @staticmethod
+    def add_arguments(command_parser: argparse.ArgumentParser) -> None:
+        subparsers = command_parser.add_subparsers(dest="subcommand", required=True)
+
+        help_msg = "Log in and acquire a token from keycloak."
+        login = subparsers.add_parser(
+            "login",
+            help=help_msg,
+            description=help_msg,
+        )
+        login.add_argument(
+            "--url",
+            help="The URL of the authentication server. If not provided, the last used URL will be tried.",
+        )
+        login.add_argument(
+            "--force-credentials",
+            "-f",
+            action="store_true",
+            help="Force a credential login even if a token is available.",
+        )
+
+        help_msg = "Synchronise an offline run with an MLflow server."
+        sync = subparsers.add_parser(
+            "sync",
+            help=help_msg,
+            description=help_msg,
+            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        )
+        sync.add_argument(
+            "--source",
+            "-s",
+            help="The MLflow logs source directory.",
+            metavar="DIR",
+            required=True,
+            default=argparse.SUPPRESS,
+        )
+        sync.add_argument(
+            "--destination",
+            "-d",
+            help="The destination MLflow tracking URI.",
+            metavar="URI",
+            required=True,
+            default=argparse.SUPPRESS,
+        )
+        sync.add_argument(
+            "--run-id",
+            "-r",
+            help="The run ID to sync.",
+            metavar="ID",
+            required=True,
+            default=argparse.SUPPRESS,
+        )
+        sync.add_argument(
+            "--experiment-name",
+            "-e",
+            help="The experiment name to sync to.",
+            metavar="NAME",
+            default="anemoi-debug",
+        )
+        sync.add_argument(
+            "--authentication",
+            "-a",
+            action="store_true",
+            help="The destination server requires authentication.",
+        )
+        sync.add_argument(
+            "--export-deleted-runs",
+            "-x",
+            action="store_true",
+        )
+        sync.add_argument(
+            "--verbose",
+            "-v",
+            action="store_true",
+        )
+
+    @staticmethod
+    def run(args: argparse.Namespace) -> None:
+        if args.subcommand == "login":
+            from anemoi.training.diagnostics.mlflow.auth import TokenAuth
+
+            url = args.url or TokenAuth.load_config().get("url")
+
+            if not url:
+                msg = "No URL provided and no past URL found. Rerun the command with --url"
+                raise ValueError(msg)
+
+            TokenAuth(url=url).login(force_credentials=args.force_credentials)
+            return
+
+        if args.subcommand == "sync":
+            from anemoi.training.diagnostics.mlflow.utils import health_check
+            from anemoi.training.utils.mlflow_sync import MlFlowSync
+
+            if args.authentication:
+                from anemoi.training.diagnostics.mlflow.auth import TokenAuth
+
+                auth = TokenAuth(url=args.destination)
+                auth.login()
+                auth.authenticate()
+
+            health_check(args.destination)
+
+            log_level = "DEBUG" if args.verbose else "INFO"
+
+            MlFlowSync(
+                args.source,
+                args.destination,
+                args.run_id,
+                args.experiment_name,
+                args.export_deleted_runs,
+                log_level,
+            ).sync()
+            return
+
+
+command = MlFlow
diff --git a/src/anemoi/training/commands/train.py b/src/anemoi/training/commands/train.py
new file mode 100644
index 00000000..705b4783
--- /dev/null
+++ b/src/anemoi/training/commands/train.py
@@ -0,0 +1,87 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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 __future__ import annotations
+
+import logging
+import os
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from anemoi.training.commands import Command
+
+if TYPE_CHECKING:
+    import argparse
+
+LOGGER = logging.getLogger(__name__)
+
+
+class Train(Command):
+    """Commands to train Anemoi models."""
+
+    accept_unknown_args = True
+
+    @staticmethod
+    def add_arguments(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+        return parser
+
+    def run(self, args: argparse.Namespace, unknown_args: list[str] | None = None) -> None:
+        # This will be picked up by the logger
+        os.environ["ANEMOI_TRAINING_CMD"] = f"{sys.argv[0]} {args.command}"
+        # Merge the known subcommands with a non-whitespace character for hydra
+        new_sysargv = self._merge_sysargv(args)
+
+        # Add the unknown arguments (belonging to hydra) to sys.argv
+        if unknown_args is not None:
+            sys.argv = [new_sysargv, *unknown_args]
+        else:
+            sys.argv = [new_sysargv]
+
+        LOGGER.info("Running anemoi training command with overrides: %s", sys.argv[1:])
+        main()
+
+    def _merge_sysargv(self, args: argparse.Namespace) -> str:
+        """Merge the sys.argv with the known subcommands to pass to hydra.
+
+        This is done for interactive DDP, which will spawn the rank > 0 processes from sys.argv[0]
+        and for hydra, which ingests sys.argv[1:]
+
+        Parameters
+        ----------
+        args : argparse.Namespace
+            args from the command line
+
+        Returns
+        -------
+        str
+            Modified sys.argv as string
+        """
+        argv = Path(sys.argv[0])
+
+        # this will turn "/env/bin/anemoi-training train" into "/env/bin/.anemoi-training-train"
+        # the dot at the beginning is intentional to not interfere with autocomplete
+        modified_sysargv = argv.with_name(f".{argv.name}-{args.command}")
+
+        if hasattr(args, "subcommand"):
+            modified_sysargv += f"-{args.subcommand}"
+        return str(modified_sysargv)
+
+
+def main() -> None:
+    # Use the environment variable to check if main is being called from the subcommand, not from the ddp entrypoint
+    if not os.environ.get("ANEMOI_TRAINING_CMD"):
+        error = "This entrypoint should not be called directly. Use `anemoi-training train` instead."
+        raise RuntimeError(error)
+
+    from anemoi.training.train.train import main as anemoi_train
+
+    anemoi_train()
+
+
+command = Train
diff --git a/anemoi/models/__init__.py b/src/anemoi/training/config/__init__.py
similarity index 94%
rename from anemoi/models/__init__.py
rename to src/anemoi/training/config/__init__.py
index 6c631a72..282d6a69 100644
--- a/anemoi/models/__init__.py
+++ b/src/anemoi/training/config/__init__.py
@@ -4,6 +4,3 @@
 # 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.
-
-
-__version__ = "0.0.1"
diff --git a/src/anemoi/training/config/config.yaml b/src/anemoi/training/config/config.yaml
new file mode 100644
index 00000000..2045da93
--- /dev/null
+++ b/src/anemoi/training/config/config.yaml
@@ -0,0 +1,9 @@
+defaults:
+- data: zarr
+- dataloader: native_grid
+- diagnostics: eval_rollout
+- hardware: example
+- graph: multi_scale
+- model: gnn
+- training: default
+- _self_
diff --git a/src/anemoi/training/config/data/zarr.yaml b/src/anemoi/training/config/data/zarr.yaml
new file mode 100644
index 00000000..1657861f
--- /dev/null
+++ b/src/anemoi/training/config/data/zarr.yaml
@@ -0,0 +1,71 @@
+format: zarr
+resolution: o96
+# Time frequency requested from dataset
+frequency: 6h
+# Time step of model (must be multiple of frequency)
+timestep: 6h
+
+# features that are not part of the forecast state
+# but are used as forcing to generate the forecast state
+forcing:
+- "cos_latitude"
+- "cos_longitude"
+- "sin_latitude"
+- "sin_longitude"
+- "cos_julian_day"
+- "cos_local_time"
+- "sin_julian_day"
+- "sin_local_time"
+- "insolation"
+- "lsm"
+- "sdor"
+- "slor"
+- "z"
+# features that are only part of the forecast state
+# but are not used as the input to the model
+diagnostic:
+- tp
+- cp
+remapped:
+
+normalizer:
+  default: "mean-std"
+  min-max:
+  max:
+  - "sdor"
+  - "slor"
+  - "z"
+  none:
+  - "cos_latitude"
+  - "cos_longitude"
+  - "sin_latitude"
+  - "sin_longitude"
+  - "cos_julian_day"
+  - "cos_local_time"
+  - "sin_julian_day"
+  - "sin_local_time"
+  - "insolation"
+  - "lsm"
+
+imputer:
+  default: "none"
+remapper:
+  default: "none"
+
+# processors including imputers and normalizers are applied in order of definition
+processors:
+  # example_imputer:
+  #   _target_: anemoi.models.preprocessing.imputer.InputImputer
+  #   _convert_: all
+  #   config: ${data.imputer}
+  normalizer:
+    _target_: anemoi.models.preprocessing.normalizer.InputNormalizer
+    _convert_: all
+    config: ${data.normalizer}
+  # remapper:
+  #   _target_: anemoi.models.preprocessing.remapper.Remapper
+  #   _convert_: all
+  #   config: ${data.remapper}
+
+# Values set in the code
+num_features: null # number of features in the forecast state
diff --git a/src/anemoi/training/config/dataloader/native_grid.yaml b/src/anemoi/training/config/dataloader/native_grid.yaml
new file mode 100644
index 00000000..35ed7e35
--- /dev/null
+++ b/src/anemoi/training/config/dataloader/native_grid.yaml
@@ -0,0 +1,58 @@
+prefetch_factor: 2
+
+num_workers:
+  training: 8
+  validation: 8
+  test: 8
+  predict: 8
+batch_size:
+  training: 2
+  validation: 4
+  test: 4
+  predict: 4
+
+# ============
+# Default effective batch_size for training is 16
+# For the o96 resolution, default per-gpu batch_size is 2 (8 gpus required)
+# The global lr is calculated as:
+# global_lr = local_lr * num_gpus_per_node * num_nodes / gpus_per_model
+# Assuming a constant effective batch_size, any change in the per_gpu batch_size
+# should come with a rescaling of the local_lr to keep a constant global_lr
+# ============
+
+# runs only N training batches [N = integer | null]
+# if null then we run through all the batches
+limit_batches:
+  training: null
+  validation: null
+  test: 20
+  predict: 20
+
+# ============
+# Dataloader definitions
+# These follow the anemoi-datasets patterns
+# You can make these as complicated for merging as you like
+# See https://anemoi-datasets.readthedocs.io
+# ============
+
+dataset: ${hardware.paths.data}/${hardware.files.dataset}
+
+training:
+  dataset: ${dataloader.dataset}
+  start: null
+  end: 2020
+  frequency: ${data.frequency}
+  drop:  []
+
+validation:
+  dataset: ${dataloader.dataset}
+  start: 2021
+  end: 2021
+  frequency: ${data.frequency}
+  drop:  []
+
+test:
+  dataset: ${dataloader.dataset}
+  start: 2022
+  end: null
+  frequency: ${data.frequency}
diff --git a/src/anemoi/training/config/debug.yaml b/src/anemoi/training/config/debug.yaml
new file mode 100644
index 00000000..5be3e9f4
--- /dev/null
+++ b/src/anemoi/training/config/debug.yaml
@@ -0,0 +1,32 @@
+defaults:
+- data: zarr
+- dataloader: native_grid
+- diagnostics: eval_rollout
+- hardware: example
+- graph: multi_scale
+- model: gnn
+- training: default
+- _self_
+
+
+### This file is for local experimentation.
+##  When you commit your changes, assign the new features and keywords
+##  to the correct defaults.
+# For example to change from default GPU count:
+# hardware:
+#   num_gpus_per_node: 1
+
+diagnostics:
+  plot:
+    enabled: False
+hardware:
+  files:
+    graph: ???
+model:
+  num_channels: 128
+dataloader:
+  limit_batches:
+    training: 100
+    validation: 100
+training:
+  max_epochs: 5
diff --git a/src/anemoi/training/config/diagnostics/eval_rollout.yaml b/src/anemoi/training/config/diagnostics/eval_rollout.yaml
new file mode 100644
index 00000000..50e9a647
--- /dev/null
+++ b/src/anemoi/training/config/diagnostics/eval_rollout.yaml
@@ -0,0 +1,100 @@
+---
+eval:
+  enabled: False
+  # use this to evaluate the model over longer rollouts, every so many validation batches
+  rollout: 12
+  frequency: 20
+plot:
+  enabled: True
+  asynchronous: True
+  frequency: 750
+  sample_idx: 0
+  per_sample: 6
+  parameters:
+  - z_500
+  - t_850
+  - u_850
+  - v_850
+  - 2t
+  - 10u
+  - 10v
+  - sp
+  - tp
+  - cp
+  #Defining the accumulation levels for precipitation related fields and the colormap
+  accumulation_levels_plot: [0, 0.05, 0.1, 0.25, 0.5, 1, 1.5, 2, 3, 4, 5, 6, 7, 100] # in mm
+  cmap_accumulation: ["#ffffff", "#04e9e7", "#019ff4", "#0300f4", "#02fd02", "#01c501", "#008e00", "#fdf802", "#e5bc00", "#fd9500", "#fd0000", "#d40000", "#bc0000", "#f800fd"]
+  precip_and_related_fields: [tp, cp]
+  # Histogram and Spectrum plots
+  parameters_histogram:
+  - z_500
+  - tp
+  - 2t
+  - 10u
+  - 10v
+  parameters_spectrum:
+  - z_500
+  - tp
+  - 2t
+  - 10u
+  - 10v
+  # group parameters by categories when visualizing contributions to the loss
+  # one-parameter groups are possible to highlight individual parameters
+  parameter_groups:
+    moisture: [tp, cp, tcw]
+    sfc_wind: [10u, 10v]
+  learned_features: False
+  longrollout:
+    enabled: False
+    rollout: [60]
+    frequency: 20 # every X epochs
+
+debug:
+  # this will detect and trace back NaNs / Infs etc. but will slow down training
+  anomaly_detection: False
+
+# activate the pytorch profiler (disable this in production)
+# remember to also activate the tensorboard logger (below)
+profiler: False
+
+checkpoint:
+  every_n_minutes:
+    save_frequency: 30 # Approximate, as this is checked at the end of training steps
+    num_models_saved: 3 # If set to k, saves the 'last' k model weights in the training.
+
+  every_n_epochs:
+    save_frequency: 1
+    num_models_saved: -1 # If set to -1, all checkpoints are kept ensuring runs can be continued/forked at any point in the training process
+
+  every_n_train_steps:
+    save_frequency: null # Does not scale with rollout
+    num_models_saved: 0
+
+log:
+  wandb:
+    enabled: False
+    offline: False
+    log_model: False
+    project: 'Anemoi'
+    entity: ???
+    # logger options (these probably come with some overhead)
+    gradients: False
+    parameters: False
+  tensorboard:
+    enabled: False
+  mlflow:
+    enabled: False
+    offline: False
+    authentication: False
+    log_model: False
+    tracking_uri: ???
+    experiment_name: 'anemoi-debug'
+    project_name: 'Anemoi'
+    system: True
+    terminal: True
+    run_name: null # If set to null, the run name will be the a random UUID
+    on_resume_create_child: True
+  interval: 100
+
+enable_progress_bar: True
+print_memory_summary: False
diff --git a/src/anemoi/training/config/graph/encoder_decoder_only.yaml b/src/anemoi/training/config/graph/encoder_decoder_only.yaml
new file mode 100644
index 00000000..b813254d
--- /dev/null
+++ b/src/anemoi/training/config/graph/encoder_decoder_only.yaml
@@ -0,0 +1,49 @@
+---
+overwrite: True
+
+data: "data"
+hidden: "hidden"
+
+nodes:
+  # Data nodes
+  data:
+    node_builder:
+      _target_: anemoi.graphs.nodes.ZarrDatasetNodes # options: ZarrDatasetNodes, NPZFileNodes
+      dataset: ${dataloader.dataset}
+    attributes: ${graph.attributes.nodes} # options: l1, l2, unit-max, unit-sum, unit-std
+  # Hidden nodes
+  hidden:
+    node_builder:
+      _target_: anemoi.graphs.nodes.NPZFileNodes # options: ZarrDatasetNodes, NPZFileNodes
+      grid_definition_path: ${hardware.paths.grids}
+      resolution: o48 # grid resolution for npz (o32, o48, ...)
+
+edges:
+# Encoder configuration
+- source_name: ${graph.data}
+  target_name: ${graph.hidden}
+  edge_builder:
+    _target_: anemoi.graphs.edges.CutOffEdges # options: KNNEdges, CutOffEdges
+    cutoff_factor: 0.6 # only for cutoff method
+  attributes: ${graph.attributes.edges}
+- source_name: ${graph.hidden}
+  # Decoder configuration
+  target_name: ${graph.data}
+  edge_builder:
+    _target_: anemoi.graphs.edges.KNNEdges # options: KNNEdges, CutOffEdges
+    num_nearest_neighbours: 3 # only for knn method
+  attributes: ${graph.attributes.edges}
+
+
+attributes:
+  nodes:
+    area_weight:
+      _target_: anemoi.graphs.nodes.attributes.AreaWeights # options: Area, Uniform
+      norm: unit-max # options: l1, l2, unit-max, unit-sum, unit-std
+  edges:
+    edge_length:
+      _target_: anemoi.graphs.edges.attributes.EdgeLength
+      norm: unit-std
+    edge_dirs:
+      _target_: anemoi.graphs.edges.attributes.EdgeDirection
+      norm: unit-std
diff --git a/src/anemoi/training/config/graph/multi_scale.yaml b/src/anemoi/training/config/graph/multi_scale.yaml
new file mode 100644
index 00000000..7e54535e
--- /dev/null
+++ b/src/anemoi/training/config/graph/multi_scale.yaml
@@ -0,0 +1,56 @@
+---
+overwrite: True
+
+data: "data"
+hidden: "hidden"
+
+nodes:
+  # Data nodes
+  data:
+    node_builder:
+      _target_: anemoi.graphs.nodes.ZarrDatasetNodes # options: ZarrDatasetNodes, NPZFileNodes
+      dataset: ${dataloader.dataset}
+    attributes: ${graph.attributes.nodes}
+  # Hidden nodes
+  hidden:
+    node_builder:
+      _target_: anemoi.graphs.nodes.TriNodes # options: ZarrDatasetNodes, NPZFileNodes, TriNodes
+      resolution: 5 # grid resolution for npz (o32, o48, ...)
+    attributes: ${graph.attributes.nodes}
+
+edges:
+# Encoder configuration
+- source_name: ${graph.data}
+  target_name: ${graph.hidden}
+  edge_builder:
+    _target_: anemoi.graphs.edges.CutOffEdges # options: KNNEdges, CutOffEdges
+    cutoff_factor: 0.6 # only for cutoff method
+  attributes: ${graph.attributes.edges}
+- source_name: ${graph.hidden}
+  # Processor configuration
+  target_name: ${graph.hidden}
+  edge_builder:
+    _target_: anemoi.graphs.edges.MultiScaleEdges
+    x_hops: 1
+  attributes: ${graph.attributes.edges}
+- source_name: ${graph.hidden}
+  # Decoder configuration
+  target_name: ${graph.data}
+  edge_builder:
+    _target_: anemoi.graphs.edges.KNNEdges # options: KNNEdges, CutOffEdges
+    num_nearest_neighbours: 3 # only for knn method
+  attributes: ${graph.attributes.edges}
+
+
+attributes:
+  nodes:
+    area_weight:
+      _target_: anemoi.graphs.nodes.attributes.AreaWeights # options: Area, Uniform
+      norm: unit-max # options: l1, l2, unit-max, unit-sum, unit-std
+  edges:
+    edge_length:
+      _target_: anemoi.graphs.edges.attributes.EdgeLength
+      norm: unit-std
+    edge_dirs:
+      _target_: anemoi.graphs.edges.attributes.EdgeDirection
+      norm: unit-std
diff --git a/src/anemoi/training/config/hardware/example.yaml b/src/anemoi/training/config/hardware/example.yaml
new file mode 100644
index 00000000..77921561
--- /dev/null
+++ b/src/anemoi/training/config/hardware/example.yaml
@@ -0,0 +1,10 @@
+---
+defaults:
+- paths: example
+- files: example
+
+# number of GPUs per node and number of nodes (for DDP)
+accelerator: auto
+num_gpus_per_node: 1
+num_nodes: 1
+num_gpus_per_model: 1
diff --git a/src/anemoi/training/config/hardware/files/example.yaml b/src/anemoi/training/config/hardware/files/example.yaml
new file mode 100644
index 00000000..0c8acb03
--- /dev/null
+++ b/src/anemoi/training/config/hardware/files/example.yaml
@@ -0,0 +1,7 @@
+dataset: ???
+graph: ???
+checkpoint:
+  every_n_epochs: anemoi-by_epoch-epoch_{epoch:03d}-step_{step:06d}
+  every_n_train_steps: anemoi-by_step-epoch_{epoch:03d}-step_{step:06d}
+  every_n_minutes: anemoi-by_time-epoch_{epoch:03d}-step_{step:06d}
+warm_start: null
diff --git a/src/anemoi/training/config/hardware/paths/example.yaml b/src/anemoi/training/config/hardware/paths/example.yaml
new file mode 100644
index 00000000..c6555f2c
--- /dev/null
+++ b/src/anemoi/training/config/hardware/paths/example.yaml
@@ -0,0 +1,12 @@
+data: ???
+grids: ???
+output: ???
+logs:
+  base: ${hardware.paths.output}logs/
+  wandb: ${hardware.paths.logs.base}
+  mlflow: ${hardware.paths.logs.base}mlflow/
+  tensorboard: ${hardware.paths.logs.base}tensorboard/
+checkpoints: ${hardware.paths.output}checkpoint/
+plots: ${hardware.paths.output}plots/
+profiler: ${hardware.paths.output}profiler/
+graph: ${hardware.paths.output}graphs/
diff --git a/src/anemoi/training/config/hardware/slurm.yaml b/src/anemoi/training/config/hardware/slurm.yaml
new file mode 100644
index 00000000..a40cd0c4
--- /dev/null
+++ b/src/anemoi/training/config/hardware/slurm.yaml
@@ -0,0 +1,10 @@
+---
+defaults:
+- paths: example
+- files: example
+
+# number of GPUs per node and number of nodes (for DDP)
+accelerator: auto
+num_gpus_per_node: ${oc.decode:${oc.env:SLURM_GPUS_PER_NODE}}
+num_nodes: ${oc.decode:${oc.env:SLURM_NNODES}}
+num_gpus_per_model: 1
diff --git a/src/anemoi/training/config/model/gnn.yaml b/src/anemoi/training/config/model/gnn.yaml
new file mode 100644
index 00000000..a01bf860
--- /dev/null
+++ b/src/anemoi/training/config/model/gnn.yaml
@@ -0,0 +1,48 @@
+activation: GELU
+num_channels: 512
+
+model:
+  _target_: anemoi.models.models.encoder_processor_decoder.AnemoiModelEncProcDec
+
+processor:
+  _target_: anemoi.models.layers.processor.GNNProcessor
+  _convert_: all
+  activation: ${model.activation}
+  trainable_size: ${model.trainable_parameters.hidden2hidden}
+  sub_graph_edge_attributes: ${model.attributes.edges}
+  num_layers: 16
+  num_chunks: 2
+  mlp_extra_layers: 0
+
+encoder:
+  _target_: anemoi.models.layers.mapper.GNNForwardMapper
+  _convert_: all
+  trainable_size: ${model.trainable_parameters.data2hidden}
+  sub_graph_edge_attributes: ${model.attributes.edges}
+  activation: ${model.activation}
+  num_chunks: 1
+  mlp_extra_layers: 0
+
+decoder:
+  _target_: anemoi.models.layers.mapper.GNNBackwardMapper
+  _convert_: all
+  trainable_size: ${model.trainable_parameters.hidden2data}
+  sub_graph_edge_attributes: ${model.attributes.edges}
+  activation: ${model.activation}
+  num_chunks: 1
+  mlp_extra_layers: 0
+
+trainable_parameters:
+  data: 8
+  hidden: 8
+  data2hidden: 8
+  hidden2data: 8
+  hidden2hidden: 8
+
+attributes:
+  edges:
+  - edge_length
+  - edge_dirs
+  nodes: []
+
+node_loss_weight: area_weight
diff --git a/src/anemoi/training/config/model/graphtransformer.yaml b/src/anemoi/training/config/model/graphtransformer.yaml
new file mode 100644
index 00000000..610de803
--- /dev/null
+++ b/src/anemoi/training/config/model/graphtransformer.yaml
@@ -0,0 +1,53 @@
+activation: GELU
+num_channels: 1024
+
+model:
+  _target_: anemoi.models.models.encoder_processor_decoder.AnemoiModelEncProcDec
+
+processor:
+  _target_: anemoi.models.layers.processor.GraphTransformerProcessor
+  _convert_: all
+  activation: ${model.activation}
+  trainable_size: ${model.trainable_parameters.hidden2hidden}
+  sub_graph_edge_attributes: ${model.attributes.edges}
+  num_layers: 16
+  num_chunks: 2
+  mlp_hidden_ratio: 4 # GraphTransformer or Transformer only
+  num_heads: 16 # GraphTransformer or Transformer only
+  dropout_p: 0.0 # GraphTransformer or Transformer only
+
+encoder:
+  _target_: anemoi.models.layers.mapper.GraphTransformerForwardMapper
+  _convert_: all
+  trainable_size: ${model.trainable_parameters.data2hidden}
+  sub_graph_edge_attributes: ${model.attributes.edges}
+  activation: ${model.activation}
+  num_chunks: 1
+  mlp_hidden_ratio: 4 # GraphTransformer or Transformer only
+  num_heads: 16 # GraphTransformer or Transformer only
+
+decoder:
+  _target_: anemoi.models.layers.mapper.GraphTransformerBackwardMapper
+  _convert_: all
+  trainable_size: ${model.trainable_parameters.hidden2data}
+  sub_graph_edge_attributes: ${model.attributes.edges}
+  activation: ${model.activation}
+  num_chunks: 1
+  mlp_hidden_ratio: 4 # GraphTransformer or Transformer only
+  num_heads: 16 # GraphTransformer or Transformer only
+
+trainable_parameters:
+  data: 8
+  hidden: 8
+  data2hidden: 8
+  hidden2data: 8
+  hidden2hidden: 8 # GNN and GraphTransformer Processor only
+
+
+attributes:
+  edges:
+  - edge_length
+  - edge_dirs
+  nodes: []
+
+node_loss_weight: area_weight
diff --git a/src/anemoi/training/config/model/transformer.yaml b/src/anemoi/training/config/model/transformer.yaml
new file mode 100644
index 00000000..7c490990
--- /dev/null
+++ b/src/anemoi/training/config/model/transformer.yaml
@@ -0,0 +1,52 @@
+activation: GELU
+num_channels: 1024
+
+model:
+  _target_: anemoi.models.models.encoder_processor_decoder.AnemoiModelEncProcDec
+
+processor:
+  _target_: anemoi.models.layers.processor.TransformerProcessor
+  _convert_: all
+  activation: ${model.activation}
+  num_layers: 16
+  num_chunks: 2
+  mlp_hidden_ratio: 4 # GraphTransformer or Transformer only
+  num_heads: 16 # GraphTransformer or Transformer only
+  window_size: 512
+  dropout_p: 0.0 # GraphTransformer
+
+encoder:
+  _target_: anemoi.models.layers.mapper.GraphTransformerForwardMapper
+  _convert_: all
+  trainable_size: ${model.trainable_parameters.data2hidden}
+  sub_graph_edge_attributes: ${model.attributes.edges}
+  activation: ${model.activation}
+  num_chunks: 1
+  mlp_hidden_ratio: 4 # GraphTransformer or Transformer only
+  num_heads: 16 # GraphTransformer or Transformer only
+
+
+decoder:
+  _target_: anemoi.models.layers.mapper.GraphTransformerBackwardMapper
+  _convert_: all
+  trainable_size: ${model.trainable_parameters.hidden2data}
+  sub_graph_edge_attributes: ${model.attributes.edges}
+  activation: ${model.activation}
+  num_chunks: 1
+  mlp_hidden_ratio: 4 # GraphTransformer or Transformer only
+  num_heads: 16 # GraphTransformer or Transformer only
+
+trainable_parameters:
+  data: 8
+  hidden: 8
+  data2hidden: 8
+  hidden2data: 8
+
+
+attributes:
+  edges:
+  - edge_length
+  - edge_dirs
+  nodes: []
+
+node_loss_weight: area_weight
diff --git a/src/anemoi/training/config/training/default.yaml b/src/anemoi/training/config/training/default.yaml
new file mode 100644
index 00000000..870eeb7a
--- /dev/null
+++ b/src/anemoi/training/config/training/default.yaml
@@ -0,0 +1,85 @@
+# resume or fork a training from a checkpoint last.ckpt or specified in hardware.files.warm_start
+run_id: null
+fork_run_id: null
+load_weights_only: null # only load model weights, do not restore optimiser states etc.
+
+# run in deterministic mode ; slows down
+deterministic: False
+
+# miscellaneous
+precision: 16-mixed
+
+# multistep input
+# 1 = single step scheme, X(t-1) used to predict X(t)
+# k > 1: multistep scheme, uses [X(t-k), X(t-k+1), ... X(t-1)] to predict X(t)
+# Deepmind use k = 2 in their model
+multistep_input: 2
+
+# gradient accumulation across K batches, K >= 1 (if K == 1 then no accumulation)
+# the effective batch size becomes num-devices * batch_size * k
+accum_grad_batches: 1
+
+# clipp gradients, 0 : don't clip, default algorithm: norm, alternative: value
+gradient_clip:
+  val: 32.
+  algorithm: value
+
+# stochastic weight averaging
+# https://pytorch.org/blog/stochastic-weight-averaging-in-pytorch/
+swa:
+  enabled: False
+  lr: 1.e-4
+
+# use ZeroRedundancyOptimizer ; saves memory for larger models
+zero_optimizer: False
+
+# dynamic rescaling of the loss gradient
+# see https://arxiv.org/pdf/2306.06079.pdf, section 4.3.2
+# don't enable this by default until it's been tested and proven beneficial
+loss_gradient_scaling: False
+
+# length of the "rollout" window (see Keisler's paper)
+rollout:
+  start: 1
+  # increase rollout every n epochs
+  epoch_increment: 0
+  # maximum rollout to use
+  max: 1
+
+max_epochs: 200
+lr:
+  rate: 0.625e-4 #local_lr
+  iterations: 300000
+  min: 3e-7 #Not scaled by #GPU
+
+# Changes in per-gpu batch_size should come with a rescaling of the local_lr
+# in order to keep a constant global_lr
+# global_lr = local_lr * num_gpus_per_node * num_nodes / gpus_per_model
+
+loss_scaling:
+  default: 1
+  pl:
+    q: 0.6 #1
+    t: 6   #1
+    u: 0.8 #0.5
+    v: 0.5 #0.33
+    w: 0.001
+    z: 12  #1
+  sfc:
+    sp: 10
+    10u: 0.1
+    10v: 0.1
+    2d: 0.5
+    tp: 0.025
+    cp: 0.0025
+
+metrics:
+- z_500
+- t_850
+- u_850
+- v_850
+
+pressure_level_scaler:
+  _target_: anemoi.training.data.scaling.ReluPressureLevelScaler
+  minimum: 0.2
+  slope: 0.001
diff --git a/src/anemoi/training/data/__init__.py b/src/anemoi/training/data/__init__.py
new file mode 100644
index 00000000..282d6a69
--- /dev/null
+++ b/src/anemoi/training/data/__init__.py
@@ -0,0 +1,6 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
diff --git a/src/anemoi/training/data/datamodule.py b/src/anemoi/training/data/datamodule.py
new file mode 100644
index 00000000..1e119892
--- /dev/null
+++ b/src/anemoi/training/data/datamodule.py
@@ -0,0 +1,203 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+import logging
+import os
+from functools import cached_property
+from typing import Callable
+
+import pytorch_lightning as pl
+from anemoi.datasets.data import open_dataset
+from anemoi.models.data_indices.collection import IndexCollection
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
+from torch.utils.data import DataLoader
+
+from anemoi.training.data.dataset import NativeGridDataset
+from anemoi.training.data.dataset import worker_init_func
+
+LOGGER = logging.getLogger(__name__)
+
+
+class AnemoiDatasetsDataModule(pl.LightningDataModule):
+    """Anemoi Datasets data module for PyTorch Lightning."""
+
+    def __init__(self, config: DictConfig) -> None:
+        """Initialize Anemoi Datasets data module.
+
+        Parameters
+        ----------
+        config : DictConfig
+            Job configuration
+
+        """
+        super().__init__()
+
+        self.config = config
+
+        # Determine the step size relative to the data frequency
+        frequency = self.config.data.frequency
+        timestep = self.config.data.timestep
+        assert (
+            isinstance(frequency, str) and isinstance(timestep, str) and frequency[-1] == "h" and timestep[-1] == "h"
+        ), f"Error in format of timestep, {timestep}, or data frequency, {frequency}"
+        assert (
+            int(timestep[:-1]) % int(frequency[:-1]) == 0
+        ), f"Timestep isn't a multiple of data frequency, {timestep}, or data frequency, {frequency}"
+        self.timeincrement = int(timestep[:-1]) // int(frequency[:-1])
+        LOGGER.info(
+            "Timeincrement set to %s for data with frequency, %s, and timestep, %s",
+            self.timeincrement,
+            frequency,
+            timestep,
+        )
+
+        self.global_rank = int(os.environ.get("SLURM_PROCID", "0"))  # global rank
+        self.model_comm_group_id = (
+            self.global_rank // self.config.hardware.num_gpus_per_model
+        )  # id of the model communication group the rank is participating in
+        self.model_comm_group_rank = (
+            self.global_rank % self.config.hardware.num_gpus_per_model
+        )  # rank within one model communication group
+        total_gpus = self.config.hardware.num_gpus_per_node * self.config.hardware.num_nodes
+        assert (
+            total_gpus
+        ) % self.config.hardware.num_gpus_per_model == 0, (
+            f"GPUs per model {self.config.hardware.num_gpus_per_model} does not divide total GPUs {total_gpus}"
+        )
+        self.model_comm_num_groups = (
+            self.config.hardware.num_gpus_per_node
+            * self.config.hardware.num_nodes
+            // self.config.hardware.num_gpus_per_model
+        )  # number of model communication groups
+        LOGGER.debug(
+            "Rank %d model communication group number %d, with local model communication group rank %d",
+            self.global_rank,
+            self.model_comm_group_id,
+            self.model_comm_group_rank,
+        )
+
+        # Set the maximum rollout to be expected
+        self.rollout = (
+            self.config.training.rollout.max
+            if self.config.training.rollout.epoch_increment > 0
+            else self.config.training.rollout.start
+        )
+
+        # Set the training end date if not specified
+        if self.config.dataloader.training.end is None:
+            LOGGER.info(
+                "No end date specified for training data, setting default before validation start date %s.",
+                self.config.dataloader.validation.start - 1,
+            )
+            self.config.dataloader.training.end = self.config.dataloader.validation.start - 1
+
+    def _check_resolution(self, resolution: str) -> None:
+        assert (
+            self.config.data.resolution.lower() == resolution.lower()
+        ), f"Network resolution {self.config.data.resolution=} does not match dataset resolution {resolution=}"
+
+    @cached_property
+    def statistics(self) -> dict:
+        return self.ds_train.statistics
+
+    @cached_property
+    def metadata(self) -> dict:
+        return self.ds_train.metadata
+
+    @cached_property
+    def data_indices(self) -> IndexCollection:
+        return IndexCollection(self.config, self.ds_train.name_to_index)
+
+    @cached_property
+    def ds_train(self) -> NativeGridDataset:
+        return self._get_dataset(
+            open_dataset(OmegaConf.to_container(self.config.dataloader.training, resolve=True)),
+            label="train",
+        )
+
+    @cached_property
+    def ds_valid(self) -> NativeGridDataset:
+        r = self.rollout
+        if self.config.diagnostics.eval.enabled:
+            r = max(r, self.config.diagnostics.eval.rollout)
+        if self.config.diagnostics.plot.get("longrollout") and self.config.diagnostics.plot.longrollout.enabled:
+            r = max(r, max(self.config.diagnostics.plot.longrollout.rollout))
+        assert self.config.dataloader.training.end < self.config.dataloader.validation.start, (
+            f"Training end date {self.config.dataloader.training.end} is not before"
+            f"validation start date {self.config.dataloader.validation.start}"
+        )
+        return self._get_dataset(
+            open_dataset(OmegaConf.to_container(self.config.dataloader.validation, resolve=True)),
+            shuffle=False,
+            rollout=r,
+            label="validation",
+        )
+
+    @cached_property
+    def ds_test(self) -> NativeGridDataset:
+        assert self.config.dataloader.training.end < self.config.dataloader.test.start, (
+            f"Training end date {self.config.dataloader.training.end} is not before"
+            f"test start date {self.config.dataloader.test.start}"
+        )
+        assert self.config.dataloader.validation.end < self.config.dataloader.test.start, (
+            f"Validation end date {self.config.dataloader.validation.end} is not before"
+            f"test start date {self.config.dataloader.test.start}"
+        )
+        return self._get_dataset(
+            open_dataset(OmegaConf.to_container(self.config.dataloader.test, resolve=True)),
+            shuffle=False,
+            label="test",
+        )
+
+    def _get_dataset(
+        self,
+        data_reader: Callable,
+        shuffle: bool = True,
+        rollout: int = 1,
+        label: str = "generic",
+    ) -> NativeGridDataset:
+        r = max(rollout, self.rollout)
+        data = NativeGridDataset(
+            data_reader=data_reader,
+            rollout=r,
+            multistep=self.config.training.multistep_input,
+            timeincrement=self.timeincrement,
+            model_comm_group_rank=self.model_comm_group_rank,
+            model_comm_group_id=self.model_comm_group_id,
+            model_comm_num_groups=self.model_comm_num_groups,
+            shuffle=shuffle,
+            label=label,
+        )
+        self._check_resolution(data.resolution)
+        return data
+
+    def _get_dataloader(self, ds: NativeGridDataset, stage: str) -> DataLoader:
+        assert stage in {"training", "validation", "test"}
+        return DataLoader(
+            ds,
+            batch_size=self.config.dataloader.batch_size[stage],
+            # number of worker processes
+            num_workers=self.config.dataloader.num_workers[stage],
+            # use of pinned memory can speed up CPU-to-GPU data transfers
+            # see https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-pinning
+            pin_memory=True,
+            # worker initializer
+            worker_init_fn=worker_init_func,
+            # prefetch batches
+            prefetch_factor=self.config.dataloader.prefetch_factor,
+            persistent_workers=True,
+        )
+
+    def train_dataloader(self) -> DataLoader:
+        return self._get_dataloader(self.ds_train, "training")
+
+    def val_dataloader(self) -> DataLoader:
+        return self._get_dataloader(self.ds_valid, "validation")
+
+    def test_dataloader(self) -> DataLoader:
+        return self._get_dataloader(self.ds_test, "test")
diff --git a/src/anemoi/training/data/dataset.py b/src/anemoi/training/data/dataset.py
new file mode 100644
index 00000000..e2aa12bd
--- /dev/null
+++ b/src/anemoi/training/data/dataset.py
@@ -0,0 +1,273 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 __future__ import annotations
+
+import logging
+import os
+import random
+from functools import cached_property
+from typing import Callable
+
+import numpy as np
+import torch
+from einops import rearrange
+from torch.utils.data import IterableDataset
+from torch.utils.data import get_worker_info
+
+from anemoi.training.utils.seeding import get_base_seed
+from anemoi.training.utils.usable_indices import get_usable_indices
+
+LOGGER = logging.getLogger(__name__)
+
+
+class NativeGridDataset(IterableDataset):
+    """Iterable dataset for AnemoI data on the arbitrary grids."""
+
+    def __init__(
+        self,
+        data_reader: Callable,
+        rollout: int = 1,
+        multistep: int = 1,
+        timeincrement: int = 1,
+        model_comm_group_rank: int = 0,
+        model_comm_group_id: int = 0,
+        model_comm_num_groups: int = 1,
+        shuffle: bool = True,
+        label: str = "generic",
+    ) -> None:
+        """Initialize (part of) the dataset state.
+
+        Parameters
+        ----------
+        data_reader : Callable
+            user function that opens and returns the zarr array data
+        rollout : int, optional
+            length of rollout window, by default 12
+        timeincrement : int, optional
+            time increment between samples, by default 1
+        multistep : int, optional
+            collate (t-1, ... t - multistep) into the input state vector, by default 1
+        model_comm_group_rank : int, optional
+            process rank in the torch.distributed group (important when running on multiple GPUs), by default 0
+        model_comm_group_id: int, optional
+            device group ID, default 0
+        model_comm_num_groups : int, optional
+            total number of device groups, by default 1
+        shuffle : bool, optional
+            Shuffle batches, by default True
+        label : str, optional
+            label for the dataset, by default "generic"
+
+        """
+        self.label = label
+
+        self.data = data_reader
+
+        self.rollout = rollout
+        self.timeincrement = timeincrement
+
+        # lazy init
+        self.n_samples_per_epoch_total: int = 0
+        self.n_samples_per_epoch_per_worker: int = 0
+
+        # DDP-relevant info
+        self.model_comm_group_rank = model_comm_group_rank
+        self.model_comm_num_groups = model_comm_num_groups
+        self.model_comm_group_id = model_comm_group_id
+        self.global_rank = int(os.environ.get("SLURM_PROCID", "0"))
+
+        # additional state vars (lazy init)
+        self.n_samples_per_worker = 0
+        self.chunk_index_range: np.ndarray | None = None
+        self.shuffle = shuffle
+
+        # Data dimensions
+        self.multi_step = multistep
+        assert self.multi_step > 0, "Multistep value must be greater than zero."
+        self.ensemble_dim: int = 2
+        self.ensemble_size = self.data.shape[self.ensemble_dim]
+
+    @cached_property
+    def statistics(self) -> dict:
+        """Return dataset statistics."""
+        return self.data.statistics
+
+    @cached_property
+    def metadata(self) -> dict:
+        """Return dataset metadata."""
+        return self.data.metadata()
+
+    @cached_property
+    def name_to_index(self) -> dict:
+        """Return dataset statistics."""
+        return self.data.name_to_index
+
+    @cached_property
+    def resolution(self) -> dict:
+        """Return dataset resolution."""
+        return self.data.resolution
+
+    @cached_property
+    def valid_date_indices(self) -> np.ndarray:
+        """Return valid date indices.
+
+        A date t is valid if we can sample the sequence
+            (t - multistep + 1, ..., t + rollout)
+        without missing data (if time_increment is 1).
+
+        If there are no missing dates, total number of valid ICs is
+        dataset length minus rollout minus additional multistep inputs
+        (if time_increment is 1).
+        """
+        return get_usable_indices(self.data.missing, len(self.data), self.rollout, self.multi_step, self.timeincrement)
+
+    def per_worker_init(self, n_workers: int, worker_id: int) -> None:
+        """Called by worker_init_func on each copy of dataset.
+
+        This initialises after the worker process has been spawned.
+
+        Parameters
+        ----------
+        n_workers : int
+            Number of workers
+        worker_id : int
+            Worker ID
+
+        """
+        self.worker_id = worker_id
+
+        # Divide this equally across shards (one shard per group!)
+        shard_size = len(self.valid_date_indices) // self.model_comm_num_groups
+        shard_start = self.model_comm_group_id * shard_size
+        shard_end = (self.model_comm_group_id + 1) * shard_size
+
+        shard_len = shard_end - shard_start
+        self.n_samples_per_worker = shard_len // n_workers
+
+        low = shard_start + worker_id * self.n_samples_per_worker
+        high = min(shard_start + (worker_id + 1) * self.n_samples_per_worker, shard_end)
+
+        LOGGER.debug(
+            "Worker %d (pid %d, global_rank %d, model comm group %d) has low/high range %d / %d",
+            worker_id,
+            os.getpid(),
+            self.global_rank,
+            self.model_comm_group_id,
+            low,
+            high,
+        )
+
+        self.chunk_index_range = self.valid_date_indices[np.arange(low, high, dtype=np.uint32)]
+
+        # each worker must have a different seed for its random number generator,
+        # otherwise all the workers will output exactly the same data
+        # should we check lightning env variable "PL_SEED_WORKERS" here?
+        # but we alwyas want to seed these anyways ...
+
+        base_seed = get_base_seed()
+
+        seed = (
+            base_seed * (self.model_comm_group_id + 1) - worker_id
+        )  # note that test, validation etc. datasets get same seed
+        torch.manual_seed(seed)
+        random.seed(seed)
+        self.rng = np.random.default_rng(seed=seed)
+        sanity_rnd = self.rng.random(1)
+
+        LOGGER.debug(
+            (
+                "Worker %d (%s, pid %d, glob. rank %d, model comm group %d, "
+                "group_rank %d, base_seed %d) using seed %d, sanity rnd %f"
+            ),
+            worker_id,
+            self.label,
+            os.getpid(),
+            self.global_rank,
+            self.model_comm_group_id,
+            self.model_comm_group_rank,
+            base_seed,
+            seed,
+            sanity_rnd,
+        )
+
+    def __iter__(self) -> torch.Tensor:
+        """Return an iterator over the dataset.
+
+        The datasets are retrieved by Anemoi Datasets from zarr files. This iterator yields
+        chunked batches for DDP and sharded training.
+
+        Currently it receives data with an ensemble dimension, which is discarded for
+        now. (Until the code is "ensemble native".)
+        """
+        if self.shuffle:
+            shuffled_chunk_indices = self.rng.choice(
+                self.chunk_index_range,
+                size=self.n_samples_per_worker,
+                replace=False,
+            )
+        else:
+            shuffled_chunk_indices = self.chunk_index_range
+
+        LOGGER.debug(
+            (
+                "Worker pid %d, label %s, worker id %d, global_rank %d, "
+                "model comm group %d, group_rank %d using indices[0:10]: %s"
+            ),
+            os.getpid(),
+            self.label,
+            self.worker_id,
+            self.global_rank,
+            self.model_comm_group_id,
+            self.model_comm_group_rank,
+            shuffled_chunk_indices[:10],
+        )
+
+        for i in shuffled_chunk_indices:
+            start = i - (self.multi_step - 1) * self.timeincrement
+            end = i + (self.rollout + 1) * self.timeincrement
+
+            x = self.data[start : end : self.timeincrement]
+            x = rearrange(x, "dates variables ensemble gridpoints -> dates ensemble gridpoints variables")
+            self.ensemble_dim = 1
+
+            yield torch.from_numpy(x)
+
+    def __repr__(self) -> str:
+        return f"""
+            {super().__repr__()}
+            Dataset: {self.data}
+            Rollout: {self.rollout}
+            Multistep: {self.multi_step}
+            Timeincrement: {self.timeincrement}
+        """
+
+
+def worker_init_func(worker_id: int) -> None:
+    """Configures each dataset worker process.
+
+    Calls WeatherBenchDataset.per_worker_init() on each dataset object.
+
+    Parameters
+    ----------
+    worker_id : int
+        Worker ID
+
+    Raises
+    ------
+    RuntimeError
+        If worker_info is None
+
+    """
+    worker_info = get_worker_info()  # information specific to each worker process
+    if worker_info is None:
+        LOGGER.error("worker_info is None! Set num_workers > 0 in your dataloader!")
+        raise RuntimeError
+    dataset_obj = worker_info.dataset  # the copy of the dataset held by this worker process.
+    dataset_obj.per_worker_init(
+        n_workers=worker_info.num_workers,
+        worker_id=worker_id,
+    )
diff --git a/src/anemoi/training/data/scaling.py b/src/anemoi/training/data/scaling.py
new file mode 100644
index 00000000..74ba9c23
--- /dev/null
+++ b/src/anemoi/training/data/scaling.py
@@ -0,0 +1,76 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+import logging
+from abc import ABC
+from abc import abstractmethod
+
+import numpy as np
+
+LOGGER = logging.getLogger(__name__)
+
+
+class BasePressureLevelScaler(ABC):
+    """Configurable method converting pressure level of variable to PTL scaling.
+
+    Scaling variables depending on pressure levels (50 to 1000).
+    """
+
+    def __init__(self, slope: float = 1.0 / 1000, minimum: float = 0.0) -> None:
+        """Initialise Scaler with slope and minimum.
+
+        Parameters
+        ----------
+        slope : float
+            Slope of the scaling function.
+        minimum : float
+            Minimum value of the scaling function.
+
+        """
+        self.slope = slope
+        self.minimum = minimum
+
+    @abstractmethod
+    def scaler(self, plev: float) -> np.ndarray: ...
+
+
+class LinearPressureLevelScaler(BasePressureLevelScaler):
+    """Linear with slope self.slope, yaxis shift by self.minimum."""
+
+    def scaler(self, plev: float) -> np.ndarray:
+        return plev * self.slope + self.minimum
+
+
+class ReluPressureLevelScaler(BasePressureLevelScaler):
+    """Linear above self.minimum, taking constant value self.minimum below."""
+
+    def scaler(self, plev: float) -> np.ndarray:
+        return max(self.minimum, plev * self.slope)
+
+
+class PolynomialPressureLevelScaler(BasePressureLevelScaler):
+    """Polynomial scaling, (slope * plev)^2, yaxis shift by self.minimum."""
+
+    def scaler(self, plev: float) -> np.ndarray:
+        return (self.slope * plev) ** 2 + self.minimum
+
+
+class NoPressureLevelScaler(BasePressureLevelScaler):
+    """Constant scaling by 1.0."""
+
+    def __init__(self, slope: float = 0.0, minimum: float = 1.0) -> None:
+        """Initialise Scaler with constant scaling of 1."""
+        assert (
+            minimum == 1.0 and slope == 0
+        ), "self.minimum must be 1.0 and self.slope 0.0 for no scaling to fit with definition of linear function."
+        super().__init__(slope=0.0, minimum=1.0)
+
+    @staticmethod
+    def scaler(plev: float) -> np.ndarray:
+        del plev  # unused
+        # no scaling, always return 1.0
+        return 1.0
diff --git a/src/anemoi/training/diagnostics/__init__.py b/src/anemoi/training/diagnostics/__init__.py
new file mode 100644
index 00000000..282d6a69
--- /dev/null
+++ b/src/anemoi/training/diagnostics/__init__.py
@@ -0,0 +1,6 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
diff --git a/src/anemoi/training/diagnostics/callbacks/__init__.py b/src/anemoi/training/diagnostics/callbacks/__init__.py
new file mode 100644
index 00000000..f2195b5f
--- /dev/null
+++ b/src/anemoi/training/diagnostics/callbacks/__init__.py
@@ -0,0 +1,1138 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+# * [WHY ARE CALLBACKS UNDER __init__.py?]
+# * This functionality will be restructured in the near future
+# * so for now callbacks are under __init__.py
+
+from __future__ import annotations
+
+import copy
+import logging
+import sys
+import time
+import traceback
+import uuid
+from abc import ABC
+from abc import abstractmethod
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import nullcontext
+from datetime import timedelta
+from functools import cached_property
+from pathlib import Path
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import Callable
+
+import matplotlib.patches as mpatches
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torchinfo
+from anemoi.utils.checkpoints import save_metadata
+from pytorch_lightning.callbacks import Callback
+from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
+from pytorch_lightning.utilities import rank_zero_only
+
+from anemoi.training.diagnostics.plots import init_plot_settings
+from anemoi.training.diagnostics.plots import plot_graph_features
+from anemoi.training.diagnostics.plots import plot_histogram
+from anemoi.training.diagnostics.plots import plot_loss
+from anemoi.training.diagnostics.plots import plot_power_spectrum
+from anemoi.training.diagnostics.plots import plot_predicted_multilevel_flat_sample
+
+if TYPE_CHECKING:
+    import pytorch_lightning as pl
+    from omegaconf import DictConfig
+    from omegaconf import OmegaConf
+
+LOGGER = logging.getLogger(__name__)
+
+
+class ParallelExecutor(ThreadPoolExecutor):
+    """Wraps parallel execution and provides accurate information about errors.
+
+    Extends ThreadPoolExecutor to preserve the original traceback and line number.
+
+    Reference: https://stackoverflow.com/questions/19309514/getting-original-line-
+    number-for-exception-in-concurrent-futures/24457608#24457608
+    """
+
+    def submit(self, fn: Any, *args, **kwargs) -> Callable:
+        """Submits the wrapped function instead of `fn`."""
+        return super().submit(self._function_wrapper, fn, *args, **kwargs)
+
+    def _function_wrapper(self, fn: Any, *args: list, **kwargs: dict) -> Callable:
+        """Wraps `fn` in order to preserve the traceback of any kind of."""
+        try:
+            return fn(*args, **kwargs)
+        except Exception as exc:
+            raise sys.exc_info()[0](traceback.format_exc()) from exc
+
+
+class BasePlotCallback(Callback, ABC):
+    """Factory for creating a callback that plots data to Experiment Logging."""
+
+    def __init__(self, config: OmegaConf) -> None:
+        """Initialise the BasePlotCallback abstract base class.
+
+        Parameters
+        ----------
+        config : OmegaConf
+            Config object
+
+        """
+        super().__init__()
+        self.config = config
+        self.save_basedir = config.hardware.paths.plots
+        self.plot_frequency = config.diagnostics.plot.frequency
+        self.post_processors = None
+        self.pre_processors = None
+        self.latlons = None
+        init_plot_settings()
+
+        self.plot = self._plot
+        self._executor = None
+
+        if self.config.diagnostics.plot.asynchronous:
+            self._executor = ParallelExecutor(max_workers=1)
+            self._error: BaseException | None = None
+            self.plot = self._async_plot
+
+    @rank_zero_only
+    def _output_figure(
+        self,
+        logger: pl.loggers.base.LightningLoggerBase,
+        fig: plt.Figure,
+        epoch: int,
+        tag: str = "gnn",
+        exp_log_tag: str = "val_pred_sample",
+    ) -> None:
+        """Figure output: save to file and/or display in notebook."""
+        if self.save_basedir is not None:
+            save_path = Path(
+                self.save_basedir,
+                "plots",
+                f"{tag}_epoch{epoch:03d}.png",
+            )
+
+            save_path.parent.mkdir(parents=True, exist_ok=True)
+            fig.savefig(save_path, dpi=100, bbox_inches="tight")
+            if self.config.diagnostics.log.wandb.enabled:
+                import wandb
+
+                logger.experiment.log({exp_log_tag: wandb.Image(fig)})
+
+            if self.config.diagnostics.log.mlflow.enabled:
+                run_id = logger.run_id
+                logger.experiment.log_artifact(run_id, str(save_path))
+
+        plt.close(fig)  # cleanup
+
+    def teardown(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str) -> None:
+        """Method is called to close the threads."""
+        del trainer, pl_module, stage  # unused
+        if self._executor is not None:
+            self._executor.shutdown(wait=True)
+
+    @abstractmethod
+    @rank_zero_only
+    def _plot(
+        *args: list,
+        **kwargs: dict,
+    ) -> None: ...
+
+    @rank_zero_only
+    def _async_plot(
+        self,
+        trainer: pl.Trainer,
+        *args: list,
+        **kwargs: dict,
+    ) -> None:
+        """To execute the plot function but ensuring we catch any errors."""
+        future = self._executor.submit(
+            self._plot,
+            trainer,
+            *args,
+            **kwargs,
+        )
+        # otherwise the error won't be thrown till the validation epoch is finished
+        try:
+            future.result()
+        except Exception:
+            LOGGER.exception("Critical error occurred in asynchronous plots.")
+            sys.exit(1)
+
+
+class RolloutEval(Callback):
+    """Evaluates the model performance over a (longer) rollout window."""
+
+    def __init__(self, config: OmegaConf) -> None:
+        """Initialize RolloutEval callback.
+
+        Parameters
+        ----------
+        config : dict
+            Dictionary with configuration settings
+
+        """
+        super().__init__()
+
+        LOGGER.debug(
+            "Setting up RolloutEval callback with rollout = %d, frequency = %d ...",
+            config.diagnostics.eval.rollout,
+            config.diagnostics.eval.frequency,
+        )
+        self.rollout = config.diagnostics.eval.rollout
+        self.frequency = config.diagnostics.eval.frequency
+
+    def _eval(
+        self,
+        pl_module: pl.LightningModule,
+        batch: torch.Tensor,
+    ) -> None:
+        loss = torch.zeros(1, dtype=batch.dtype, device=pl_module.device, requires_grad=False)
+        metrics = {}
+
+        # start rollout
+        batch = pl_module.model.pre_processors(batch, in_place=False)
+        x = batch[
+            :,
+            0 : pl_module.multi_step,
+            ...,
+            pl_module.data_indices.internal_data.input.full,
+        ]  # (bs, multi_step, latlon, nvar)
+        assert (
+            batch.shape[1] >= self.rollout + pl_module.multi_step
+        ), "Batch length not sufficient for requested rollout length!"
+
+        with torch.no_grad():
+            for rollout_step in range(self.rollout):
+                y_pred = pl_module(x)  # prediction at rollout step rollout_step, shape = (bs, latlon, nvar)
+                y = batch[
+                    :,
+                    pl_module.multi_step + rollout_step,
+                    ...,
+                    pl_module.data_indices.internal_data.output.full,
+                ]  # target, shape = (bs, latlon, nvar)
+                # y includes the auxiliary variables, so we must leave those out when computing the loss
+                loss += pl_module.loss(y_pred, y)
+
+                x = pl_module.advance_input(x, y_pred, batch, rollout_step)
+
+                metrics_next, _ = pl_module.calculate_val_metrics(y_pred, y, rollout_step)
+                metrics.update(metrics_next)
+
+            # scale loss
+            loss *= 1.0 / self.rollout
+            self._log(pl_module, loss, metrics, batch.shape[0])
+
+    def _log(self, pl_module: pl.LightningModule, loss: torch.Tensor, metrics: dict, bs: int) -> None:
+        pl_module.log(
+            f"val_r{self.rollout}_wmse",
+            loss,
+            on_epoch=True,
+            on_step=True,
+            prog_bar=False,
+            logger=pl_module.logger_enabled,
+            batch_size=bs,
+            sync_dist=False,
+            rank_zero_only=True,
+        )
+        for mname, mvalue in metrics.items():
+            pl_module.log(
+                f"val_r{self.rollout}_" + mname,
+                mvalue,
+                on_epoch=True,
+                on_step=False,
+                prog_bar=False,
+                logger=pl_module.logger_enabled,
+                batch_size=bs,
+                sync_dist=False,
+                rank_zero_only=True,
+            )
+
+    @rank_zero_only
+    def on_validation_batch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        outputs: list,
+        batch: torch.Tensor,
+        batch_idx: int,
+    ) -> None:
+        del outputs  # outputs are not used
+        if batch_idx % self.frequency == 0:
+            precision_mapping = {
+                "16-mixed": torch.float16,
+                "bf16-mixed": torch.bfloat16,
+            }
+            prec = trainer.precision
+            dtype = precision_mapping.get(prec)
+            context = torch.autocast(device_type=batch.device.type, dtype=dtype) if dtype is not None else nullcontext()
+
+            with context:
+                self._eval(pl_module, batch)
+
+
+class LongRolloutPlots(BasePlotCallback):
+    """Evaluates the model performance over a (longer) rollout window."""
+
+    def __init__(self, config) -> None:
+        """Initialize RolloutEval callback.
+
+        Parameters
+        ----------
+        config : dict
+            Dictionary with configuration settings
+        """
+        super().__init__(config)
+
+        LOGGER.debug(
+            "Setting up callback for plots with long rollout: rollout = %d, frequency = every %d epoch ...",
+            config.diagnostics.plot.longrollout.rollout,
+            config.diagnostics.plot.longrollout.frequency,
+        )
+        self.rollout = config.diagnostics.plot.longrollout.rollout
+        self.eval_frequency = config.diagnostics.plot.longrollout.frequency
+        self.sample_idx = self.config.diagnostics.plot.sample_idx
+
+    @rank_zero_only
+    def _plot(
+        self,
+        trainer,
+        pl_module: pl.LightningModule,
+        batch: torch.Tensor,
+        batch_idx,
+        epoch,
+    ) -> None:
+
+        start_time = time.time()
+
+        logger = trainer.logger
+
+        # Build dictionary of inidicies and parameters to be plotted
+        plot_parameters_dict = {
+            pl_module.data_indices.model.output.name_to_index[name]: (
+                name,
+                name not in self.config.data.get("diagnostic", []),
+            )
+            for name in self.config.diagnostics.plot.parameters
+        }
+
+        if self.post_processors is None:
+            # Copy to be used across all the training cycle
+            self.post_processors = copy.deepcopy(pl_module.model.post_processors).cpu()
+        if self.latlons is None:
+            self.latlons = np.rad2deg(pl_module.latlons_data.clone().cpu().numpy())
+        local_rank = pl_module.local_rank
+
+        batch = pl_module.model.pre_processors(batch, in_place=False)
+        # prepare input tensor for rollout from preprocessed batch
+        x = batch[
+            :,
+            0 : pl_module.multi_step,
+            ...,
+            pl_module.data_indices.internal_data.input.full,
+        ]  # (bs, multi_step, latlon, nvar)
+        assert (
+            batch.shape[1] >= max(self.rollout) + pl_module.multi_step
+        ), "Batch length not sufficient for requested rollout length!"
+
+        # prepare input tensor for plotting
+        input_tensor_0 = batch[
+            self.sample_idx,
+            pl_module.multi_step - 1,
+            ...,
+            pl_module.data_indices.internal_data.output.full,
+        ].cpu()
+        data_0 = self.post_processors(input_tensor_0).numpy()
+
+        # start rollout
+        with torch.no_grad():
+            for rollout_step in range(max(self.rollout)):
+                y_pred = pl_module(x)  # prediction at rollout step rollout_step, shape = (bs, latlon, nvar)
+
+                x = pl_module.advance_input(x, y_pred, batch, rollout_step)
+
+                if (rollout_step + 1) in self.rollout:
+                    # prepare true output tensor for plotting
+                    input_tensor_rollout_step = batch[
+                        self.sample_idx,
+                        pl_module.multi_step + rollout_step,  # (pl_module.multi_step - 1) + (rollout_step + 1)
+                        ...,
+                        pl_module.data_indices.internal_data.output.full,
+                    ].cpu()
+                    data_rollout_step = self.post_processors(input_tensor_rollout_step).numpy()
+
+                    # prepare predicted output tensor for plotting
+                    output_tensor = self.post_processors(
+                        y_pred[self.sample_idx : self.sample_idx + 1, ...].cpu()
+                    ).numpy()
+
+                    fig = plot_predicted_multilevel_flat_sample(
+                        plot_parameters_dict,
+                        self.config.diagnostics.plot.per_sample,
+                        self.latlons,
+                        self.config.diagnostics.plot.get("accumulation_levels_plot", None),
+                        self.config.diagnostics.plot.get("cmap_accumulation", None),
+                        data_0.squeeze(),
+                        data_rollout_step.squeeze(),
+                        output_tensor[0, 0, :, :],  # rolloutstep, first member
+                        # force_global_view=self.show_entire_globe,
+                    )
+
+                    self._output_figure(
+                        logger,
+                        fig,
+                        epoch=epoch,
+                        tag=f"gnn_pred_val_sample_rstep{rollout_step:03d}_batch{batch_idx:04d}_rank0",
+                        exp_log_tag=f"val_pred_sample_rstep{rollout_step:03d}_rank{local_rank:01d}",
+                    )
+        LOGGER.info(f"Time taken to plot samples after longer rollout: {int(time.time() - start_time)} seconds")
+
+    @rank_zero_only
+    def on_validation_batch_end(self, trainer, pl_module, output, batch, batch_idx) -> None:
+        if (batch_idx) % self.plot_frequency == 0 and (trainer.current_epoch + 1) % self.eval_frequency == 0:
+            precision_mapping = {
+                "16-mixed": torch.float16,
+                "bf16-mixed": torch.bfloat16,
+            }
+            prec = trainer.precision
+            dtype = precision_mapping.get(prec)
+            context = torch.autocast(device_type=batch.device.type, dtype=dtype) if dtype is not None else nullcontext()
+
+            with context:
+                self._plot(trainer, pl_module, batch, batch_idx, epoch=trainer.current_epoch)
+
+
+class GraphTrainableFeaturesPlot(BasePlotCallback):
+    """Visualize the trainable features defined at the data and hidden graph nodes.
+
+    TODO: How best to visualize the learned edge embeddings? Offline, perhaps - using code from @Simon's notebook?
+    """
+
+    def __init__(self, config: OmegaConf) -> None:
+        """Initialise the GraphTrainableFeaturesPlot callback.
+
+        Parameters
+        ----------
+        config : OmegaConf
+            Config object
+
+        """
+        super().__init__(config)
+        self._graph_name_data = config.graph.data
+        self._graph_name_hidden = config.graph.hidden
+
+    @rank_zero_only
+    def _plot(
+        self,
+        trainer: pl.Trainer,
+        latlons: np.ndarray,
+        features: np.ndarray,
+        epoch: int,
+        tag: str,
+        exp_log_tag: str,
+    ) -> None:
+        fig = plot_graph_features(latlons, features)
+        self._output_figure(trainer.logger, fig, epoch=epoch, tag=tag, exp_log_tag=exp_log_tag)
+
+    @rank_zero_only
+    def on_validation_epoch_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+
+        model = pl_module.model.module.model if hasattr(pl_module.model, "module") else pl_module.model.model
+        graph = pl_module.graph_data.cpu().detach()
+        epoch = trainer.current_epoch
+
+        if model.trainable_data is not None:
+            data_coords = np.rad2deg(graph[(self._graph_name_data, "to", self._graph_name_data)].ecoords_rad.numpy())
+
+            self.plot(
+                trainer,
+                data_coords,
+                model.trainable_data.trainable.cpu().detach().numpy(),
+                epoch=epoch,
+                tag="trainable_data",
+                exp_log_tag="trainable_data",
+            )
+
+        if model.trainable_hidden is not None:
+            hidden_coords = np.rad2deg(
+                graph[(self._graph_name_hidden, "to", self._graph_name_hidden)].hcoords_rad.numpy(),
+            )
+
+            self.plot(
+                trainer,
+                hidden_coords,
+                model.trainable_hidden.trainable.cpu().detach().numpy(),
+                epoch=epoch,
+                tag="trainable_hidden",
+                exp_log_tag="trainable_hidden",
+            )
+
+
+class PlotLoss(BasePlotCallback):
+    """Plots the unsqueezed loss over rollouts."""
+
+    def __init__(self, config: OmegaConf) -> None:
+        """Initialise the PlotLoss callback.
+
+        Parameters
+        ----------
+        config : OmegaConf
+            Object with configuration settings
+
+        """
+        super().__init__(config)
+        self.parameter_names = None
+        self.parameter_groups = self.config.diagnostics.plot.parameter_groups
+        if self.parameter_groups is None:
+            self.parameter_groups = {}
+
+    @cached_property
+    def sort_and_color_by_parameter_group(self) -> tuple[np.ndarray, np.ndarray, dict, list]:
+        """Sort parameters by group and prepare colors."""
+
+        def automatically_determine_group(name: str) -> str:
+            # first prefix of parameter name is group name
+            parts = name.split("_")
+            return parts[0]
+
+        # group parameters by their determined group name for > 15 parameters
+        if len(self.parameter_names) <= 15:
+            # for <= 15 parameters, keep the full name of parameters
+            parameters_to_groups = np.array(self.parameter_names)
+            sort_by_parameter_group = np.arange(len(self.parameter_names), dtype=int)
+        else:
+            parameters_to_groups = np.array(
+                [
+                    next(
+                        (
+                            group_name
+                            for group_name, group_parameters in self.parameter_groups.items()
+                            if name in group_parameters
+                        ),
+                        automatically_determine_group(name),
+                    )
+                    for name in self.parameter_names
+                ],
+            )
+
+            unique_group_list, group_inverse, group_counts = np.unique(
+                parameters_to_groups,
+                return_inverse=True,
+                return_counts=True,
+            )
+
+            # join parameter groups that appear only once and are not given in config-file
+            unique_group_list = np.array(
+                [
+                    unique_group_list[tn] if count > 1 or unique_group_list[tn] in self.parameter_groups else "other"
+                    for tn, count in enumerate(group_counts)
+                ],
+            )
+            parameters_to_groups = unique_group_list[group_inverse]
+            unique_group_list, group_inverse = np.unique(parameters_to_groups, return_inverse=True)
+
+            # sort parameters by groups
+            sort_by_parameter_group = np.argsort(group_inverse, kind="stable")
+
+        # apply new order to parameters
+        sorted_parameter_names = np.array(self.parameter_names)[sort_by_parameter_group]
+        parameters_to_groups = parameters_to_groups[sort_by_parameter_group]
+        unique_group_list, group_inverse, group_counts = np.unique(
+            parameters_to_groups,
+            return_inverse=True,
+            return_counts=True,
+        )
+
+        # get a color per group and project to parameter list
+        cmap = "tab10" if len(unique_group_list) <= 10 else "tab20"
+        if len(unique_group_list) > 20:
+            LOGGER.warning("More than 20 groups detected, but colormap has only 20 colors.")
+        # if all groups have count 1 use black color
+        bar_color_per_group = (
+            np.tile("k", len(group_counts))
+            if not np.any(group_counts - 1)
+            else plt.get_cmap(cmap)(np.linspace(0, 1, len(unique_group_list)))
+        )
+
+        # set x-ticks
+        x_tick_positions = np.cumsum(group_counts) - group_counts / 2 - 0.5
+        xticks = dict(zip(unique_group_list, x_tick_positions))
+
+        legend_patches = []
+        for group_idx, group in enumerate(unique_group_list):
+            text_label = f"{group}: "
+            string_length = len(text_label)
+            for ii in np.where(group_inverse == group_idx)[0]:
+                text_label += sorted_parameter_names[ii] + ", "
+                string_length += len(sorted_parameter_names[ii]) + 2
+                if string_length > 50:
+                    # linebreak after 50 characters
+                    text_label += "\n"
+                    string_length = 0
+            legend_patches.append(mpatches.Patch(color=bar_color_per_group[group_idx], label=text_label[:-2]))
+
+        return sort_by_parameter_group, bar_color_per_group[group_inverse], xticks, legend_patches
+
+    @rank_zero_only
+    def _plot(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.Lightning_module,
+        outputs: list[torch.Tensor],
+        batch: torch.Tensor,
+        batch_idx: int,
+        epoch: int,
+    ) -> None:
+        logger = trainer.logger
+
+        parameter_names = list(pl_module.data_indices.internal_model.output.name_to_index.keys())
+        parameter_positions = list(pl_module.data_indices.internal_model.output.name_to_index.values())
+        # reorder parameter_names by position
+        self.parameter_names = [parameter_names[i] for i in np.argsort(parameter_positions)]
+
+        batch = pl_module.model.pre_processors(batch, in_place=False)
+        for rollout_step in range(pl_module.rollout):
+            y_hat = outputs[1][rollout_step]
+            y_true = batch[
+                :, pl_module.multi_step + rollout_step, ..., pl_module.data_indices.internal_data.output.full
+            ]
+            loss = pl_module.loss(y_hat, y_true, squash=False).cpu().numpy()
+
+            sort_by_parameter_group, colors, xticks, legend_patches = self.sort_and_color_by_parameter_group
+            fig = plot_loss(loss[sort_by_parameter_group], colors, xticks, legend_patches)
+
+            self._output_figure(
+                logger,
+                fig,
+                epoch=epoch,
+                tag=f"loss_rstep_rstep{rollout_step:02d}_rank{pl_module.local_rank:01d}",
+                exp_log_tag=f"loss_sample_rstep{rollout_step:02d}_rank{pl_module.local_rank:01d}",
+            )
+
+    def on_validation_batch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        outputs: list[torch.Tensor],
+        batch: torch.Tensor,
+        batch_idx: int,
+    ) -> None:
+        if batch_idx % self.plot_frequency == 0:
+            self.plot(trainer, pl_module, outputs, batch, batch_idx, epoch=trainer.current_epoch)
+
+
+class PlotSample(BasePlotCallback):
+    """Plots a post-processed sample: input, target and prediction."""
+
+    def __init__(self, config: OmegaConf) -> None:
+        """Initialise the PlotSample callback.
+
+        Parameters
+        ----------
+        config : OmegaConf
+            Config object
+
+        """
+        super().__init__(config)
+        self.sample_idx = self.config.diagnostics.plot.sample_idx
+        self.precip_and_related_fields = self.config.diagnostics.plot.precip_and_related_fields
+        LOGGER.info(f"Using defined accumulation colormap for fields: {self.precip_and_related_fields}")
+
+    @rank_zero_only
+    def _plot(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.Lightning_module,
+        outputs: list[torch.Tensor],
+        batch: torch.Tensor,
+        batch_idx: int,
+        epoch: int,
+    ) -> None:
+        logger = trainer.logger
+
+        # Build dictionary of indices and parameters to be plotted
+        diagnostics = [] if self.config.data.diagnostic is None else self.config.data.diagnostic
+        plot_parameters_dict = {
+            pl_module.data_indices.model.output.name_to_index[name]: (name, name not in diagnostics)
+            for name in self.config.diagnostics.plot.parameters
+        }
+
+        # When running in Async mode, it might happen that in the last epoch these tensors
+        # have been moved to the cpu (and then the denormalising would fail as the 'input_tensor' would be on CUDA
+        # but internal ones would be on the cpu), The lines below allow to address this problem
+        if self.post_processors is None:
+            # Copy to be used across all the training cycle
+            self.post_processors = copy.deepcopy(pl_module.model.post_processors).cpu()
+        if self.latlons is None:
+            self.latlons = np.rad2deg(pl_module.latlons_data.clone().cpu().numpy())
+        local_rank = pl_module.local_rank
+
+        batch = pl_module.model.pre_processors(batch, in_place=False)
+        input_tensor = batch[
+            self.sample_idx,
+            pl_module.multi_step - 1 : pl_module.multi_step + pl_module.rollout + 1,
+            ...,
+            pl_module.data_indices.internal_data.output.full,
+        ].cpu()
+        data = self.post_processors(input_tensor).numpy()
+
+        output_tensor = self.post_processors(
+            torch.cat(tuple(x[self.sample_idx : self.sample_idx + 1, ...].cpu() for x in outputs[1])),
+            in_place=False,
+        ).numpy()
+
+        for rollout_step in range(pl_module.rollout):
+            fig = plot_predicted_multilevel_flat_sample(
+                plot_parameters_dict,
+                self.config.diagnostics.plot.per_sample,
+                self.latlons,
+                self.config.diagnostics.plot.accumulation_levels_plot,
+                self.config.diagnostics.plot.cmap_accumulation,
+                data[0, ...].squeeze(),
+                data[rollout_step + 1, ...].squeeze(),
+                output_tensor[rollout_step, ...],
+                precip_and_related_fields=self.precip_and_related_fields,
+            )
+
+            self._output_figure(
+                logger,
+                fig,
+                epoch=epoch,
+                tag=f"gnn_pred_val_sample_rstep{rollout_step:02d}_batch{batch_idx:04d}_rank0",
+                exp_log_tag=f"val_pred_sample_rstep{rollout_step:02d}_rank{local_rank:01d}",
+            )
+
+    def on_validation_batch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.Lightning_module,
+        outputs: list[torch.Tensor],
+        batch: torch.Tensor,
+        batch_idx: int,
+    ) -> None:
+        if batch_idx % self.plot_frequency == 0:
+            self.plot(trainer, pl_module, outputs, batch, batch_idx, epoch=trainer.current_epoch)
+
+
+class PlotAdditionalMetrics(BasePlotCallback):
+    """Plots TP related metric comparing target and prediction.
+
+    The actual increment (output - input) is plot for prognostic variables while the output is plot for diagnostic ones.
+
+    - Power Spectrum
+    - Histograms
+    """
+
+    def __init__(self, config: OmegaConf) -> None:
+        """Initialise the PlotAdditionalMetrics callback.
+
+        Parameters
+        ----------
+        config : OmegaConf
+            Config object
+
+        """
+        super().__init__(config)
+        self.sample_idx = self.config.diagnostics.plot.sample_idx
+        self.precip_and_related_fields = self.config.diagnostics.plot.precip_and_related_fields
+        LOGGER.info(f"Using precip histogram plotting method for fields: {self.precip_and_related_fields}")
+
+    @rank_zero_only
+    def _plot(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        outputs: list,
+        batch: torch.Tensor,
+        batch_idx: int,
+        epoch: int,
+    ) -> None:
+        logger = trainer.logger
+
+        # When running in Async mode, it might happen that in the last epoch these tensors
+        # have been moved to the cpu (and then the denormalising would fail as the 'input_tensor' would be on CUDA
+        # but internal ones would be on the cpu), The lines below allow to address this problem
+        if self.pre_processors is None:
+            # Copy to be used across all the training cycle
+            self.pre_processors = copy.deepcopy(pl_module.model.pre_processors).cpu()
+        if self.post_processors is None:
+            # Copy to be used across all the training cycle
+            self.post_processors = copy.deepcopy(pl_module.model.post_processors).cpu()
+        if self.latlons is None:
+            self.latlons = np.rad2deg(pl_module.latlons_data.clone().cpu().numpy())
+        local_rank = pl_module.local_rank
+        batch = pl_module.model.pre_processors(batch, in_place=False)
+        input_tensor = batch[
+            self.sample_idx,
+            pl_module.multi_step - 1 : pl_module.multi_step + pl_module.rollout + 1,
+            ...,
+            pl_module.data_indices.internal_data.output.full,
+        ].cpu()
+        data = self.post_processors(input_tensor).numpy()
+        output_tensor = self.post_processors(
+            torch.cat(tuple(x[self.sample_idx : self.sample_idx + 1, ...].cpu() for x in outputs[1])),
+            in_place=False,
+        ).numpy()
+
+        for rollout_step in range(pl_module.rollout):
+            if self.config.diagnostics.plot.parameters_histogram is not None:
+                # Build dictionary of inidicies and parameters to be plotted
+
+                diagnostics = [] if self.config.data.diagnostic is None else self.config.data.diagnostic
+                plot_parameters_dict_histogram = {
+                    pl_module.data_indices.model.output.name_to_index[name]: (name, name not in diagnostics)
+                    for name in self.config.diagnostics.plot.parameters_histogram
+                }
+
+                fig = plot_histogram(
+                    plot_parameters_dict_histogram,
+                    data[0, ...].squeeze(),
+                    data[rollout_step + 1, ...].squeeze(),
+                    output_tensor[rollout_step, ...],
+                    precip_and_related_fields=self.precip_and_related_fields,
+                )
+
+                self._output_figure(
+                    logger,
+                    fig,
+                    epoch=epoch,
+                    tag=f"gnn_pred_val_histo_rstep_{rollout_step:02d}_batch{batch_idx:04d}_rank0",
+                    exp_log_tag=f"val_pred_histo_rstep_{rollout_step:02d}_rank{local_rank:01d}",
+                )
+
+            if self.config.diagnostics.plot.parameters_spectrum is not None:
+                # Build dictionary of inidicies and parameters to be plotted
+                diagnostics = [] if self.config.data.diagnostic is None else self.config.data.diagnostic
+
+                plot_parameters_dict_spectrum = {
+                    pl_module.data_indices.model.output.name_to_index[name]: (name, name not in diagnostics)
+                    for name in self.config.diagnostics.plot.parameters_spectrum
+                }
+
+                fig = plot_power_spectrum(
+                    plot_parameters_dict_spectrum,
+                    self.latlons,
+                    data[0, ...].squeeze(),
+                    data[rollout_step + 1, ...].squeeze(),
+                    output_tensor[rollout_step, ...],
+                )
+
+                self._output_figure(
+                    logger,
+                    fig,
+                    epoch=epoch,
+                    tag=f"gnn_pred_val_spec_rstep_{rollout_step:02d}_batch{batch_idx:04d}_rank0",
+                    exp_log_tag=f"val_pred_spec_rstep_{rollout_step:02d}_rank{local_rank:01d}",
+                )
+
+    def on_validation_batch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        outputs: list[torch.Tensor],
+        batch: torch.Tensor,
+        batch_idx: int,
+    ) -> None:
+        if batch_idx % self.plot_frequency == 0:
+            self.plot(trainer, pl_module, outputs, batch, batch_idx, epoch=trainer.current_epoch)
+
+
+class ParentUUIDCallback(Callback):
+    """A callback that retrieves the parent UUID for a model, if it is a child model."""
+
+    def __init__(self, config: OmegaConf) -> None:
+        """Initialise the ParentUUIDCallback callback.
+
+        Parameters
+        ----------
+        config : OmegaConf
+            Config object
+
+        """
+        super().__init__()
+        self.config = config
+
+    def on_load_checkpoint(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        checkpoint: torch.nn.Module,
+    ) -> None:
+        del trainer  # unused
+        pl_module.hparams["metadata"]["parent_uuid"] = checkpoint["hyper_parameters"]["metadata"]["uuid"]
+
+
+class AnemoiCheckpoint(ModelCheckpoint):
+    """A checkpoint callback that saves the model after every validation epoch."""
+
+    def __init__(self, config: OmegaConf, **kwargs: dict) -> None:
+        """Initialise the AnemoiCheckpoint callback.
+
+        Parameters
+        ----------
+        config : OmegaConf
+            Config object
+        kwargs : dict
+            Additional keyword arguments for Pytorch ModelCheckpoint
+
+        """
+        super().__init__(**kwargs)
+        self.config = config
+        self.start = time.time()
+        self._model_metadata = None
+        self._tracker_metadata = None
+        self._tracker_name = None
+
+    @staticmethod
+    def _torch_drop_down(trainer: pl.Trainer) -> torch.nn.Module:
+        # Get the model from the DataParallel wrapper, for single and multi-gpu cases
+        assert hasattr(trainer, "model"), "Trainer has no attribute 'model'! Is the Pytorch Lightning version correct?"
+        return trainer.model.module.model if hasattr(trainer.model, "module") else trainer.model.model
+
+    @rank_zero_only
+    def model_metadata(self, model: torch.nn.Module) -> dict:
+        if self._model_metadata is not None:
+            return self._model_metadata
+
+        self._model_metadata = {
+            "model": model.__class__.__name__,
+            "trainable_parameters": sum(p.numel() for p in model.parameters() if p.requires_grad),
+            "total_parameters": sum(p.numel() for p in model.parameters()),
+            "summary": repr(
+                torchinfo.summary(
+                    model,
+                    depth=50,
+                    verbose=0,
+                    row_settings=["var_names"],
+                ),
+            ),
+        }
+
+        return self._model_metadata
+
+    def tracker_metadata(self, trainer: pl.Trainer) -> dict:
+        if self._tracker_metadata is not None:
+            return {self._tracker_name: self._tracker_metadata}
+
+        if self.config.diagnostics.log.wandb.enabled:
+            self._tracker_name = "wand"
+            import wandb
+
+            run = wandb.run
+            if run is not None:
+                self._tracker_metadata = {
+                    "id": run.id,
+                    "name": run.name,
+                    "url": run.url,
+                    "project": run.project,
+                }
+            return {self._tracker_name: self._tracker_metadata}
+
+        if self.config.diagnostics.log.mlflow.enabled:
+            self._tracker_name = "mlflow"
+
+            from anemoi.training.diagnostics.mlflow.logger import AnemoiMLflowLogger
+
+            mlflow_logger = next(logger for logger in trainer.loggers if isinstance(logger, AnemoiMLflowLogger))
+            run_id = mlflow_logger.run_id
+            run = mlflow_logger._mlflow_client.get_run(run_id)
+
+            if run is not None:
+                self._tracker_metadata = {
+                    "id": run.info.run_id,
+                    "name": run.info.run_name,
+                    "url": run.info.artifact_uri,
+                    "project": run.info.experiment_id,
+                }
+            return {self._tracker_name: self._tracker_metadata}
+
+        return {}
+
+    def _remove_checkpoint(self, trainer: "pl.Trainer", filepath: str) -> None:
+        """Calls the strategy to remove the checkpoint file."""
+        super()._remove_checkpoint(trainer, filepath)
+        trainer.strategy.remove_checkpoint(self._get_inference_checkpoint_filepath(filepath))
+
+    def _get_inference_checkpoint_filepath(self, filepath: str) -> str:
+        """Defines the filepath for the inference checkpoint."""
+        return Path(filepath).parent / Path("inference-" + str(Path(filepath).name))
+
+    def _save_checkpoint(self, trainer: pl.Trainer, lightning_checkpoint_filepath: str) -> None:
+        if trainer.is_global_zero:
+            model = self._torch_drop_down(trainer)
+
+            # We want a different uuid each time we save the model
+            # so we can tell them apart in the catalogue (i.e. different epochs)
+            checkpoint_uuid = str(uuid.uuid4())
+            trainer.lightning_module._hparams["metadata"]["uuid"] = checkpoint_uuid
+
+            trainer.lightning_module._hparams["metadata"]["model"] = self.model_metadata(model)
+            trainer.lightning_module._hparams["metadata"]["tracker"] = self.tracker_metadata(trainer)
+
+            trainer.lightning_module._hparams["metadata"]["training"] = {
+                "current_epoch": trainer.current_epoch,
+                "global_step": trainer.global_step,
+                "elapsed_time": time.time() - self.start,
+            }
+
+            Path(lightning_checkpoint_filepath).parent.mkdir(parents=True, exist_ok=True)
+
+            save_config = model.config
+            model.config = None
+
+            tmp_metadata = model.metadata
+            model.metadata = None
+
+            metadata = dict(**tmp_metadata)
+
+            inference_checkpoint_filepath = self._get_inference_checkpoint_filepath(lightning_checkpoint_filepath)
+
+            torch.save(model, inference_checkpoint_filepath)
+
+            save_metadata(inference_checkpoint_filepath, metadata)
+
+            model.config = save_config
+            model.metadata = tmp_metadata
+
+            self._last_global_step_saved = trainer.global_step
+
+        trainer.strategy.barrier()
+
+        # saving checkpoint used for pytorch-lightning based training
+        trainer.save_checkpoint(lightning_checkpoint_filepath, self.save_weights_only)
+
+        self._last_global_step_saved = trainer.global_step
+        self._last_checkpoint_saved = lightning_checkpoint_filepath
+
+        if trainer.is_global_zero:
+            from weakref import proxy
+
+            # save metadata for the training checkpoint in the same format as inference
+            save_metadata(lightning_checkpoint_filepath, metadata)
+
+            # notify loggers
+            for logger in trainer.loggers:
+                logger.after_save_checkpoint(proxy(self))
+
+
+def get_callbacks(config: DictConfig) -> list:  # noqa: C901
+    """Setup callbacks for PyTorch Lightning trainer.
+
+    Parameters
+    ----------
+    config : DictConfig
+        Job configuration
+
+    Returns
+    -------
+    List
+        A list of PyTorch Lightning callbacks
+
+    """
+    checkpoint_settings = {
+        "dirpath": config.hardware.paths.checkpoints,
+        "verbose": False,
+        # save weights, optimizer states, LR-schedule states, hyperparameters etc.
+        # https://pytorch-lightning.readthedocs.io/en/stable/common/checkpointing_basic.html#contents-of-a-checkpoint
+        "save_weights_only": False,
+        "auto_insert_metric_name": False,
+        # save after every validation epoch, if we've improved
+        "save_on_train_epoch_end": False,
+        "enable_version_counter": False,
+    }
+
+    ckpt_frequency_save_dict = {}
+    for key, frequency_dict in config.diagnostics.checkpoint.items():
+        frequency = frequency_dict["save_frequency"]
+        n_saved = frequency_dict["num_models_saved"]
+        if key == "every_n_minutes" and frequency_dict["save_frequency"] is not None:
+            target = "train_time_interval"
+            frequency = timedelta(minutes=frequency_dict["save_frequency"])
+        else:
+            target = key
+        ckpt_frequency_save_dict[target] = (config.hardware.files.checkpoint[key], frequency, n_saved)
+
+    trainer_callbacks = []
+    if not config.diagnostics.profiler:
+        for save_key, (name, save_frequency, save_n_models) in ckpt_frequency_save_dict.items():
+            if save_frequency is not None:
+                LOGGER.debug("Checkpoint callback at %s = %s ...", save_key, save_frequency)
+                trainer_callbacks.extend(
+                    # save_top_k: the save_top_k flag can either save the best or the last k checkpoints
+                    # depending on the monitor flag on ModelCheckpoint.
+                    # See https://lightning.ai/docs/pytorch/stable/common/checkpointing_intermediate.html for reference
+                    [
+                        AnemoiCheckpoint(
+                            config=config,
+                            filename=name,
+                            save_last=True,
+                            **{save_key: save_frequency},
+                            # if save_top_k == k, last k models saved; if save_top_k == -1, all models are saved
+                            save_top_k=save_n_models,
+                            monitor="step",
+                            mode="max",
+                            **checkpoint_settings,
+                        ),
+                    ],
+                )
+            else:
+                LOGGER.debug("Not setting up a checkpoint callback with %s", save_key)
+    else:
+        # the tensorboard logger + pytorch profiler cause pickling errors when writing checkpoints
+        LOGGER.warning("Profiling is enabled - will not write any training or inference model checkpoints!")
+
+    if any([config.diagnostics.log.wandb.enabled, config.diagnostics.log.mlflow.enabled]):
+        from pytorch_lightning.callbacks import LearningRateMonitor
+
+        trainer_callbacks.append(
+            LearningRateMonitor(
+                logging_interval="step",
+                log_momentum=False,
+            ),
+        )
+
+    if config.diagnostics.eval.enabled:
+        trainer_callbacks.append(RolloutEval(config))
+
+    if config.diagnostics.plot.enabled:
+        trainer_callbacks.extend(
+            [
+                PlotLoss(config),
+                PlotSample(config),
+            ],
+        )
+        if (config.diagnostics.plot.parameters_histogram or config.diagnostics.plot.parameters_spectrum) is not None:
+            trainer_callbacks.extend([PlotAdditionalMetrics(config)])
+        if config.diagnostics.plot.get("longrollout") and config.diagnostics.plot.longrollout.enabled:
+            trainer_callbacks.extend([LongRolloutPlots(config)])
+
+    if config.training.swa.enabled:
+        from pytorch_lightning.callbacks.stochastic_weight_avg import StochasticWeightAveraging
+
+        trainer_callbacks.append(
+            StochasticWeightAveraging(
+                swa_lrs=config.training.swa.lr,
+                swa_epoch_start=min(
+                    int(0.75 * config.training.max_epochs),
+                    config.training.max_epochs - 1,
+                ),
+                annealing_epochs=max(int(0.25 * config.training.max_epochs), 1),
+                annealing_strategy="cos",
+                device=None,
+            ),
+        )
+
+    trainer_callbacks.append(ParentUUIDCallback(config))
+
+    if config.diagnostics.plot.learned_features:
+        LOGGER.debug("Setting up a callback to plot the trainable graph node features ...")
+        trainer_callbacks.append(GraphTrainableFeaturesPlot(config))
+    return trainer_callbacks
diff --git a/src/anemoi/training/diagnostics/continents.json b/src/anemoi/training/diagnostics/continents.json
new file mode 100644
index 00000000..96c46d4d
--- /dev/null
+++ b/src/anemoi/training/diagnostics/continents.json
@@ -0,0 +1,50 @@
+{"type":"FeatureCollection", "features":[
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[141,-2.6],[142.74,-3.29],[144.58,-3.86],[145.27,-4.37],[145.83,-4.88],[145.98,-5.47],[147.65,-6.08],[147.89,-6.61],[146.97,-6.72],[147.19,-7.39],[148.08,-8.04],[148.73,-9.1],[149.31,-9.07],[149.27,-9.51],[150.04,-9.68],[149.74,-9.87],[150.8,-10.29],[150.69,-10.58],[150.03,-10.65],[149.78,-10.39],[148.92,-10.28],[147.91,-10.13],[147.14,-9.49],[146.57,-8.94],[146.05,-8.07],[144.74,-7.63],[143.9,-7.92],[143.29,-8.25],[143.41,-8.98],[142.63,-9.33],[142.07,-9.16],[141.03,-9.12],[140.14,-8.3],[139.13,-8.1],[138.88,-8.38],[137.61,-8.41],[138.04,-7.6],[138.67,-7.32],[138.41,-6.23],[137.93,-5.39],[135.99,-4.55],[135.16,-4.46],[133.66,-3.54],[133.37,-4.02],[132.98,-4.11],[132.76,-3.75],[132.75,-3.31],[131.99,-2.82],[133.07,-2.46],[133.78,-2.48],[133.7,-2.21],[132.23,-2.21],[131.84,-1.62],[130.94,-1.43],[130.52,-0.94],[131.87,-0.7],[132.38,-0.37],[133.99,-0.78],[134.14,-1.15],[134.42,-2.77],[135.46,-3.37],[136.29,-2.31],[137.44,-1.7],[138.33,-1.7],[139.18,-2.05],[139.93,-2.41],[141,-2.6]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[114.2,4.53],[114.6,4.9],[115.45,5.45],[116.22,6.14],[116.73,6.92],[117.13,6.93],[117.64,6.42],[117.69,5.99],[118.35,5.71],[119.18,5.41],[119.11,5.02],[118.44,4.97],[118.62,4.48],[117.88,4.14],[117.31,3.23],[118.05,2.29],[117.88,1.83],[119,0.9],[117.81,0.78],[117.48,0.1],[117.52,-0.8],[116.56,-1.49],[116.53,-2.48],[116.15,-4.01],[116,-3.66],[114.86,-4.11],[114.47,-3.5],[113.76,-3.44],[113.26,-3.12],[112.07,-3.48],[111.7,-2.99],[111.05,-3.05],[110.22,-2.93],[110.07,-1.59],[109.57,-1.31],[109.09,-0.46],[108.95,0.42],[109.07,1.34],[109.66,2.01],[110.4,1.66],[111.17,1.85],[111.37,2.7],[111.8,2.89],[113,3.1],[113.71,3.89],[114.2,4.53]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-88.15,74.39],[-89.76,74.52],[-92.42,74.84],[-92.77,75.39],[-92.89,75.88],[-93.89,76.32],[-95.96,76.44],[-97.12,76.75],[-96.75,77.16],[-94.68,77.1],[-93.57,76.78],[-91.61,76.78],[-90.74,76.45],[-90.97,76.07],[-89.82,75.85],[-89.19,75.61],[-87.84,75.57],[-86.38,75.48],[-84.79,75.7],[-82.75,75.78],[-81.13,75.71],[-80.06,75.34],[-79.83,74.92],[-80.46,74.66],[-81.95,74.44],[-83.23,74.56],[-86.1,74.41],[-88.15,74.39]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-82.27,23.19],[-81.4,23.12],[-80.62,23.11],[-79.68,22.77],[-79.28,22.4],[-78.35,22.51],[-77.99,22.28],[-77.15,21.66],[-76.52,21.21],[-76.19,21.22],[-75.6,21.02],[-75.67,20.74],[-74.93,20.69],[-74.18,20.28],[-74.3,20.05],[-74.96,19.92],[-75.63,19.87],[-76.32,19.95],[-77.76,19.86],[-77.09,20.41],[-77.49,20.67],[-78.14,20.74],[-78.48,21.03],[-78.72,21.6],[-79.28,21.56],[-80.22,21.83],[-80.52,22.04],[-81.82,22.19],[-82.17,22.39],[-81.8,22.64],[-82.78,22.69],[-83.49,22.17],[-83.91,22.15],[-84.05,21.91],[-84.55,21.8],[-84.97,21.9],[-84.45,22.2],[-84.23,22.57],[-83.78,22.79],[-83.27,22.98],[-82.51,23.08],[-82.27,23.19]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-55.6,51.32],[-56.13,50.69],[-56.8,49.81],[-56.14,50.15],[-55.47,49.94],[-55.82,49.59],[-54.94,49.31],[-54.47,49.56],[-53.48,49.25],[-53.79,48.52],[-53.09,48.69],[-52.96,48.16],[-52.65,47.54],[-53.07,46.66],[-53.52,46.62],[-54.18,46.81],[-53.96,47.63],[-54.24,47.75],[-55.4,46.88],[-56,46.92],[-55.29,47.39],[-56.25,47.63],[-57.33,47.57],[-59.27,47.6],[-59.42,47.9],[-58.8,48.25],[-59.23,48.52],[-58.39,49.13],[-57.36,50.72],[-56.74,51.29],[-55.87,51.63],[-55.41,51.59],[-55.6,51.32]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-83.88,65.11],[-82.79,64.77],[-81.64,64.46],[-81.55,63.98],[-80.82,64.06],[-80.1,63.73],[-80.99,63.41],[-82.55,63.65],[-83.11,64.1],[-84.1,63.57],[-85.52,63.05],[-85.87,63.64],[-87.22,63.54],[-86.35,64.04],[-86.22,64.82],[-85.88,65.74],[-85.16,65.66],[-84.98,65.22],[-84.46,65.37],[-83.88,65.11]]},"properties":{"scalerank":1,"featurecla":"Coastline","min_zoom":0.5}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-78.77,72.35],[-77.82,72.75],[-75.61,72.24],[-74.23,71.77],[-74.1,71.33],[-72.24,71.56],[-71.2,70.92],[-68.79,70.53],[-67.91,70.12],[-66.97,69.19],[-68.81,68.72],[-66.45,68.07],[-64.86,67.85],[-63.42,66.93],[-61.85,66.86],[-62.16,66.16],[-63.92,65],[-65.15,65.43],[-66.72,66.39],[-68.02,66.26],[-68.14,65.69],[-67.09,65.11],[-65.73,64.65],[-65.32,64.38],[-64.67,63.39],[-65.01,62.67],[-66.28,62.95],[-68.78,63.75],[-67.37,62.88],[-66.33,62.28],[-66.17,61.93],[-68.88,62.33],[-71.02,62.91],[-72.24,63.4],[-71.89,63.68],[-73.38,64.19],[-74.83,64.68],[-74.82,64.39],[-77.71,64.23],[-78.56,64.57],[-77.9,65.31],[-76.02,65.33],[-73.96,65.45],[-74.29,65.81],[-73.94,66.31],[-72.65,67.28],[-72.93,67.73],[-73.31,68.07],[-74.84,68.55],[-76.87,68.89],[-76.23,69.15],[-77.29,69.77],[-78.17,69.83],[-78.96,70.17],[-79.49,69.87],[-81.31,69.74],[-84.94,69.97],[-87.06,70.26],[-88.68,70.41],[-89.51,70.76],[-88.47,71.22],[-89.89,71.22],[-90.21,72.24],[-89.44,73.13],[-88.41,73.54],[-85.83,73.8],[-86.56,73.16],[-85.77,72.53],[-84.85,73.34],[-82.32,73.75],[-80.6,72.72],[-80.75,72.06],[-78.77,72.35]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-100.44,72.71],[-101.54,73.36],[-100.36,73.84],[-99.16,73.63],[-97.38,73.76],[-97.12,73.47],[-98.05,72.99],[-96.54,72.56],[-96.72,71.66],[-98.36,71.27],[-99.32,71.36],[-100.01,71.74],[-102.48,72.48],[-102.48,72.83],[-100.44,72.71]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-107.82,75.85],[-106.93,76.01],[-105.88,75.97],[-105.7,75.48],[-106.31,75.01],[-109.7,74.85],[-112.22,74.42],[-113.74,74.39],[-113.87,74.72],[-111.79,75.16],[-116.31,75.04],[-117.71,75.22],[-116.35,76.2],[-115.4,76.48],[-112.59,76.14],[-110.81,75.55],[-109.07,75.47],[-110.5,76.43],[-109.58,76.79],[-108.55,76.68],[-108.21,76.2],[-107.82,75.85]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-121.54,74.45],[-120.11,74.24],[-117.56,74.19],[-116.58,73.9],[-115.51,73.48],[-116.77,73.22],[-119.22,72.52],[-120.46,71.82],[-120.46,71.38],[-123.09,70.9],[-123.62,71.34],[-125.93,71.87],[-125.5,72.29],[-124.81,73.02],[-123.94,73.68],[-124.92,74.29],[-121.54,74.45]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-125.42,49.95],[-124.92,49.48],[-123.92,49.06],[-123.51,48.51],[-124.01,48.37],[-125.66,48.83],[-125.95,49.18],[-126.85,49.53],[-127.03,49.81],[-128.06,49.99],[-128.44,50.54],[-128.36,50.77],[-127.31,50.55],[-126.7,50.4],[-125.76,50.3],[-125.42,49.95]]},"properties":{"scalerank":1,"featurecla":"Coastline","min_zoom":1}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[49.54,-12.47],[49.81,-12.9],[50.06,-13.56],[50.22,-14.76],[50.48,-15.23],[50.38,-15.71],[50.2,-16],[49.86,-15.41],[49.67,-15.71],[49.86,-16.45],[49.77,-16.88],[49.5,-17.11],[49.44,-17.95],[49.04,-19.12],[48.55,-20.5],[47.93,-22.39],[47.55,-23.78],[47.1,-24.94],[46.28,-25.18],[45.41,-25.6],[44.83,-25.35],[44.04,-24.99],[43.76,-24.46],[43.7,-23.57],[43.35,-22.78],[43.25,-22.06],[43.43,-21.34],[43.89,-21.16],[43.9,-20.83],[44.37,-20.07],[44.46,-19.44],[44.23,-18.96],[44.04,-18.33],[43.96,-17.41],[44.31,-16.85],[44.45,-16.22],[44.94,-16.18],[45.5,-15.97],[45.87,-15.79],[46.31,-15.78],[46.88,-15.21],[47.71,-14.59],[48.01,-14.09],[47.87,-13.66],[48.29,-13.78],[48.85,-13.09],[48.86,-12.49],[49.19,-12.04],[49.54,-12.47]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-48.66,-78.05],[-48.15,-78.05],[-46.66,-77.83],[-45.15,-78.05],[-43.92,-78.48],[-43.49,-79.09],[-43.37,-79.52],[-43.33,-80.03],[-44.88,-80.34],[-46.51,-80.59],[-48.39,-80.83],[-50.48,-81.03],[-52.85,-80.97],[-54.16,-80.63],[-53.99,-80.22],[-51.85,-79.95],[-50.99,-79.61],[-50.36,-79.18],[-49.91,-78.81],[-49.31,-78.46],[-48.66,-78.05]]},"properties":{"scalerank":1,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-73.92,-71.27],[-73.23,-71.15],[-72.07,-71.19],[-71.78,-70.68],[-71.72,-70.31],[-71.74,-69.51],[-71.17,-69.04],[-70.25,-68.88],[-69.72,-69.25],[-69.49,-69.62],[-69.06,-70.07],[-68.73,-70.51],[-68.45,-70.96],[-68.33,-71.41],[-68.51,-71.8],[-68.78,-72.17],[-69.96,-72.31],[-71.08,-72.5],[-72.39,-72.48],[-71.9,-72.09],[-73.07,-72.23],[-74.19,-72.37],[-74.95,-72.07],[-75.01,-71.66],[-73.92,-71.27]]},"properties":{"scalerank":1,"featurecla":"Coastline","min_zoom":1}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[151.3,-5.84],[150.75,-6.08],[150.24,-6.32],[149.71,-6.32],[148.89,-6.03],[148.32,-5.75],[148.4,-5.44],[149.3,-5.58],[149.85,-5.51],[150,-5.03],[150.14,-5],[150.24,-5.53],[150.81,-5.46],[151.09,-5.11],[151.65,-4.76],[151.54,-4.17],[152.14,-4.15],[152.34,-4.31],[152.32,-4.87],[151.98,-5.48],[151.46,-5.56],[151.3,-5.84]]},"properties":{"scalerank":1,"featurecla":"Coastline","min_zoom":1}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[176.89,-40.07],[176.51,-40.6],[176.01,-41.29],[175.24,-41.69],[175.07,-41.43],[174.65,-41.28],[175.23,-40.46],[174.9,-39.91],[173.82,-39.51],[173.85,-39.15],[174.57,-38.8],[174.74,-38.03],[174.7,-37.38],[174.29,-36.71],[174.32,-36.53],[173.84,-36.12],[173.05,-35.24],[172.64,-34.53],[173.01,-34.45],[173.55,-35.01],[174.33,-35.27],[174.61,-36.16],[175.34,-37.21],[175.36,-36.53],[175.81,-36.8],[175.96,-37.56],[176.76,-37.88],[177.44,-37.96],[178.01,-37.58],[178.52,-37.7],[178.27,-38.58],[177.97,-39.17],[177.21,-39.15],[176.94,-39.45],[177.03,-39.88],[176.89,-40.07]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[169.67,-43.56],[170.52,-43.03],[171.13,-42.51],[171.57,-41.77],[171.95,-41.51],[172.1,-40.96],[172.8,-40.49],[173.02,-40.92],[173.25,-41.33],[173.96,-40.93],[174.25,-41.35],[174.25,-41.77],[173.88,-42.23],[173.22,-42.97],[172.71,-43.37],[173.08,-43.85],[172.31,-43.87],[171.45,-44.24],[171.19,-44.9],[170.62,-45.91],[169.83,-46.36],[169.33,-46.64],[168.41,-46.62],[167.76,-46.29],[166.68,-46.22],[166.51,-45.85],[167.05,-45.11],[168.3,-44.12],[168.95,-43.94],[169.67,-43.56]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[147.69,-40.81],[148.29,-40.88],[148.36,-42.06],[148.02,-42.41],[147.91,-43.21],[147.56,-42.94],[146.87,-43.63],[146.66,-43.58],[146.05,-43.55],[145.43,-42.69],[145.3,-42.03],[144.72,-41.16],[144.74,-40.7],[145.4,-40.79],[146.36,-41.14],[146.91,-41],[147.69,-40.81]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[126.15,-32.22],[125.09,-32.73],[124.22,-32.96],[124.03,-33.48],[123.66,-33.89],[122.81,-33.91],[122.18,-34],[121.3,-33.82],[120.58,-33.93],[119.89,-33.98],[119.3,-34.51],[119.01,-34.46],[118.51,-34.75],[118.02,-35.06],[117.3,-35.03],[116.63,-35.03],[115.56,-34.39],[115.03,-34.2],[115.05,-33.62],[115.55,-33.49],[115.71,-33.26],[115.68,-32.9],[115.8,-32.21],[115.69,-31.61],[115.16,-30.6],[115,-30.03],[115.04,-29.46],[114.64,-28.81],[114.62,-28.52],[114.17,-28.12],[114.05,-27.33],[113.48,-26.54],[113.34,-26.12],[113.78,-26.55],[113.44,-25.62],[113.94,-25.91],[114.23,-26.3],[114.22,-25.79],[113.72,-25],[113.63,-24.68],[113.39,-24.38],[113.5,-23.81],[113.71,-23.56],[113.84,-23.06],[113.74,-22.48],[114.15,-21.76],[114.23,-22.52],[114.65,-21.83],[115.46,-21.5],[115.95,-21.07],[116.71,-20.7],[117.17,-20.62],[117.44,-20.75],[118.23,-20.37],[118.84,-20.26],[118.99,-20.04],[119.25,-19.95],[119.81,-19.98],[120.86,-19.68],[121.4,-19.24],[121.66,-18.71],[122.24,-18.2],[122.29,-17.8],[122.31,-17.25],[123.01,-16.41],[123.43,-17.27],[123.86,-17.07],[123.5,-16.6],[123.82,-16.11],[124.26,-16.33],[124.38,-15.57],[124.93,-15.08],[125.17,-14.68],[125.67,-14.51],[125.69,-14.23],[126.13,-14.35],[126.14,-14.1],[126.58,-13.95],[127.07,-13.82],[127.8,-14.28],[128.36,-14.87],[128.99,-14.88],[129.62,-14.97],[129.41,-14.42],[129.89,-13.62],[130.34,-13.36],[130.18,-13.11],[130.62,-12.54],[131.22,-12.18],[131.74,-12.3],[132.58,-12.11],[132.56,-11.6],[131.82,-11.27],[132.36,-11.13],[133.02,-11.38],[133.55,-11.79],[134.39,-12.04],[134.68,-11.94],[135.3,-12.25],[135.88,-11.96],[136.26,-12.05],[136.49,-11.86],[136.95,-12.35],[136.69,-12.89],[136.31,-13.29],[135.96,-13.32],[136.08,-13.72],[135.78,-14.22],[135.43,-14.72],[135.5,-15],[136.3,-15.55],[137.07,-15.87],[137.58,-16.22],[138.3,-16.81],[138.59,-16.81],[139.11,-17.06],[139.26,-17.37],[140.22,-17.71],[140.88,-17.37],[141.07,-16.83],[141.27,-16.39],[141.4,-15.84],[141.7,-15.04],[141.56,-14.56],[141.64,-14.27],[141.52,-13.7],[141.65,-12.94],[141.84,-12.74],[141.69,-12.41],[141.93,-11.88],[142.12,-11.33],[142.14,-11.04],[142.52,-10.67],[142.8,-11.16],[142.87,-11.78],[143.12,-11.91],[143.16,-12.33],[143.52,-12.83],[143.6,-13.4],[143.56,-13.76],[143.92,-14.55],[144.56,-14.17],[144.89,-14.59],[145.37,-14.98],[145.27,-15.43],[145.49,-16.29],[145.64,-16.78],[145.89,-16.91],[146.16,-17.76],[146.06,-18.28],[146.39,-18.96],[147.47,-19.48],[148.18,-19.96],[148.85,-20.39],[148.72,-20.63],[149.29,-21.26],[149.68,-22.34],[150.08,-22.12],[150.48,-22.56],[150.73,-22.4],[150.9,-23.46],[151.61,-24.08],[152.07,-24.46],[152.86,-25.27],[153.14,-26.07],[153.16,-26.64],[153.09,-27.26],[153.57,-28.11],[153.51,-29],[153.34,-29.46],[153.07,-30.35],[153.09,-30.92],[152.89,-31.64],[152.45,-32.55],[151.71,-33.04],[151.34,-33.82],[151.01,-34.31],[150.71,-35.17],[150.33,-35.67],[150.08,-36.42],[149.95,-37.11],[150,-37.43],[149.42,-37.77],[148.3,-37.81],[147.38,-38.22],[146.92,-38.61],[146.32,-39.04],[145.49,-38.59],[144.88,-38.42],[145.03,-37.9],[144.49,-38.09],[143.61,-38.81],[142.75,-38.54],[142.18,-38.38],[141.61,-38.31],[140.64,-38.02],[139.99,-37.4],[139.81,-36.64],[139.57,-36.14],[139.08,-35.73],[138.12,-35.61],[138.45,-35.13],[138.21,-34.38],[137.72,-35.08],[136.83,-35.26],[137.35,-34.71],[137.5,-34.13],[137.89,-33.64],[137.81,-32.9],[137,-33.75],[136.37,-34.09],[135.99,-34.89],[135.21,-34.48],[135.24,-33.95],[134.61,-33.22],[134.09,-32.85],[134.27,-32.62],[132.99,-32.01],[132.29,-31.98],[131.33,-31.5],[129.54,-31.59],[128.24,-31.95],[127.1,-32.28],[126.15,-32.22]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[122.93,0.88],[124.08,0.92],[125.07,1.64],[125.24,1.42],[124.44,0.43],[123.69,0.24],[122.72,0.43],[121.06,0.38],[120.18,0.24],[120.04,-0.52],[120.94,-1.41],[121.48,-0.96],[123.34,-0.62],[123.26,-1.08],[122.82,-0.93],[122.39,-1.52],[121.51,-1.9],[122.45,-3.19],[122.27,-3.53],[123.17,-4.68],[123.16,-5.34],[122.63,-5.63],[122.24,-5.28],[122.72,-4.46],[121.74,-4.85],[121.49,-4.57],[121.62,-4.19],[120.9,-3.6],[120.97,-2.63],[120.31,-2.93],[120.39,-4.1],[120.43,-5.53],[119.8,-5.67],[119.37,-5.38],[119.65,-4.46],[119.5,-3.49],[119.08,-3.49],[118.77,-2.8],[119.18,-2.15],[119.32,-1.35],[119.83,0.15],[120.04,0.57],[120.89,1.31],[121.67,1.01],[122.93,0.88]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[108.49,-6.42],[108.62,-6.78],[110.54,-6.88],[110.76,-6.47],[112.61,-6.95],[112.98,-7.59],[114.48,-7.78],[115.71,-8.37],[114.56,-8.75],[113.46,-8.35],[112.56,-8.38],[111.52,-8.3],[110.59,-8.12],[109.43,-7.74],[108.69,-7.64],[108.28,-7.77],[106.45,-7.35],[106.28,-6.92],[105.37,-6.85],[106.05,-5.9],[107.27,-5.95],[108.07,-6.35],[108.49,-6.42]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[104.37,-1.08],[104.54,-1.78],[104.89,-2.34],[105.62,-2.43],[106.11,-3.06],[105.86,-4.31],[105.82,-5.85],[104.71,-5.87],[103.87,-5.04],[102.58,-4.22],[102.16,-3.61],[101.4,-2.8],[100.9,-2.05],[100.14,-0.65],[99.26,0.18],[98.97,1.04],[98.6,1.82],[97.7,2.45],[97.18,3.31],[96.42,3.87],[95.38,4.97],[95.29,5.48],[95.94,5.44],[97.48,5.25],[98.37,4.27],[99.14,3.59],[99.69,3.17],[100.64,2.1],[101.66,2.08],[102.5,1.4],[103.08,0.56],[103.84,0.1],[103.44,-0.71],[104.01,-1.06],[104.37,-1.08]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[126.38,8.41],[126.48,7.75],[126.54,7.19],[126.2,6.27],[125.83,7.29],[125.36,6.79],[125.68,6.05],[125.4,5.58],[124.22,6.16],[123.94,6.89],[124.24,7.36],[123.61,7.83],[123.3,7.42],[122.83,7.46],[122.09,6.9],[121.92,7.19],[122.31,8.03],[122.94,8.32],[123.49,8.69],[123.84,8.24],[124.6,8.51],[124.76,8.96],[125.47,8.99],[125.41,9.76],[126.22,9.29],[126.31,8.78],[126.38,8.41]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[141.88,39.18],[140.96,38.17],[140.98,37.14],[140.6,36.34],[140.77,35.84],[140.25,35.14],[138.98,34.67],[137.22,34.61],[135.79,33.46],[135.12,33.85],[135.08,34.6],[133.34,34.38],[132.16,33.9],[130.99,33.89],[132,33.15],[131.33,31.45],[130.69,31.03],[130.2,31.42],[130.45,32.32],[129.81,32.61],[129.41,33.3],[130.35,33.6],[130.88,34.23],[131.88,34.75],[132.62,35.43],[134.61,35.73],[135.68,35.53],[136.72,37.3],[137.39,36.83],[138.86,37.83],[139.43,38.22],[140.05,39.44],[139.88,40.56],[140.31,41.2],[141.37,41.38],[141.91,39.99],[141.88,39.18]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[144.61,43.96],[145.32,44.38],[145.54,43.26],[144.06,42.99],[143.18,42],[141.61,42.68],[141.07,41.58],[139.96,41.57],[139.82,42.56],[140.31,43.33],[141.38,43.39],[141.67,44.77],[141.97,45.55],[143.14,44.51],[143.91,44.17],[144.61,43.96]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-4.21,58.55],[-3.01,58.64],[-4.07,57.55],[-3.06,57.69],[-1.96,57.68],[-2.22,56.87],[-3.12,55.97],[-2.09,55.91],[-1.11,54.62],[-0.43,54.46],[0.18,53.33],[0.47,52.93],[1.68,52.74],[1.56,52.1],[1.05,51.81],[1.45,51.29],[0.55,50.77],[-0.79,50.77],[-2.49,50.5],[-2.96,50.7],[-3.62,50.23],[-4.54,50.34],[-5.25,49.96],[-5.78,50.16],[-4.31,51.21],[-3.41,51.43],[-4.98,51.59],[-5.27,51.99],[-4.22,52.3],[-4.77,52.84],[-4.58,53.5],[-3.09,53.4],[-2.95,53.98],[-3.63,54.62],[-4.84,54.79],[-5.08,55.06],[-4.72,55.51],[-5.05,55.78],[-5.59,55.31],[-5.64,56.28],[-6.15,56.79],[-5.79,57.82],[-5.01,58.63],[-4.21,58.55]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-14.51,66.46],[-14.74,65.81],[-13.61,65.13],[-14.91,64.36],[-17.79,63.68],[-18.66,63.5],[-19.97,63.64],[-22.76,63.96],[-21.78,64.4],[-23.96,64.89],[-22.18,65.08],[-22.23,65.38],[-24.33,65.61],[-23.65,66.26],[-22.13,66.41],[-20.58,65.73],[-19.06,66.28],[-17.8,65.99],[-16.17,66.53],[-14.51,66.46]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[142.91,53.7],[143.26,52.74],[143.24,51.76],[143.65,50.75],[144.65,48.98],[143.17,49.31],[142.56,47.86],[143.53,46.84],[143.51,46.14],[142.75,46.74],[142.09,45.97],[141.91,46.81],[142.02,47.78],[141.9,48.86],[142.14,49.62],[142.18,50.95],[141.59,51.94],[141.68,53.3],[142.61,53.76],[142.21,54.23],[142.65,54.37],[142.91,53.7]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[122.34,18.22],[122.17,17.81],[122.52,17.09],[122.25,16.26],[121.66,15.93],[121.51,15.12],[121.73,14.33],[122.26,14.22],[122.7,14.34],[123.95,13.78],[123.86,13.24],[124.18,13],[124.08,12.54],[123.3,13.03],[122.93,13.55],[122.67,13.19],[122.03,13.78],[121.13,13.64],[120.63,13.86],[120.68,14.27],[120.99,14.53],[120.69,14.76],[120.56,14.4],[120.07,14.97],[119.92,15.41],[119.88,16.36],[120.29,16.03],[120.39,17.6],[120.72,18.51],[121.32,18.5],[121.94,18.22],[122.25,18.48],[122.34,18.22]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-77.35,8.67],[-76.84,8.64],[-76.09,9.34],[-75.67,9.44],[-75.66,9.77],[-75.48,10.62],[-74.91,11.08],[-74.28,11.1],[-74.2,11.31],[-73.41,11.23],[-72.63,11.73],[-72.24,11.96],[-71.75,12.44],[-71.4,12.38],[-71.14,12.11],[-71.33,11.78],[-71.36,11.54],[-71.95,11.42],[-71.62,10.97],[-71.63,10.45],[-72.07,9.87],[-71.7,9.07],[-71.26,9.14],[-71.04,9.86],[-71.35,10.21],[-71.4,10.97],[-70.16,11.38],[-70.29,11.85],[-69.94,12.16],[-69.58,11.46],[-68.88,11.44],[-68.23,10.89],[-68.19,10.55],[-67.3,10.55],[-66.23,10.65],[-65.66,10.2],[-64.89,10.08],[-64.33,10.39],[-64.32,10.64],[-63.08,10.7],[-61.88,10.72],[-62.73,10.42],[-62.39,9.95],[-61.59,9.87],[-60.83,9.38],[-60.67,8.58],[-60.15,8.6],[-59.76,8.37],[-59.1,8],[-58.48,7.35],[-58.45,6.83],[-58.08,6.81],[-57.54,6.32],[-57.15,5.97],[-55.95,5.77],[-55.84,5.95],[-55.03,6.03],[-53.96,5.76],[-53.62,5.65],[-52.88,5.41],[-51.82,4.57],[-51.66,4.16],[-51.3,4.12],[-51.07,3.65],[-50.51,1.9],[-49.97,1.74],[-49.95,1.05],[-50.7,0.22],[-50.39,-0.08],[-48.62,-0.24],[-48.58,-1.24],[-47.82,-0.58],[-46.57,-0.94],[-44.91,-1.55],[-44.42,-2.14],[-44.58,-2.69],[-43.42,-2.38],[-41.47,-2.91],[-39.98,-2.87],[-38.5,-3.7],[-37.22,-4.82],[-36.45,-5.11],[-35.6,-5.15],[-35.24,-5.46],[-34.9,-6.74],[-34.73,-7.34],[-35.13,-9],[-35.64,-9.65],[-37.05,-11.04],[-37.68,-12.17],[-38.42,-13.04],[-38.67,-13.06],[-38.95,-13.79],[-38.88,-15.67],[-39.16,-17.21],[-39.27,-17.87],[-39.58,-18.26],[-39.76,-19.6],[-40.77,-20.9],[-40.94,-21.94],[-41.75,-22.37],[-41.99,-22.97],[-43.07,-22.97],[-44.65,-23.35],[-45.35,-23.8],[-46.47,-24.09],[-47.65,-24.89],[-48.5,-25.88],[-48.64,-26.62],[-48.47,-27.18],[-48.66,-28.19],[-48.89,-28.67],[-49.59,-29.22],[-50.7,-30.98],[-51.58,-31.78],[-52.26,-32.25],[-52.71,-33.2],[-53.37,-33.77],[-53.81,-34.4],[-54.94,-34.95],[-55.67,-34.75],[-56.22,-34.86],[-57.14,-34.43],[-57.82,-34.46],[-58.43,-33.91],[-58.5,-34.43],[-57.23,-35.29],[-57.36,-35.98],[-56.74,-36.41],[-56.79,-36.9],[-57.75,-38.18],[-59.23,-38.72],[-61.24,-38.93],[-62.34,-38.83],[-62.13,-39.42],[-62.33,-40.17],[-62.15,-40.68],[-62.75,-41.03],[-63.77,-41.17],[-64.73,-40.8],[-65.12,-41.06],[-64.98,-42.06],[-64.3,-42.36],[-63.76,-42.04],[-63.46,-42.56],[-64.38,-42.87],[-65.18,-43.5],[-65.33,-44.5],[-65.57,-45.04],[-66.51,-45.04],[-67.29,-45.55],[-67.58,-46.3],[-66.6,-47.03],[-65.64,-47.24],[-65.99,-48.13],[-67.17,-48.7],[-67.82,-49.87],[-68.73,-50.26],[-69.14,-50.73],[-68.82,-51.77],[-68.15,-52.35],[-68.57,-52.3],[-69.46,-52.29],[-69.94,-52.54],[-70.85,-52.9],[-71.01,-53.83],[-71.43,-53.86],[-72.56,-53.53],[-73.7,-52.84],[-74.95,-52.26]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-77.88,7.22],[-77.48,6.69],[-77.32,5.85],[-77.53,5.58],[-77.31,4.67],[-77.5,4.09],[-77.13,3.85],[-77.51,3.33],[-77.93,2.7],[-78.43,2.63],[-78.66,2.27],[-78.62,1.77],[-78.99,1.69],[-78.86,1.38],[-79.54,0.98],[-80.09,0.77],[-80.02,0.36],[-80.4,-0.28],[-80.58,-0.91],[-80.93,-1.06],[-80.76,-1.97],[-80.97,-2.25],[-80.37,-2.69],[-79.99,-2.22],[-79.77,-2.66],[-80.3,-3.4],[-81.1,-4.04],[-81.41,-4.74],[-80.93,-5.69],[-81.25,-6.14],[-80.54,-6.54],[-79.76,-7.19],[-79.45,-7.93],[-79.04,-8.39],[-78.09,-10.38],[-77.11,-12.22],[-76.26,-13.54],[-76.42,-13.82],[-76.01,-14.65],[-75.24,-15.27],[-73.44,-16.36],[-71.46,-17.36],[-71.38,-17.77],[-70.37,-18.35],[-70.16,-19.76],[-70.09,-21.39],[-70.4,-23.63],[-70.72,-25.71],[-70.91,-27.64],[-71.49,-28.86],[-71.37,-30.1],[-71.67,-30.92],[-71.44,-32.42],[-71.86,-33.91],[-72.55,-35.51],[-73.17,-37.12],[-73.59,-37.16],[-73.51,-38.28],[-73.22,-39.26],[-73.68,-39.94],[-74.02,-41.79],[-74.33,-43.22],[-73.7,-43.37],[-73.39,-42.12],[-72.72,-42.38],[-73.24,-44.45],[-74.35,-44.1],[-74.69,-45.76],[-75.64,-46.65],[-74.13,-46.94],[-75.18,-47.71],[-75.61,-48.67],[-75.48,-50.38],[-74.98,-51.04],[-75.26,-51.63],[-74.95,-52.26]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-74.66,-52.84],[-73.84,-53.05],[-72.43,-53.72],[-71.11,-54.07],[-70.59,-53.62],[-70.27,-52.93],[-69.35,-52.52],[-68.63,-52.64],[-68.25,-53.1],[-67.75,-53.85],[-66.45,-54.45],[-65.05,-54.7],[-65.5,-55.2],[-66.45,-55.25],[-66.96,-54.9],[-67.29,-55.3],[-68.15,-55.61],[-68.64,-55.58],[-69.23,-55.5],[-69.96,-55.2],[-71.01,-55.05],[-72.26,-54.5],[-73.29,-53.96],[-74.66,-52.84]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[53.51,73.75],[55.9,74.63],[55.63,75.08],[57.87,75.61],[61.17,76.25],[64.5,76.44],[66.21,76.81],[68.16,76.94],[68.85,76.54],[68.18,76.23],[64.64,75.74],[61.58,75.26],[58.48,74.31],[56.99,73.33],[55.42,72.37],[55.62,71.54],[57.54,70.72],[56.94,70.63],[53.68,70.76],[53.41,71.21],[51.6,71.47],[51.46,72.01],[52.48,72.23],[52.44,72.77],[54.43,73.63],[53.51,73.75]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[15.14,79.67],[15.52,80.02],[16.99,80.05],[18.25,79.7],[21.54,78.96],[19.03,78.56],[18.47,77.83],[17.59,77.64],[17.12,76.81],[15.91,76.77],[13.76,77.38],[14.67,77.74],[13.17,78.02],[11.22,78.87],[10.44,79.65],[13.17,80.01],[13.72,79.66],[15.14,79.67]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-77.88,7.22],[-78.21,7.51],[-78.43,8.05],[-78.18,8.32],[-78.44,8.39],[-78.62,8.72],[-79.12,9],[-79.56,8.93],[-79.76,8.58],[-80.16,8.33],[-80.38,8.3],[-80.48,8.09],[-80,7.55],[-80.28,7.42],[-80.42,7.27],[-80.89,7.22],[-81.06,7.82],[-81.19,7.65],[-81.52,7.71],[-81.72,8.11],[-82.13,8.18],[-82.39,8.29],[-82.61,8.29],[-82.82,8.29],[-82.85,8.07],[-82.97,8.23],[-83.51,8.45],[-83.71,8.66],[-83.6,8.83],[-83.63,9.05],[-83.91,9.29],[-84.3,9.49],[-84.65,9.62],[-84.71,9.91],[-84.98,10.09],[-84.91,9.8],[-85.11,9.56],[-85.34,9.83],[-85.66,9.93],[-85.8,10.13],[-85.79,10.44],[-85.66,10.75],[-85.8,10.82],[-85.94,10.9],[-85.71,11.09],[-86.06,11.4],[-86.53,11.81],[-86.75,12.14],[-87.17,12.46],[-87.67,12.91],[-87.56,13.06],[-87.39,12.91],[-87.32,12.98],[-87.49,13.3],[-87.64,13.34],[-87.79,13.38],[-87.9,13.15],[-88.48,13.16],[-88.84,13.26],[-89.26,13.46],[-89.81,13.52],[-90.1,13.74],[-90.61,13.91],[-91.23,13.93],[-91.69,14.13],[-92.23,14.54],[-93.36,15.62],[-93.88,15.94],[-94.69,16.2],[-95.25,16.13],[-96.05,15.75],[-96.56,15.65],[-97.26,15.92],[-98.01,16.11],[-98.95,16.57],[-99.7,16.71],[-100.83,17.17],[-101.67,17.65],[-101.92,17.92],[-102.48,17.98],[-103.5,18.29],[-103.92,18.75],[-104.99,19.32],[-105.49,19.95],[-105.73,20.43],[-105.4,20.53],[-105.5,20.82],[-105.27,21.08],[-105.27,21.42],[-105.6,21.87],[-105.69,22.27],[-106.03,22.77],[-106.91,23.77],[-107.92,24.55],[-108.4,25.17],[-109.26,25.58],[-109.44,25.82],[-109.29,26.44],[-109.8,26.68],[-110.39,27.16],[-110.64,27.86],[-111.18,27.94],[-111.76,28.47],[-112.23,28.95],[-112.27,29.27],[-112.81,30.02],[-113.16,30.79],[-113.15,31.17],[-113.87,31.57],[-114.21,31.52],[-114.78,31.8],[-114.94,31.39],[-114.77,30.91],[-114.67,30.16],[-114.33,29.75],[-113.59,29.06],[-113.42,28.83],[-113.27,28.75],[-113.14,28.41],[-112.96,28.43],[-112.76,27.78],[-112.46,27.53],[-112.24,27.17],[-111.62,26.66],[-111.28,25.73],[-110.99,25.29],[-110.71,24.83],[-110.66,24.3],[-110.17,24.27],[-109.77,23.81],[-109.41,23.36],[-109.43,23.19],[-109.85,22.82],[-110.03,22.82],[-110.3,23.43],[-110.95,24],[-111.67,24.48],[-112.18,24.74],[-112.15,25.47],[-112.3,26.01],[-112.78,26.32],[-113.46,26.77],[-113.6,26.64],[-113.85,26.9],[-114.47,27.14],[-115.06,27.72],[-114.98,27.8],[-114.57,27.74],[-114.2,28.12],[-114.16,28.57],[-114.93,29.28],[-115.52,29.56],[-115.89,30.18],[-116.26,30.84],[-116.72,31.64],[-117.13,32.54],[-117.3,33.05],[-117.94,33.62],[-118.41,33.74],[-118.52,34.03],[-119.08,34.08],[-119.44,34.35],[-120.37,34.45],[-120.62,34.61],[-120.74,35.16],[-121.71,36.16],[-122.55,37.55],[-122.51,37.78],[-122.95,38.11],[-123.73,38.95],[-123.87,39.77],[-124.4,40.31],[-124.18,41.14],[-124.21,42],[-124.53,42.77],[-124.14,43.71],[-124.02,44.62],[-123.9,45.52],[-124.08,46.86],[-124.4,47.72],[-124.69,48.18],[-124.57,48.38],[-123.12,48.04],[-122.59,47.1],[-122.34,47.36],[-122.5,48.18],[-122.84,49],[-122.97,49],[-124.91,49.98],[-125.62,50.42],[-127.44,50.83],[-127.99,51.72],[-127.85,52.33],[-129.13,52.76],[-129.31,53.56],[-130.51,54.29],[-130.54,54.8],[-131.09,55.18],[-131.97,55.5],[-132.25,56.37],[-133.54,57.18],[-134.08,58.12],[-135.04,58.19],[-136.63,58.21],[-137.8,58.5],[-139.87,59.54],[-140.83,59.73],[-142.57,60.08],[-143.96,60],[-145.93,60.46],[-147.11,60.88],[-148.22,60.67],[-148.02,59.98],[-148.57,59.91],[-149.73,59.71],[-150.61,59.37],[-151.72,59.16],[-151.86,59.74],[-151.41,60.73],[-150.35,61.03],[-150.62,61.28],[-151.9,60.73],[-152.58,60.06],[-154.02,59.35],[-153.29,58.86],[-154.23,58.15],[-155.31,57.73],[-156.31,57.42],[-156.56,56.98],[-158.12,56.46],[-158.43,55.99],[-159.6,55.57],[-160.29,55.64],[-161.22,55.36],[-162.24,55.02],[-163.07,54.69],[-164.79,54.4],[-164.94,54.57],[-163.85,55.04],[-162.87,55.35],[-161.8,55.89],[-160.56,56.01],[-160.07,56.42],[-158.68,57.02],[-158.46,57.22],[-157.72,57.57],[-157.55,58.33],[-157.04,58.92],[-158.19,58.62],[-158.52,58.79],[-159.06,58.42],[-159.71,58.93],[-159.98,58.57],[-160.36,59.07],[-161.36,58.67],[-161.97,58.67],[-162.05,59.27],[-161.87,59.63],[-162.52,59.99],[-163.82,59.8],[-164.66,60.27],[-165.35,60.51],[-165.35,61.07],[-166.12,61.5],[-165.73,62.07],[-164.92,62.63],[-164.56,63.15],[-163.75,63.22],[-163.07,63.06],[-162.26,63.54],[-161.53,63.46],[-160.77,63.77],[-160.96,64.22],[-161.52,64.4],[-160.78,64.79],[-161.39,64.78],[-162.45,64.56],[-162.76,64.34],[-163.55,64.56],[-164.96,64.45],[-166.43,64.69],[-166.85,65.09],[-168.11,65.67],[-166.71,66.09],[-164.47,66.58],[-163.65,66.58],[-163.79,66.08],[-161.68,66.12],[-162.49,66.74],[-163.72,67.12],[-164.43,67.62],[-165.39,68.04],[-166.76,68.36],[-166.2,68.88],[-164.43,68.92],[-163.17,69.37],[-162.93,69.86],[-161.91,70.33],[-160.93,70.45],[-159.04,70.89],[-158.12,70.82],[-156.58,71.36],[-155.07,71.15],[-154.34,70.7],[-153.9,70.89],[-152.21,70.83],[-152.27,70.6],[-150.74,70.43],[-149.72,70.53],[-147.61,70.21],[-145.69,70.12],[-144.92,69.99],[-143.59,70.15],[-142.07,69.85],[-140.99,69.71],[-139.12,69.47],[-137.55,68.99],[-136.5,68.9],[-135.63,69.32],[-134.41,69.63],[-132.93,69.51],[-131.43,69.94],[-129.79,70.19],[-129.11,69.78],[-128.36,70.01],[-128.14,70.48],[-127.45,70.38],[-125.76,69.48],[-124.42,70.16],[-124.29,69.4],[-123.06,69.56],[-122.68,69.86],[-121.47,69.8],[-119.94,69.38],[-117.6,69.01],[-116.23,68.84],[-115.25,68.91],[-113.9,68.4],[-115.3,67.9],[-113.5,67.69],[-110.8,67.81],[-109.95,67.98],[-108.88,67.38],[-107.79,67.89],[-108.81,68.31],[-108.17,68.65],[-106.95,68.7],[-106.15,68.8],[-105.34,68.56],[-104.34,68.02],[-103.22,68.1],[-101.45,67.65],[-99.9,67.81],[-98.44,67.78],[-98.56,68.4],[-97.67,68.58],[-96.12,68.24],[-96.13,67.29],[-95.49,68.09],[-94.68,68.06],[-94.23,69.07],[-95.3,69.69],[-96.47,70.09],[-96.39,71.19],[-95.21,71.92],[-93.89,71.76],[-92.88,71.32],[-91.52,70.19],[-92.41,69.7],[-90.55,69.5],[-90.55,68.47],[-89.22,69.26],[-88.02,68.62],[-88.32,67.87],[-87.35,67.2],[-86.31,67.92],[-85.58,68.78],[-85.52,69.88],[-84.1,69.81],[-82.62,69.66],[-81.28,69.16],[-81.22,68.67],[-81.96,68.13],[-81.26,67.6],[-81.39,67.11],[-83.34,66.41],[-84.74,66.26],[-85.77,66.56],[-86.07,66.06],[-87.03,65.21],[-87.32,64.78],[-88.48,64.1],[-89.91,64.03],[-90.7,63.61],[-90.77,62.96],[-91.93,62.84],[-93.16,62.02],[-94.24,60.9],[-94.63,60.11],[-94.68,58.95],[-93.22,58.78],[-92.76,57.85],[-92.3,57.09],[-90.9,57.28],[-89.04,56.85],[-88.04,56.47],[-87.32,56],[-86.07,55.72],[-85.01,55.3],[-83.36,55.24],[-82.27,55.15],[-82.44,54.28],[-82.13,53.28],[-81.4,52.16],[-79.91,51.21],[-79.14,51.53],[-78.6,52.56],[-79.12,54.14],[-79.83,54.67],[-78.23,55.14],[-77.1,55.84],[-76.54,56.53],[-76.62,57.2],[-77.3,58.05],[-78.52,58.8],[-77.34,59.85],[-77.77,60.76],[-78.11,62.32],[-77.41,62.55],[-75.7,62.28],[-74.67,62.18],[-73.84,62.44],[-72.91,62.11],[-71.68,61.53],[-71.37,61.14],[-69.59,61.06],[-69.62,60.22],[-69.29,58.96],[-68.37,58.8],[-67.65,58.21],[-66.2,58.77],[-65.25,59.87],[-64.58,60.34],[-63.8,59.44],[-62.5,58.17],[-61.4,56.97],[-61.8,56.34],[-60.47,55.78],[-59.57,55.2],[-57.98,54.95],[-57.33,54.63],[-56.94,53.78],[-56.16,53.65],[-55.76,53.27],[-55.68,52.15],[-56.41,51.77],[-57.13,51.42],[-58.77,51.06],[-60.03,50.24],[-61.72,50.08],[-63.86,50.29],[-65.36,50.3],[-66.4,50.23],[-67.24,49.51],[-68.51,49.07],[-69.95,47.74],[-71.1,46.82],[-70.26,46.99],[-68.65,48.3],[-66.55,49.13],[-65.06,49.23],[-64.17,48.74],[-65.12,48.07],[-64.8,46.99],[-64.47,46.24],[-63.17,45.74],[-61.52,45.88],[-60.52,47.01],[-60.45,46.28],[-59.8,45.92],[-61.04,45.27],[-63.25,44.67],[-64.25,44.27],[-65.36,43.55],[-66.12,43.62],[-66.16,44.47],[-64.43,45.29],[-66.03,45.26],[-67.14,45.14],[-66.96,44.81],[-68.03,44.33],[-69.06,43.98],[-70.12,43.68],[-70.65,43.09],[-70.81,42.87],[-70.82,42.33],[-70.49,41.8],[-70.08,41.78],[-70.18,42.14],[-69.88,41.92],[-69.97,41.64],[-70.64,41.47],[-71.12,41.49],[-71.85,41.32],[-72.29,41.27],[-72.88,41.22],[-73.71,40.93],[-72.24,41.12],[-71.95,40.93],[-73.35,40.63],[-73.98,40.63],[-73.95,40.75],[-74.26,40.47],[-73.96,40.43],[-74.18,39.71],[-74.91,38.94],[-74.98,39.2],[-75.2,39.25],[-75.53,39.5],[-75.32,38.96],[-75.07,38.78],[-75.06,38.4],[-75.38,38.02],[-75.94,37.22],[-76.03,37.26],[-75.72,37.94],[-76.23,38.32],[-76.35,39.15],[-76.54,38.72],[-76.33,38.08],[-76.99,38.24],[-76.3,37.92],[-76.26,36.97],[-75.97,36.9],[-75.87,36.55],[-75.73,35.55],[-76.36,34.81],[-77.4,34.51],[-78.05,33.93],[-78.55,33.86],[-79.06,33.49],[-79.2,33.16],[-80.3,32.51],[-80.86,32.03],[-81.34,31.44],[-81.49,30.73],[-81.31,30.04],[-80.98,29.18],[-80.54,28.47],[-80.53,28.04],[-80.06,26.88],[-80.09,26.21],[-80.13,25.82],[-80.38,25.21],[-80.68,25.08],[-81.17,25.2],[-81.33,25.64],[-81.71,25.87],[-82.24,26.73],[-82.71,27.5],[-82.86,27.89],[-82.65,28.55],[-82.93,29.1],[-83.71,29.94],[-84.1,30.09],[-85.11,29.64],[-85.29,29.69],[-85.77,30.15],[-86.4,30.4],[-87.53,30.27],[-88.42,30.38],[-89.18,30.32],[-89.59,30.16],[-89.41,29.89],[-89.43,29.49],[-89.22,29.29],[-89.41,29.16],[-89.78,29.31],[-90.15,29.12],[-90.88,29.15],[-91.63,29.68],[-92.5,29.55],[-93.23,29.78],[-93.85,29.71],[-94.69,29.48],[-95.6,28.74],[-96.59,28.31],[-97.14,27.83],[-97.37,27.38],[-97.38,26.69],[-97.33,26.21],[-97.14,25.87],[-97.53,24.99],[-97.7,24.27],[-97.78,22.93],[-97.87,22.44],[-97.7,21.9],[-97.39,21.41],[-97.19,20.64],[-96.53,19.89],[-96.29,19.32],[-95.9,18.83],[-94.84,18.56],[-94.43,18.14],[-93.55,18.42],[-92.79,18.52],[-92.04,18.7],[-91.41,18.88],[-90.77,19.28],[-90.53,19.87],[-90.45,20.71],[-90.28,21],[-89.6,21.26],[-88.54,21.49],[-87.66,21.46],[-87.05,21.54],[-86.81,21.33],[-86.85,20.85],[-87.38,20.26],[-87.62,19.65],[-87.44,19.47],[-87.59,19.04],[-87.84,18.26],[-88.09,18.52],[-88.3,18.5],[-88.3,18.35],[-88.11,18.35],[-88.12,18.08],[-88.29,17.64],[-88.2,17.49],[-88.3,17.13],[-88.24,17.04],[-88.36,16.53],[-88.55,16.27],[-88.73,16.23],[-88.93,15.89],[-88.6,15.71],[-88.52,15.86],[-88.19,15.72],[-88.12,15.69],[-87.9,15.86],[-87.62,15.88],[-87.52,15.8],[-87.37,15.85],[-86.9,15.76],[-86.44,15.78],[-86.12,15.89],[-86,16.01],[-85.68,15.95],[-85.44,15.89],[-85.18,15.91],[-84.98,16],[-84.53,15.86],[-84.37,15.84],[-84.06,15.65],[-83.77,15.42],[-83.41,15.27],[-83.15,15],[-83.23,14.9],[-83.28,14.68],[-83.18,14.31],[-83.41,13.97],[-83.52,13.57],[-83.55,13.13],[-83.5,12.87],[-83.47,12.42],[-83.63,12.32],[-83.72,11.89],[-83.65,11.63],[-83.86,11.37],[-83.81,11.1],[-83.66,10.94],[-83.59,10.79],[-83.4,10.4],[-83.02,9.99],[-82.55,9.57]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-82.55,9.57],[-82.19,9.21],[-82.21,9],[-81.81,8.95],[-81.71,9.03],[-81.44,8.79],[-80.95,8.86],[-80.52,9.11],[-79.91,9.31],[-79.57,9.61],[-79.02,9.55],[-79.06,9.45],[-78.5,9.42],[-78.06,9.25],[-77.73,8.95],[-77.35,8.67]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-71.71,19.71],[-71.59,19.88],[-71.38,19.9],[-70.81,19.88],[-70.21,19.62],[-69.95,19.65],[-69.77,19.29],[-69.22,19.31],[-69.25,19.02],[-68.81,18.98],[-68.32,18.61],[-68.69,18.21],[-69.16,18.42],[-69.62,18.38],[-69.95,18.43],[-70.13,18.25],[-70.52,18.18],[-70.67,18.43],[-71,18.28],[-71.4,17.6],[-71.66,17.76],[-71.71,18.04],[-72.37,18.21],[-72.84,18.15],[-73.45,18.22],[-73.92,18.03],[-74.46,18.34],[-74.37,18.66],[-73.45,18.53],[-72.69,18.45],[-72.33,18.67],[-72.79,19.1],[-72.78,19.48],[-73.42,19.64],[-73.19,19.92],[-72.58,19.87],[-71.71,19.71]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-16.26,19.1],[-16.38,19.59],[-16.28,20.09],[-16.54,20.57],[-17.06,21],[-17.02,21.42],[-16.97,21.89],[-16.59,22.16],[-16.26,22.68],[-16.33,23.02],[-15.98,23.72],[-15.43,24.36],[-15.09,24.52],[-14.82,25.1],[-14.8,25.64],[-14.44,26.25],[-13.77,26.62],[-13.14,27.64],[-12.62,28.04],[-11.69,28.15],[-10.9,28.83],[-10.4,29.1],[-9.56,29.93],[-9.81,31.18],[-9.43,32.04],[-9.3,32.56],[-8.66,33.24],[-7.65,33.7],[-6.91,34.11],[-6.24,35.15],[-5.93,35.76],[-5.19,35.76],[-4.59,35.33],[-3.64,35.4],[-2.6,35.18],[-2.17,35.17],[-1.21,35.71],[-0.13,35.89],[0.5,36.3],[1.47,36.61],[3.16,36.78],[4.82,36.87],[5.32,36.72],[6.26,37.11],[7.33,37.12],[7.74,36.89],[8.42,36.95],[9.51,37.35],[10.21,37.23],[10.18,36.72],[11.03,37.09],[11.1,36.9],[10.6,36.41],[10.59,35.95],[10.94,35.7],[10.81,34.83],[10.15,34.33],[10.34,33.79],[10.86,33.77],[11.11,33.29],[11.49,33.14],[12.66,32.79],[13.08,32.88],[13.92,32.71],[15.25,32.27],[15.71,31.38],[16.61,31.18],[18.02,30.76],[19.09,30.27],[19.57,30.53],[20.05,30.99],[19.82,31.75],[20.13,32.24],[20.85,32.71],[21.54,32.84],[22.9,32.64],[23.24,32.19],[23.61,32.19],[23.93,32.02],[24.92,31.9],[25.16,31.57],[26.5,31.59],[27.46,31.32],[28.45,31.03],[28.91,30.87],[29.68,31.19],[30.1,31.47],[30.98,31.56],[31.69,31.43],[31.96,30.93],[32.19,31.26],[32.99,31.02],[33.77,30.97],[34.27,31.22],[34.56,31.55],[34.49,31.61],[34.75,32.07],[34.96,32.83],[35.1,33.08],[35.13,33.09],[35.48,33.91],[35.98,34.61],[36,34.64],[35.91,35.41],[36.15,35.82],[35.78,36.27],[36.16,36.65],[35.55,36.57],[34.71,36.8],[34.03,36.22],[32.51,36.11],[31.7,36.64],[30.62,36.68],[30.39,36.26],[29.7,36.14],[28.73,36.68],[27.64,36.66],[27.05,37.65],[26.32,38.21],[26.8,38.99],[26.17,39.46],[27.28,40.42],[28.82,40.46],[29.24,41.22],[31.15,41.09],[32.35,41.74],[33.51,42.02],[35.17,42.04],[36.91,41.34],[38.35,40.95],[39.51,41.1],[40.37,41.01],[41.55,41.54],[41.7,41.96],[41.45,42.65],[40.88,43.01],[40.32,43.13],[39.96,43.43],[38.68,44.28],[37.54,44.66],[36.68,45.24],[37.4,45.4],[38.23,46.24],[37.67,46.64],[39.15,47.04],[39.12,47.26],[38.22,47.1],[37.43,47.02],[36.76,46.7],[35.82,46.65],[34.96,46.27],[35.02,45.65],[35.51,45.41],[36.53,45.47],[36.33,45.11],[35.24,44.94],[33.88,44.36],[33.33,44.56],[33.55,45.03],[32.45,45.33],[32.63,45.52],[33.59,45.85],[33.3,46.08],[31.74,46.33],[31.68,46.71],[30.75,46.58],[30.38,46.03],[29.6,45.29],[29.63,45.04],[29.14,44.82],[28.84,44.91],[28.56,43.71],[28.04,43.29],[27.67,42.58],[28,42.01],[28.12,41.62],[28.99,41.3],[28.81,41.05],[27.62,41],[27.19,40.69],[26.36,40.15],[26.04,40.62],[26.06,40.82],[25.45,40.85],[24.93,40.95],[23.71,40.69],[24.41,40.12],[23.9,39.96],[23.34,39.96],[22.81,40.48],[22.63,40.26],[22.85,39.66],[23.35,39.19],[22.97,38.97],[23.53,38.51],[24.03,38.22],[24.04,37.66],[23.12,37.92],[23.41,37.41],[22.77,37.31],[23.15,36.42],[22.49,36.41],[21.67,36.84],[21.3,37.64],[21.12,38.31],[20.73,38.77],[20.22,39.34],[20.15,39.62],[19.98,39.69],[19.96,39.92],[19.41,40.25],[19.32,40.73],[19.4,41.41],[19.54,41.72],[19.37,41.88],[19.16,41.96],[18.88,42.28],[18.45,42.48],[17.51,42.85],[16.93,43.21],[16.02,43.51],[15.17,44.24],[15.38,44.32],[14.92,44.74],[14.9,45.08],[14.26,45.23],[13.95,44.8],[13.66,45.14],[13.68,45.48],[13.72,45.5],[13.94,45.59],[13.14,45.74],[12.33,45.38],[12.38,44.89],[12.26,44.6],[12.59,44.09],[13.53,43.59],[14.03,42.76],[15.14,41.96],[15.93,41.96],[16.17,41.74],[15.89,41.54],[16.79,41.18],[17.52,40.88],[18.38,40.36],[18.48,40.17],[18.29,39.81],[17.74,40.28],[16.87,40.44],[16.45,39.8],[17.17,39.42],[17.05,38.9],[16.64,38.84],[16.1,37.99],[15.68,37.91],[15.69,38.21],[15.89,38.75],[16.11,38.96],[15.72,39.54],[15.41,40.05],[15,40.17],[14.7,40.6],[14.06,40.79],[13.63,41.19],[12.89,41.25],[12.11,41.7],[11.19,42.36],[10.51,42.93],[10.2,43.92],[9.7,44.04],[8.89,44.37],[8.43,44.23],[7.85,43.77],[7.44,43.69],[6.53,43.13],[4.56,43.4],[3.1,43.08],[2.99,42.47],[3.04,41.89],[2.09,41.23],[0.81,41.01],[0.72,40.68],[0.11,40.12],[-0.28,39.31],[0.11,38.74],[-0.47,38.29],[-0.68,37.64],[-1.44,37.44],[-2.15,36.67],[-3.42,36.66],[-4.37,36.68],[-5,36.32],[-5.38,35.95],[-5.87,36.03],[-6.24,36.37],[-6.52,36.94],[-7.45,37.1],[-7.86,36.84],[-8.38,36.98],[-8.9,36.87],[-8.75,37.65],[-8.84,38.27],[-9.29,38.36],[-9.53,38.74],[-9.45,39.39],[-9.05,39.76],[-8.98,40.16],[-8.77,40.76],[-8.79,41.18],[-8.99,41.54],[-9.03,41.88],[-8.98,42.59],[-9.39,43.03],[-7.98,43.75],[-6.75,43.57],[-5.41,43.57],[-4.35,43.4],[-3.52,43.46],[-1.9,43.42],[-1.38,44.02],[-1.19,46.01],[-2.23,47.06],[-2.96,47.57],[-4.49,47.95],[-4.59,48.68],[-3.3,48.9],[-1.62,48.64],[-1.93,49.78],[-0.99,49.35],[1.34,50.13],[1.64,50.95],[2.51,51.15],[3.31,51.35],[3.83,51.62],[4.71,53.09],[6.07,53.51],[6.91,53.48],[7.1,53.69],[7.94,53.75],[8.12,53.53],[8.8,54.02],[8.57,54.4],[8.53,54.96],[8.12,55.52],[8.09,56.54],[8.26,56.81],[8.54,57.11],[9.42,57.17],[9.78,57.45],[10.58,57.73],[10.55,57.22],[10.25,56.89],[10.37,56.61],[10.91,56.46],[10.67,56.08],[10.37,56.19],[9.65,55.47],[9.92,54.98],[9.94,54.6],[10.95,54.36],[10.94,54.01],[11.96,54.2],[12.52,54.47],[13.65,54.08],[14.12,53.76],[14.8,54.05],[16.36,54.51],[17.62,54.85],[18.62,54.68],[18.7,54.44],[19.66,54.43],[19.89,54.87],[21.27,55.19],[21.06,56.03],[21.09,56.78],[21.58,57.41],[22.52,57.75],[23.32,57.01],[24.12,57.03],[24.31,57.79],[24.43,58.38],[24.06,58.26],[23.43,58.61],[23.34,59.19],[24.6,59.47],[25.86,59.61],[26.95,59.45],[27.98,59.48],[29.12,60.03],[28.07,60.5],[26.26,60.42],[24.5,60.06],[22.87,59.85],[22.29,60.39],[21.32,60.72],[21.54,61.71],[21.06,62.61],[21.54,63.19],[22.44,63.82],[24.73,64.9],[25.4,65.11],[25.29,65.53],[23.9,66.01],[22.18,65.72],[21.21,65.03],[21.37,64.41],[19.78,63.61],[17.85,62.75],[17.12,61.34],[17.83,60.64],[18.79,60.08],[17.87,58.95],[16.83,58.72],[16.45,57.04],[15.88,56.1],[14.67,56.2],[14.1,55.41],[12.94,55.36],[12.63,56.31],[11.79,57.44],[11.03,58.86],[10.36,59.47],[8.38,58.31],[7.05,58.08],[5.67,58.59],[5.31,59.66],[4.99,61.97],[5.91,62.61],[8.55,63.45],[10.53,64.49],[12.36,65.88],[14.76,67.81],[16.44,68.56],[19.18,69.82],[21.38,70.26],[23.02,70.2],[24.55,71.03],[26.37,70.99],[28.17,71.19],[31.29,70.45],[30.01,70.19],[31.1,69.56],[32.13,69.91],[33.78,69.3],[36.51,69.06],[40.29,67.93],[41.06,67.46],[41.13,66.79],[40.02,66.27],[38.38,66],[33.92,66.76],[33.18,66.63],[34.81,65.9],[34.88,65.44],[34.94,64.41],[36.23,64.11],[37.01,63.85],[37.14,64.33],[36.54,64.76],[37.18,65.14],[39.59,64.52],[40.44,64.76],[39.76,65.5],[42.09,66.48],[43.02,66.42],[43.95,66.07],[44.53,66.76],[43.7,67.35],[44.17,67.96],[43.45,68.57],[46.25,68.25],[46.82,67.69],[45.56,67.57],[45.56,67.01],[46.35,66.67],[47.89,66.88],[48.14,67.52],[53.72,68.86],[54.47,68.81],[53.49,68.2],[54.75,68.09],[55.44,68.44],[57.32,68.47],[58.8,68.88],[59.94,68.28],[61.08,68.94],[60.03,69.52],[60.55,69.85],[63.5,69.55],[64.89,69.23],[68.51,68.09],[69.18,68.62],[68.16,69.14],[68.14,69.36],[66.93,69.45],[67.26,69.93],[66.72,70.71],[66.69,71.03],[68.54,71.93],[69.2,72.84],[69.94,73.04],[72.59,72.78],[72.8,72.22],[71.85,71.41],[72.47,71.09],[72.79,70.39],[72.56,69.02],[73.67,68.41],[73.24,67.74],[71.28,66.32],[72.42,66.17],[72.82,66.53],[73.92,66.79],[74.19,67.28],[75.05,67.76],[74.47,68.33],[74.94,68.99],[73.84,69.07],[73.6,69.63],[74.4,70.63],[73.1,71.45],[74.89,72.12],[74.66,72.83],[75.16,72.85],[75.68,72.3],[75.29,71.34],[76.36,71.15],[75.9,71.87],[77.58,72.27],[79.65,72.32],[81.5,71.75],[80.61,72.58],[80.51,73.65],[82.25,73.85],[84.66,73.81],[86.82,73.94],[86.01,74.46],[87.17,75.12],[88.32,75.14],[90.26,75.64],[92.9,75.77],[93.23,76.05],[95.86,76.14],[96.68,75.92],[98.92,76.45],[100.76,76.43],[101.04,76.86],[101.99,77.29],[104.35,77.7],[106.07,77.37],[104.71,77.13],[106.97,76.97],[107.24,76.48],[108.15,76.72],[111.08,76.71],[113.33,76.22],[114.13,75.85],[113.89,75.33],[112.78,75.03],[110.15,74.48],[109.4,74.18],[110.64,74.04],[112.12,73.79],[113.02,73.98],[113.53,73.34],[113.97,73.59],[115.57,73.75],[118.78,73.59],[119.02,73.12],[123.2,72.97],[123.26,73.74],[125.38,73.56],[126.98,73.57],[128.59,73.04],[129.05,72.4],[128.46,71.98],[129.72,71.19],[131.29,70.79],[132.25,71.84],[133.86,71.39],[135.56,71.66],[137.5,71.35],[138.23,71.63],[139.87,71.49],[139.15,72.42],[140.47,72.85],[149.5,72.2],[150.35,71.61],[152.97,70.84],[157.01,71.03],[159,70.87],[159.83,70.45],[159.71,69.72],[160.94,69.44],[162.28,69.64],[164.05,69.67],[165.94,69.47],[167.86,69.57],[169.58,68.69],[170.82,69.01],[170.01,69.65],[170.45,70.1],[173.64,69.82],[175.72,69.88],[178.6,69.4],[180,68.96]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[180,64.98],[178.71,64.53],[177.41,64.61],[178.31,64.08],[178.91,63.25],[179.37,62.98],[179.49,62.57],[179.23,62.3],[177.36,62.52],[174.57,61.77],[173.68,61.65],[172.15,60.95],[170.7,60.34],[170.33,59.88],[168.9,60.57],[166.29,59.79],[165.84,60.16],[164.88,59.73],[163.54,59.87],[163.22,59.21],[162.02,58.24],[162.05,57.84],[163.19,57.62],[163.06,56.16],[162.13,56.12],[161.7,55.29],[162.12,54.86],[160.37,54.34],[160.02,53.2],[158.53,52.96],[158.23,51.94],[156.79,51.01],[156.42,51.7],[155.99,53.16],[155.43,55.38],[155.91,56.77],[156.76,57.36],[156.81,57.83],[158.36,58.06],[160.15,59.31],[161.87,60.34],[163.67,61.14],[164.47,62.55],[163.26,62.47],[162.66,61.64],[160.12,60.54],[159.3,61.77],[156.72,61.43],[154.22,59.76],[155.04,59.14],[152.81,58.88],[151.27,58.78],[151.34,59.5],[149.78,59.66],[148.54,59.16],[145.49,59.34],[142.2,59.04],[138.96,57.09],[135.13,54.73],[136.7,54.6],[137.19,53.98],[138.16,53.76],[138.8,54.25],[139.9,54.19],[141.35,53.09],[141.38,52.24],[140.6,51.24],[140.51,50.05],[140.06,48.45],[138.55,47],[138.22,46.31],[136.86,45.14],[135.52,43.99],[134.87,43.4],[133.54,42.81],[132.91,42.8],[132.28,43.28],[130.94,42.55],[130.78,42.22],[130.4,42.28],[129.97,41.94],[129.67,41.6],[129.71,40.88],[129.19,40.66],[129.01,40.49],[128.63,40.19],[127.97,40.03],[127.53,39.76],[127.5,39.32],[127.39,39.21],[127.78,39.05],[128.35,38.61],[129.21,37.43],[129.46,36.78],[129.47,35.63],[129.09,35.08],[128.19,34.89],[127.39,34.48],[126.49,34.39],[126.37,34.93],[126.56,35.68],[126.12,36.73],[126.86,36.89],[126.17,37.75],[125.69,37.94],[125.57,37.75],[125.28,37.67],[125.24,37.86],[124.98,37.95],[124.71,38.11],[124.99,38.55],[125.22,38.67],[125.13,38.85],[125.39,39.39],[125.32,39.55],[124.74,39.66],[124.27,39.93],[122.87,39.64],[122.13,39.17],[121.05,38.9],[121.59,39.36],[121.38,39.75],[122.17,40.42],[121.64,40.95],[120.77,40.59],[119.64,39.9],[119.02,39.25],[118.04,39.2],[117.53,38.74],[118.06,38.06],[118.88,37.9],[118.91,37.45],[119.7,37.16],[120.82,37.87],[121.71,37.48],[122.36,37.45],[122.52,36.93],[121.1,36.65],[120.64,36.11],[119.66,35.61],[119.15,34.91],[120.23,34.36],[120.62,33.38],[121.23,32.46],[121.91,31.69],[121.89,30.95],[121.26,30.68],[121.5,30.14],[122.09,29.83],[121.94,29.02],[121.68,28.23],[121.13,28.14],[120.4,27.05],[119.59,25.74],[118.66,24.55],[117.28,23.62],[115.89,22.78],[114.76,22.67],[114.15,22.22],[113.81,22.55],[113.24,22.05],[111.84,21.55],[110.79,21.4],[110.51,20.57],[110.44,20.34],[109.89,20.28],[109.63,21.01],[109.86,21.4],[108.52,21.72],[108.05,21.55],[106.72,20.7],[105.88,19.75],[105.66,19.06],[106.43,18],[107.36,16.7],[108.27,16.08],[108.88,15.28],[109.34,13.43],[109.2,11.67],[108.37,11.01],[107.22,10.36],[106.41,9.53],[105.16,8.6],[104.8,9.24],[105.08,9.92],[104.33,10.49],[103.5,10.63],[103.09,11.15],[102.58,12.19],[101.69,12.65],[100.83,12.63],[100.98,13.41],[100.1,13.41],[100.02,12.31],[99.48,10.85],[99.15,9.96],[99.22,9.24],[99.87,9.21],[100.28,8.3],[100.46,7.43],[101.02,6.86],[101.62,6.74],[102.14,6.22],[102.37,6.13],[102.96,5.52],[103.38,4.86],[103.44,4.18],[103.33,3.73],[103.43,3.38],[103.5,2.79],[103.85,2.52],[104.25,1.63],[104.23,1.29],[103.52,1.23],[102.57,1.97],[101.39,2.76],[101.27,3.27],[100.7,3.94],[100.56,4.77],[100.2,5.31],[100.31,6.04],[100.09,6.46],[99.69,6.85],[99.52,7.34],[98.99,7.91],[98.5,8.38],[98.34,7.79],[98.15,8.35],[98.26,8.97],[98.55,9.93],[98.46,10.68],[98.76,11.44],[98.43,12.03],[98.51,13.12],[98.1,13.64],[97.78,14.84],[97.6,16.1],[97.16,16.93],[96.51,16.43],[95.37,15.71],[94.81,15.8],[94.19,16.04],[94.53,17.28],[94.32,18.21],[93.54,19.37],[93.66,19.73],[93.08,19.86],[92.37,20.67],[92.08,21.19],[92.03,21.7],[91.83,22.18],[91.42,22.77],[90.5,22.81],[90.59,22.39],[90.27,21.84],[89.85,22.04],[89.7,21.86],[89.42,21.97],[89.03,22.06],[88.89,21.69],[88.21,21.7],[86.98,21.5],[87.03,20.74],[86.5,20.15],[85.06,19.48],[83.94,18.3],[83.19,17.67],[82.19,17.02],[82.19,16.56],[81.69,16.31],[80.79,15.95],[80.32,15.9],[80.03,15.14],[80.23,13.84],[80.29,13.01],[79.86,12.06],[79.86,10.36],[79.34,10.31],[78.89,9.55],[79.19,9.22],[78.28,8.93],[77.94,8.25],[77.54,7.97],[76.59,8.9],[76.13,10.3],[75.75,11.31],[75.4,11.78],[74.86,12.74],[74.62,13.99],[74.44,14.62],[73.53,15.99],[73.12,17.93],[72.82,19.21],[72.82,20.42],[72.63,21.36],[71.18,20.76],[70.47,20.88],[69.16,22.09],[69.64,22.45],[69.35,22.84],[68.18,23.69],[67.44,23.94],[67.15,24.66],[66.37,25.43],[64.53,25.24],[62.91,25.22],[61.5,25.08],[59.62,25.38],[58.53,25.61],[57.4,25.74],[56.97,26.97],[56.49,27.14],[55.72,26.96],[54.72,26.48],[53.49,26.81],[52.48,27.58],[51.52,27.87],[50.85,28.81],[50.12,30.15],[49.58,29.99],[48.94,30.32],[48.57,29.93],[47.97,29.98],[48.18,29.53],[48.09,29.31],[48.42,28.55],[48.81,27.69],[49.3,27.46],[49.47,27.11],[50.15,26.69],[50.21,26.28],[50.11,25.94],[50.24,25.61],[50.53,25.33],[50.66,25],[50.81,24.75],[50.74,25.48],[51.01,26.01],[51.29,26.11],[51.59,25.8],[51.61,25.22],[51.39,24.63],[51.58,24.25],[51.76,24.29],[51.79,24.02],[52.58,24.18],[53.4,24.15],[54.01,24.12],[54.69,24.8],[55.44,25.44],[56.07,26.06],[56.36,26.4],[56.49,26.31],[56.39,25.9],[56.26,25.71],[56.4,24.92],[56.85,24.24],[57.4,23.88],[58.14,23.75],[58.73,23.57],[59.18,22.99],[59.45,22.66],[59.81,22.53],[59.81,22.31],[59.44,21.71],[59.28,21.43],[58.86,21.11],[58.49,20.43],[58.03,20.48],[57.83,20.24],[57.67,19.74],[57.79,19.07],[57.69,18.94],[57.23,18.95],[56.61,18.57],[56.51,18.09],[56.28,17.88],[55.66,17.88],[55.27,17.63],[55.27,17.23],[54.79,16.95],[54.24,17.04],[53.57,16.71],[53.11,16.65],[52.39,16.38],[52.19,15.94],[52.17,15.6],[51.17,15.18],[49.57,14.71],[48.68,14],[48.24,13.95],[47.94,14.01],[47.35,13.59],[46.72,13.4],[45.88,13.35],[45.63,13.29],[45.41,13.03],[45.14,12.95],[44.99,12.7],[44.49,12.72],[44.18,12.59],[43.48,12.64],[43.22,13.22],[43.25,13.77],[43.09,14.06],[42.89,14.8],[42.6,15.21],[42.81,15.26],[42.7,15.72],[42.82,15.91],[42.78,16.35],[42.65,16.77],[42.35,17.08],[42.27,17.47],[41.75,17.83],[41.22,18.67],[40.94,19.49],[40.25,20.17],[39.8,20.34],[39.14,21.29],[39.02,21.99],[39.07,22.58],[38.49,23.69],[38.02,24.08],[37.48,24.29],[37.15,24.86],[37.21,25.08],[36.93,25.6],[36.64,25.83],[36.25,26.57],[35.64,27.38],[35.13,28.06],[34.63,28.06],[34.79,28.61],[34.83,28.96],[34.96,29.36],[34.92,29.5],[34.64,29.1],[34.43,28.34],[34.15,27.82],[33.92,27.65],[33.59,27.97],[33.14,28.42],[32.42,29.85],[32.32,29.76],[32.73,28.71],[33.35,27.7],[34.1,26.14],[34.47,25.6],[34.8,25.03],[35.69,23.93],[35.49,23.75],[35.53,23.1],[36.69,22.2],[36.87,22],[37.19,21.02],[36.97,20.84],[37.11,19.81],[37.48,18.61],[37.86,18.37],[38.41,18],[38.99,16.84],[39.27,15.92],[39.81,15.44],[41.18,14.49],[41.73,13.92],[42.28,13.34],[42.59,13],[43.08,12.7],[43.32,12.39],[43.29,11.97],[42.72,11.74],[43.15,11.46],[43.47,11.28],[43.67,10.86],[44.12,10.45],[44.61,10.44],[45.56,10.7],[46.65,10.82],[47.53,11.13],[48.02,11.19],[48.38,11.38],[48.95,11.41],[49.27,11.43],[49.73,11.58],[50.26,11.68],[50.73,12.02],[51.11,12.02],[51.13,11.75],[51.04,11.17],[51.05,10.64],[50.83,10.28],[50.55,9.2],[50.07,8.08],[49.45,6.8],[48.59,5.34],[47.74,4.22],[46.56,2.86],[45.56,2.05],[44.07,1.05],[43.14,0.29],[42.04,-0.92],[41.81,-1.45],[41.59,-1.68],[40.88,-2.08],[40.64,-2.5],[40.26,-2.57],[40.12,-3.28],[39.8,-3.68],[39.6,-4.35],[39.2,-4.68],[38.74,-5.91],[38.8,-6.48],[39.44,-6.84],[39.47,-7.1],[39.19,-7.7],[39.25,-8.01],[39.19,-8.49],[39.54,-9.11],[39.95,-10.1],[40.32,-10.32],[40.48,-10.77],[40.44,-11.76],[40.56,-12.64],[40.6,-14.2],[40.78,-14.69],[40.48,-15.41],[40.09,-16.1],[39.45,-16.72],[38.54,-17.1],[37.41,-17.59],[36.28,-18.66],[35.9,-18.84],[35.2,-19.55],[34.79,-19.78],[34.7,-20.5],[35.18,-21.25],[35.37,-21.84],[35.39,-22.14],[35.56,-22.09],[35.53,-23.07],[35.37,-23.54],[35.61,-23.71],[35.46,-24.12],[35.04,-24.48],[34.22,-24.82],[33.01,-25.36],[32.57,-25.73],[32.66,-26.15],[32.92,-26.22],[32.83,-26.74],[32.58,-27.47],[32.46,-28.3],[32.2,-28.75],[31.52,-29.26],[31.33,-29.4],[30.9,-29.91],[30.62,-30.42],[30.06,-31.14],[28.93,-32.17],[28.22,-32.77],[27.46,-33.23],[26.42,-33.61],[25.91,-33.67],[25.78,-33.94],[25.17,-33.8],[24.68,-33.99],[23.59,-33.79],[22.99,-33.92],[22.57,-33.86],[21.54,-34.26],[20.69,-34.42],[20.07,-34.8],[19.62,-34.82],[19.19,-34.46],[18.86,-34.44],[18.42,-34],[18.38,-34.14],[18.24,-33.87],[18.25,-33.28],[17.93,-32.61],[18.25,-32.43],[18.22,-31.66],[17.57,-30.73],[17.06,-29.88],[17.06,-29.88],[16.34,-28.58],[15.6,-27.82],[15.21,-27.09],[14.99,-26.12],[14.74,-25.39],[14.41,-23.85],[14.39,-22.66],[14.26,-22.11],[13.87,-21.7],[13.35,-20.87],[12.83,-19.67],[12.61,-19.05],[11.79,-18.07],[11.73,-17.3],[11.64,-16.67],[11.78,-15.79],[12.12,-14.88],[12.18,-14.45],[12.5,-13.55],[12.74,-13.14],[13.31,-12.48],[13.63,-12.04],[13.74,-11.3],[13.69,-10.73],[13.39,-10.37],[13.12,-9.77],[12.88,-9.17],[12.93,-8.96],[13.24,-8.56],[12.93,-7.6],[12.73,-6.93],[12.23,-6.29],[12.32,-6.1],[12.18,-5.79],[11.91,-5.04],[11.09,-3.98],[10.07,-2.97],[9.41,-2.14],[8.8,-1.11],[8.83,-0.78],[9.05,-0.46],[9.29,0.27],[9.49,1.01],[9.31,1.16],[9.65,2.28],[9.8,3.07],[9.4,3.73],[8.95,3.9],[8.74,4.35],[8.49,4.5],[8.5,4.77],[7.46,4.41],[7.08,4.46],[6.7,4.24],[5.9,4.26],[5.36,4.89],[5.03,5.61],[4.33,6.27],[3.57,6.26],[2.69,6.26],[1.87,6.14],[1.06,5.93],[-0.51,5.34],[-1.06,5],[-1.96,4.71],[-2.86,4.99],[-3.31,4.98],[-4.01,5.18],[-4.65,5.17],[-5.83,4.99],[-6.53,4.71],[-7.52,4.34],[-7.71,4.36],[-7.97,4.36],[-9,4.83],[-9.91,5.59],[-10.77,6.14],[-11.44,6.79],[-11.71,6.86],[-12.43,7.26],[-12.95,7.8],[-13.12,8.16],[-13.25,8.9],[-13.69,9.49],[-14.07,9.89],[-14.33,10.02],[-14.58,10.21],[-14.69,10.66],[-14.84,10.88],[-15.13,11.04],[-15.66,11.46],[-16.09,11.52],[-16.31,11.81],[-16.31,11.96],[-16.61,12.17],[-16.68,12.38],[-16.84,13.15],[-16.71,13.59],[-17.13,14.37],[-17.63,14.73],[-17.19,14.92],[-16.7,15.62],[-16.46,16.14],[-16.55,16.67],[-16.27,17.17],[-16.15,18.11],[-16.26,19.1]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[125.95,-8.43],[126.64,-8.4],[126.96,-8.27],[127.34,-8.4],[126.97,-8.67],[125.93,-9.11],[125.09,-9.39],[124.44,-10.14],[123.58,-10.36],[123.46,-10.24],[123.55,-9.9],[123.98,-9.29],[124.97,-8.89],[125.09,-8.66],[125.95,-8.43]]},"properties":{"scalerank":1,"featurecla":"Coastline","min_zoom":1}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-180,-84.71],[-179.94,-84.72],[-179.06,-84.14],[-177.26,-84.45],[-176.08,-84.1],[-175.83,-84.12],[-174.38,-84.53],[-173.12,-84.12],[-172.89,-84.06],[-169.95,-83.88],[-169,-84.12],[-168.53,-84.24],[-167.02,-84.57],[-164.18,-84.83],[-161.93,-85.14],[-158.07,-85.37],[-155.19,-85.1],[-150.94,-85.3],[-148.53,-85.61],[-145.89,-85.32],[-143.11,-85.04],[-142.89,-84.57],[-146.83,-84.53],[-150.06,-84.3],[-150.9,-83.9],[-153.59,-83.69],[-153.41,-83.24],[-153.04,-82.83],[-152.67,-82.45],[-152.86,-82.04],[-154.53,-81.77],[-155.29,-81.42],[-156.84,-81.1],[-154.41,-81.16],[-152.1,-81],[-150.65,-81.34],[-148.87,-81.04],[-147.22,-80.67],[-146.42,-80.34],[-146.77,-79.93],[-148.06,-79.65],[-149.53,-79.36],[-151.59,-79.3],[-153.39,-79.16],[-155.33,-79.06],[-155.98,-78.69],[-157.27,-78.38],[-158.05,-78.03],[-158.37,-76.89],[-157.88,-76.99],[-156.97,-77.3],[-155.33,-77.2],[-153.74,-77.07],[-152.92,-77.5],[-151.33,-77.4],[-150,-77.18],[-148.75,-76.91],[-147.61,-76.58],[-146.1,-76.48],[-146.14,-76.11],[-146.5,-75.73],[-146.2,-75.38],[-144.91,-75.2],[-144.32,-75.54],[-142.79,-75.34],[-141.64,-75.09],[-140.21,-75.07],[-138.86,-74.97],[-137.51,-74.73],[-136.43,-74.52],[-135.21,-74.3],[-134.43,-74.36],[-133.75,-74.44],[-132.26,-74.3],[-130.93,-74.48],[-129.55,-74.46],[-128.24,-74.32],[-126.89,-74.42],[-125.4,-74.52],[-124.01,-74.48],[-122.56,-74.5],[-121.07,-74.52],[-119.7,-74.48],[-118.68,-74.19],[-117.47,-74.03],[-116.22,-74.24],[-115.02,-74.07],[-113.94,-73.71],[-113.3,-74.03],[-112.95,-74.38],[-112.3,-74.71],[-111.26,-74.42],[-110.07,-74.79],[-108.71,-74.91],[-107.56,-75.18],[-106.15,-75.13],[-104.88,-74.95],[-103.37,-74.99],[-102.02,-75.13],[-100.65,-75.3],[-100.12,-74.87],[-100.76,-74.54],[-101.25,-74.19],[-102.55,-74.11],[-103.11,-73.73],[-103.33,-73.36],[-103.68,-72.62],[-102.92,-72.75],[-101.61,-72.81],[-100.31,-72.75],[-99.14,-72.91],[-98.12,-73.21],[-97.69,-73.56],[-96.34,-73.62],[-95.04,-73.48],[-93.67,-73.28],[-92.44,-73.17],[-91.42,-73.4],[-90.09,-73.32],[-89.23,-72.56],[-88.42,-73.01],[-87.27,-73.19],[-86.01,-73.09],[-85.19,-73.48],[-83.88,-73.52],[-82.67,-73.64],[-81.47,-73.85],[-80.69,-73.48],[-80.3,-73.13],[-79.3,-73.52],[-77.93,-73.42],[-76.91,-73.64],[-76.22,-73.97],[-74.89,-73.87],[-73.85,-73.66],[-72.83,-73.4],[-71.62,-73.26],[-70.21,-73.15],[-68.94,-73.01],[-67.96,-72.79],[-67.37,-72.48],[-67.13,-72.05],[-67.25,-71.64],[-67.56,-71.25],[-67.92,-70.85],[-68.23,-70.46],[-68.49,-70.11],[-68.54,-69.72],[-68.45,-69.33],[-67.98,-68.95],[-67.58,-68.54],[-67.43,-68.15],[-67.62,-67.72],[-67.74,-67.33],[-67.25,-66.88],[-66.7,-66.58],[-66.06,-66.21],[-65.37,-65.9],[-64.57,-65.6],[-64.18,-65.17],[-63.63,-64.9],[-63,-64.64],[-62.04,-64.58],[-61.41,-64.27],[-60.71,-64.07],[-59.89,-63.96],[-59.16,-63.7],[-58.59,-63.39],[-57.81,-63.27],[-57.22,-63.53],[-57.6,-63.86],[-58.61,-64.15],[-59.05,-64.37],[-59.79,-64.21],[-60.61,-64.31],[-61.3,-64.54],[-62.02,-64.8],[-62.51,-65.09],[-62.65,-65.48],[-62.59,-65.86],[-62.12,-66.19],[-62.81,-66.43],[-63.75,-66.5],[-64.29,-66.84],[-64.88,-67.15],[-65.51,-67.58],[-65.67,-67.95],[-65.31,-68.37],[-64.78,-68.68],[-63.96,-68.91],[-63.2,-69.23],[-62.79,-69.62],[-62.57,-69.99],[-62.28,-70.38],[-61.81,-70.72],[-61.51,-71.09],[-61.38,-72.01],[-61.08,-72.38],[-61,-72.77],[-60.69,-73.17],[-60.83,-73.7],[-61.38,-74.11],[-61.96,-74.44],[-63.3,-74.58],[-63.75,-74.93],[-64.35,-75.26],[-65.86,-75.64],[-67.19,-75.79],[-68.45,-76.01],[-69.8,-76.22],[-70.6,-76.63],[-72.21,-76.67],[-73.97,-76.63],[-75.56,-76.71],[-77.24,-76.71],[-76.93,-77.1],[-75.4,-77.28],[-74.28,-77.56],[-73.66,-77.91],[-74.77,-78.22],[-76.5,-78.12],[-77.93,-78.38],[-77.98,-78.79],[-78.02,-79.18],[-76.85,-79.51],[-76.63,-79.89],[-75.36,-80.26],[-73.24,-80.42],[-71.44,-80.69],[-70.01,-81],[-68.19,-81.32],[-65.7,-81.47],[-63.26,-81.75],[-61.55,-82.04],[-59.69,-82.38],[-58.71,-82.85],[-58.22,-83.22],[-57.01,-82.87],[-55.36,-82.57],[-53.62,-82.26],[-51.54,-82],[-49.76,-81.73],[-47.27,-81.71],[-44.83,-81.85],[-42.81,-82.08],[-42.16,-81.65],[-40.77,-81.36],[-38.24,-81.34],[-36.27,-81.12],[-34.39,-80.91],[-32.31,-80.77],[-30.1,-80.59],[-28.55,-80.34],[-29.25,-79.99],[-29.69,-79.63],[-29.69,-79.26],[-31.62,-79.3],[-33.68,-79.46],[-35.64,-79.46],[-35.91,-79.08],[-35.78,-78.34],[-35.33,-78.12],[-33.9,-77.89],[-32.21,-77.65],[-31,-77.36],[-29.78,-77.07],[-28.88,-76.67],[-27.51,-76.5],[-26.16,-76.36],[-25.47,-76.28],[-23.93,-76.24],[-22.46,-76.11],[-21.22,-75.91],[-20.01,-75.67],[-18.91,-75.44],[-17.52,-75.13],[-16.64,-74.79],[-15.7,-74.5],[-15.41,-74.11],[-16.47,-73.87],[-16.11,-73.46],[-15.45,-73.15],[-14.41,-72.95],[-13.31,-72.72],[-12.29,-72.4],[-11.51,-72.01],[-11.02,-71.54],[-10.3,-71.27],[-9.1,-71.32],[-8.61,-71.66],[-7.42,-71.7],[-7.38,-71.32],[-6.87,-70.93],[-5.79,-71.03],[-5.54,-71.4],[-4.34,-71.46],[-3.05,-71.29],[-1.8,-71.17],[-0.66,-71.23],[-0.23,-71.64],[0.87,-71.3],[1.89,-71.13],[3.02,-70.99],[4.14,-70.85],[5.16,-70.62],[6.27,-70.46],[7.14,-70.25],[7.74,-69.89],[8.49,-70.15],[9.53,-70.01],[10.25,-70.48],[10.82,-70.83],[11.95,-70.64],[12.4,-70.25],[13.42,-69.97],[14.73,-70.03],[15.13,-70.4],[15.95,-70.03],[17.03,-69.91],[18.2,-69.87],[19.26,-69.89],[20.38,-70.01],[21.45,-70.07],[21.92,-70.4],[22.57,-70.7],[23.67,-70.52],[24.84,-70.48],[25.98,-70.48],[27.09,-70.46],[28.09,-70.32],[29.15,-70.21],[30.03,-69.93],[30.97,-69.76],[31.99,-69.66],[32.75,-69.38],[33.3,-68.84],[33.87,-68.5],[34.91,-68.66],[35.3,-69.01],[36.16,-69.25],[37.2,-69.17],[37.91,-69.52],[38.65,-69.78],[39.67,-69.54],[40.02,-69.11],[40.92,-68.93],[41.96,-68.6],[42.94,-68.46],[44.11,-68.27],[44.9,-68.05],[45.72,-67.82],[46.5,-67.6],[47.44,-67.72],[48.34,-67.37],[48.99,-67.09],[49.93,-67.11],[50.75,-66.88],[50.95,-66.52],[51.79,-66.25],[52.61,-66.05],[53.61,-65.9],[54.53,-65.82],[55.41,-65.88],[56.36,-65.97],[57.16,-66.25],[57.26,-66.68],[58.14,-67.01],[58.74,-67.29],[59.94,-67.41],[60.61,-67.68],[61.43,-67.95],[62.39,-68.01],[63.19,-67.82],[64.05,-67.41],[64.99,-67.62],[65.97,-67.74],[66.91,-67.86],[67.89,-67.93],[68.89,-67.93],[69.71,-68.97],[69.67,-69.23],[69.56,-69.68],[68.6,-69.93],[67.81,-70.31],[67.95,-70.7],[69.07,-70.68],[68.93,-71.07],[68.42,-71.44],[67.95,-71.85],[68.71,-72.17],[69.87,-72.26],[71.02,-72.09],[71.57,-71.7],[71.91,-71.32],[72.45,-71.01],[73.08,-70.72],[73.34,-70.36],[73.86,-69.87],[74.49,-69.78],[75.63,-69.74],[76.63,-69.62],[77.64,-69.46],[78.13,-69.07],[78.43,-68.7],[79.11,-68.33],[80.09,-68.07],[80.94,-67.88],[81.48,-67.54],[82.05,-67.37],[82.78,-67.21],[83.78,-67.31],[84.68,-67.21],[85.66,-67.09],[86.75,-67.15],[87.48,-66.88],[87.99,-66.21],[88.36,-66.48],[88.83,-66.95],[89.67,-67.15],[90.63,-67.23],[91.59,-67.11],[92.61,-67.19],[93.55,-67.21],[94.18,-67.11],[95.02,-67.17],[95.78,-67.39],[96.68,-67.25],[97.76,-67.25],[98.68,-67.11],[99.72,-67.25],[100.38,-66.92],[100.89,-66.58],[101.58,-66.31],[102.83,-65.56],[103.48,-65.7],[104.24,-65.97],[104.91,-66.33],[106.18,-66.93],[107.16,-66.95],[108.08,-66.95],[109.16,-66.84],[110.24,-66.7],[111.06,-66.43],[111.74,-66.13],[112.86,-66.09],[113.6,-65.88],[114.39,-66.07],[114.9,-66.39],[115.6,-66.7],[116.7,-66.66],[117.38,-66.92],[118.58,-67.17],[119.83,-67.27],[120.87,-67.19],[121.65,-66.88],[122.32,-66.56],[123.22,-66.48],[124.12,-66.62],[125.16,-66.72],[126.1,-66.56],[127,-66.56],[127.88,-66.66],[128.8,-66.76],[129.7,-66.58],[130.78,-66.43],[131.8,-66.39],[132.94,-66.39],[133.86,-66.29],[134.76,-66.21],[135.03,-65.72],[135.07,-65.31],[135.7,-65.58],[135.87,-66.03],[136.21,-66.45],[136.62,-66.78],[137.46,-66.95],[138.6,-66.9],[139.91,-66.88],[140.81,-66.82],[142.12,-66.82],[143.06,-66.8],[144.37,-66.84],[145.49,-66.92],[146.2,-67.23],[146,-67.6],[146.65,-67.9],[147.72,-68.13],[148.84,-68.39],[150.13,-68.56],[151.48,-68.72],[152.5,-68.87],[153.64,-68.89],[154.28,-68.56],[155.17,-68.84],[155.93,-69.15],[156.81,-69.38],[158.03,-69.48],[159.18,-69.6],[159.67,-69.99],[160.81,-70.23],[161.57,-70.58],[162.69,-70.74],[163.84,-70.72],[164.92,-70.78],[166.11,-70.76],[167.31,-70.83],[168.43,-70.97],[169.46,-71.21],[170.5,-71.4],[171.21,-71.7],[171.09,-72.09],[170.56,-72.44],[170.11,-72.89],[169.76,-73.24],[169.29,-73.66],[167.98,-73.81],[167.39,-74.17],[166.09,-74.38],[165.64,-74.77],[164.96,-75.15],[164.23,-75.46],[163.82,-75.87],[163.57,-76.24],[163.47,-76.69],[163.49,-77.07],[164.06,-77.46],[164.27,-77.83],[164.74,-78.18],[166.6,-78.32],[167,-78.75],[165.19,-78.91],[163.67,-79.12],[161.77,-79.16],[160.92,-79.73],[160.75,-80.2],[160.32,-80.57],[159.79,-80.95],[161.12,-81.28],[161.63,-81.69],[162.49,-82.06],[163.71,-82.4],[165.1,-82.71],[166.6,-83.02],[168.9,-83.34],[169.4,-83.83],[172.28,-84.04],[172.48,-84.12],[173.22,-84.41],[175.99,-84.16],[178.28,-84.47],[180,-84.71]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-180,68.96],[-177.55,68.2],[-174.93,67.21],[-175.01,66.58],[-174.34,66.34],[-174.57,67.06],[-171.86,66.91],[-169.9,65.98],[-170.89,65.54],[-172.53,65.44],[-172.56,64.46],[-172.96,64.25],[-173.89,64.28],[-174.65,64.63],[-175.98,64.92],[-176.21,65.36],[-177.22,65.52],[-178.36,65.39],[-178.9,65.74],[-178.69,66.11],[-179.88,65.87],[-179.43,65.4],[-180,64.98]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-155.4,20.08],[-155.22,19.99],[-155.06,19.86],[-154.81,19.51],[-154.83,19.45],[-155.22,19.24],[-155.54,19.08],[-155.69,18.92],[-155.94,19.06],[-155.91,19.34],[-156.07,19.7],[-156.02,19.81],[-155.85,19.98],[-155.92,20.17],[-155.86,20.27],[-155.79,20.25],[-155.4,20.08]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0.5}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[46.68,44.61],[47.68,45.64],[48.65,45.81],[49.1,46.4],[50.03,46.61],[51.19,47.05],[52.04,46.8],[53.04,46.85],[53.22,46.23],[53.04,45.26],[52.17,45.41],[51.32,45.25],[51.28,44.51],[50.31,44.61],[50.34,44.28],[50.89,44.03],[51.34,43.13],[52.5,42.79],[52.69,42.44],[52.45,42.03],[52.5,41.78],[52.81,41.14],[52.92,41.87],[53.72,42.12],[54.01,41.55],[54.74,40.95],[53.86,40.63],[52.92,40.88],[52.69,40.03],[53.36,39.98],[53.1,39.29],[53.88,38.95],[53.74,37.91],[53.92,37.2],[53.83,36.97],[52.26,36.7],[50.84,36.87],[50.15,37.37],[49.2,37.58],[48.88,38.32],[48.86,38.82],[49.22,39.05],[49.4,39.4],[49.57,40.18],[50.39,40.26],[50.08,40.53],[49.62,40.57],[49.11,41.28],[48.58,41.81],[47.49,42.99],[47.59,43.66],[46.68,44.61]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-106.52,73.08],[-105.4,72.67],[-104.77,71.7],[-104.46,70.99],[-102.79,70.5],[-100.98,70.02],[-101.09,69.58],[-102.73,69.5],[-102.09,69.12],[-102.43,68.75],[-104.24,68.91],[-105.96,69.18],[-107.12,69.12],[-109,68.78],[-111.53,68.63],[-113.31,68.54],[-113.85,69.01],[-115.22,69.28],[-116.11,69.17],[-117.34,69.96],[-116.67,70.07],[-115.13,70.24],[-113.72,70.19],[-112.42,70.37],[-114.35,70.6],[-116.49,70.52],[-117.9,70.54],[-118.43,70.91],[-116.11,71.31],[-117.66,71.3],[-119.4,71.56],[-118.56,72.31],[-117.87,72.71],[-115.19,73.31],[-114.17,73.12],[-114.67,72.65],[-112.44,72.96],[-111.05,72.45],[-109.92,72.96],[-109.01,72.63],[-108.19,71.65],[-107.69,72.07],[-108.4,73.09],[-107.52,73.24],[-106.52,73.08]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-96.02,80.6],[-95.32,80.91],[-94.3,80.98],[-94.74,81.21],[-92.41,81.26],[-91.13,80.72],[-87.81,80.32],[-87.02,79.66],[-85.81,79.34],[-87.19,79.04],[-89.04,78.29],[-90.8,78.22],[-92.88,78.34],[-93.95,78.75],[-93.94,79.11],[-93.15,79.38],[-94.97,79.37],[-96.08,79.71],[-96.71,80.16],[-96.02,80.6]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-91.59,81.89],[-90.1,82.09],[-88.93,82.12],[-86.97,82.28],[-85.5,82.65],[-84.26,82.6],[-83.18,82.32],[-82.42,82.86],[-81.1,83.02],[-79.31,83.13],[-76.25,83.17],[-75.72,83.06],[-72.83,83.23],[-70.67,83.17],[-68.5,83.11],[-65.83,83.03],[-63.68,82.9],[-61.85,82.63],[-61.89,82.36],[-64.33,81.93],[-66.75,81.73],[-67.66,81.5],[-65.48,81.51],[-67.84,80.9],[-69.47,80.62],[-71.18,79.8],[-73.24,79.63],[-73.88,79.43],[-76.91,79.32],[-75.53,79.2],[-76.22,79.02],[-75.39,78.53],[-76.34,78.18],[-77.89,77.9],[-78.36,77.51],[-79.76,77.21],[-79.62,76.98],[-77.91,77.02],[-77.89,76.78],[-80.56,76.18],[-83.17,76.45],[-86.11,76.3],[-87.6,76.42],[-89.49,76.47],[-89.62,76.95],[-87.77,77.18],[-88.26,77.9],[-87.65,77.97],[-84.98,77.54],[-86.34,78.18],[-87.96,78.37],[-87.15,78.76],[-85.38,79],[-85.09,79.35],[-86.51,79.74],[-86.93,80.25],[-84.2,80.21],[-83.41,80.1],[-81.85,80.46],[-84.1,80.58],[-87.6,80.52],[-89.37,80.86],[-90.2,81.26],[-91.37,81.55],[-91.59,81.89]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}},
+{"type":"Feature","geometry":{"type":"LineString","coordinates":[[-46.76,82.63],[-43.41,83.23],[-39.9,83.18],[-38.62,83.55],[-35.09,83.65],[-27.1,83.52],[-20.85,82.73],[-22.69,82.34],[-26.52,82.3],[-31.9,82.2],[-31.4,82.02],[-27.86,82.13],[-24.84,81.79],[-22.9,82.09],[-22.07,81.73],[-23.17,81.15],[-20.62,81.52],[-15.77,81.91],[-12.77,81.72],[-12.21,81.29],[-16.29,80.58],[-16.85,80.35],[-20.05,80.18],[-17.73,80.13],[-18.9,79.4],[-19.7,78.75],[-19.67,77.64],[-18.47,76.99],[-20.04,76.94],[-21.68,76.63],[-19.83,76.1],[-19.6,75.25],[-20.67,75.16],[-19.37,74.3],[-21.59,74.22],[-20.43,73.82],[-20.76,73.46],[-22.17,73.31],[-23.57,73.31],[-22.31,72.63],[-22.3,72.18],[-24.28,72.6],[-24.79,72.33],[-23.44,72.08],[-22.13,71.47],[-21.75,70.66],[-23.54,70.47],[-24.31,70.86],[-25.54,71.43],[-25.2,70.75],[-26.36,70.23],[-23.73,70.18],[-22.35,70.13],[-25.03,69.26],[-27.75,68.47],[-30.67,68.13],[-31.78,68.12],[-32.81,67.74],[-34.2,66.68],[-36.35,65.98],[-37.04,65.94],[-38.38,65.69],[-39.81,65.46],[-40.67,64.84],[-40.68,64.14],[-41.19,63.48],[-42.82,62.68],[-42.42,61.9],[-42.87,61.07],[-43.38,60.1],[-44.79,60.04],[-46.26,60.85],[-48.26,60.86],[-49.23,61.41],[-49.9,62.38],[-51.63,63.63],[-52.14,64.28],[-52.28,65.18],[-53.66,66.1],[-53.3,66.84],[-53.97,67.19],[-52.98,68.36],[-51.48,68.73],[-51.08,69.15],[-50.87,69.93],[-52.01,69.57],[-52.56,69.43],[-53.46,69.28],[-54.68,69.61],[-54.75,70.29],[-54.36,70.82],[-53.43,70.84],[-51.39,70.57],[-53.11,71.2],[-54,71.55],[-55,71.41],[-55.83,71.65],[-54.72,72.59],[-55.33,72.96],[-56.12,73.65],[-57.32,74.71],[-58.6,75.1],[-58.59,75.52],[-61.27,76.1],[-63.39,76.18],[-66.06,76.13],[-68.5,76.06],[-69.66,76.38],[-71.4,77.01],[-68.78,77.32],[-66.76,77.38],[-71.04,77.64],[-73.3,78.04],[-73.16,78.43],[-69.37,78.91],[-65.71,79.39],[-65.32,79.76],[-68.02,80.12],[-67.15,80.52],[-63.69,81.21],[-62.23,81.32],[-62.65,81.77],[-60.28,82.03],[-57.21,82.19],[-54.13,82.2],[-53.04,81.89],[-50.39,82.44],[-48,82.06],[-46.6,81.99],[-44.52,81.66],[-46.9,82.2],[-46.76,82.63]]},"properties":{"scalerank":0,"featurecla":"Coastline","min_zoom":0}}
+]}
diff --git a/src/anemoi/training/diagnostics/logger.py b/src/anemoi/training/diagnostics/logger.py
new file mode 100644
index 00000000..599dad36
--- /dev/null
+++ b/src/anemoi/training/diagnostics/logger.py
@@ -0,0 +1,158 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
+
+if TYPE_CHECKING:
+    import pytorch_lightning as pl
+
+LOGGER = logging.getLogger(__name__)
+
+
+def get_mlflow_logger(config: DictConfig) -> None:
+    if not config.diagnostics.log.mlflow.enabled:
+        LOGGER.debug("MLFlow logging is disabled.")
+        return None
+
+    from anemoi.training.diagnostics.mlflow.logger import AnemoiMLflowLogger
+    from anemoi.training.diagnostics.mlflow.logger import get_mlflow_run_params
+
+    resumed = config.training.run_id is not None
+    forked = config.training.fork_run_id is not None
+
+    save_dir = config.hardware.paths.logs.mlflow
+    tracking_uri = config.diagnostics.log.mlflow.tracking_uri
+    offline = config.diagnostics.log.mlflow.offline
+
+    if (resumed or forked) and (offline):  # when resuming or forking offline -
+        # tracking_uri = ${hardware.paths.logs.mlflow}
+        tracking_uri = save_dir
+    # create directory if it does not exist
+    Path(config.hardware.paths.logs.mlflow).mkdir(parents=True, exist_ok=True)
+    run_id, run_name, tags = get_mlflow_run_params(config, tracking_uri)
+
+    log_hyperparams = True
+    if resumed and not config.diagnostics.log.mlflow.on_resume_create_child:
+        LOGGER.info(
+            (
+                "Resuming run without creating child run - MLFlow logs will not update the"
+                "initial runs hyperparameters with those of the resumed run."
+                "To update the initial run's hyperparameters, set "
+                "`diagnostics.log.mlflow.on_resume_create_child: True`."
+            ),
+        )
+        log_hyperparams = False
+
+    logger = AnemoiMLflowLogger(
+        experiment_name=config.diagnostics.log.mlflow.experiment_name,
+        tracking_uri=tracking_uri,
+        save_dir=save_dir,
+        run_name=run_name,
+        run_id=run_id,
+        log_model=config.diagnostics.log.mlflow.log_model,
+        offline=offline,
+        tags=tags,
+        resumed=resumed,
+        forked=forked,
+        log_hyperparams=log_hyperparams,
+        authentication=config.diagnostics.log.mlflow.authentication,
+    )
+    config_params = OmegaConf.to_container(config, resolve=True)
+
+    logger.log_hyperparams(config_params)
+
+    if config.diagnostics.log.mlflow.terminal:
+        logger.log_terminal_output(artifact_save_dir=config.hardware.paths.plots)
+    if config.diagnostics.log.mlflow.system:
+        logger.log_system_metrics()
+
+    return logger
+
+
+def get_tensorboard_logger(config: DictConfig) -> pl.loggers.TensorBoardLogger | None:
+    """Setup TensorBoard experiment logger.
+
+    Parameters
+    ----------
+    config : DictConfig
+        Job configuration
+
+    Returns
+    -------
+    Optional[pl.loggers.TensorBoardLogger]
+        Logger object, or None
+
+    """
+    if not config.diagnostics.log.tensorboard.enabled:
+        LOGGER.debug("Tensorboard logging is disabled.")
+        return None
+
+    from pytorch_lightning.loggers import TensorBoardLogger
+
+    return TensorBoardLogger(
+        save_dir=config.hardware.paths.logs.tensorboard,
+        log_graph=False,
+    )
+
+
+def get_wandb_logger(config: DictConfig, model: pl.LightningModule) -> pl.loggers.WandbLogger | None:
+    """Setup Weights & Biases experiment logger.
+
+    Parameters
+    ----------
+    config : DictConfig
+        Job configuration
+    model: GraphForecaster
+        Model to watch
+
+    Returns
+    -------
+    Optional[pl.loggers.WandbLogger]
+        Logger object
+
+    Raises
+    ------
+    ImportError
+        If `wandb` is not installed
+
+    """
+    if not config.diagnostics.log.wandb.enabled:
+        LOGGER.debug("Weights & Biases logging is disabled.")
+        return None
+
+    try:
+        from pytorch_lightning.loggers.wandb import WandbLogger
+    except ImportError as err:
+        msg = "To activate W&B logging, please install `wandb` as an optional dependency."
+        raise ImportError(msg) from err
+
+    logger = WandbLogger(
+        project=config.diagnostics.log.wandb.project,
+        entity=config.diagnostics.log.wandb.entity,
+        id=config.training.run_id,
+        save_dir=config.hardware.paths.logs.wandb,
+        offline=config.diagnostics.log.wandb.offline,
+        log_model=config.diagnostics.log.wandb.log_model,
+        resume=config.training.run_id is not None,
+    )
+    logger.log_hyperparams(OmegaConf.to_container(config, resolve=True))
+    if config.diagnostics.log.wandb.gradients or config.diagnostics.log.wandb.parameters:
+        if config.diagnostics.log.wandb.gradients and config.diagnostics.log.wandb.parameters:
+            log_ = "all"
+        elif config.diagnostics.log.wandb.gradients:
+            log_ = "gradients"
+        else:
+            log_ = "parameters"
+        logger.watch(model, log=log_, log_freq=config.diagnostics.log.interval, log_graph=False)
+
+    return logger
diff --git a/src/anemoi/training/diagnostics/maps.py b/src/anemoi/training/diagnostics/maps.py
new file mode 100644
index 00000000..265d814b
--- /dev/null
+++ b/src/anemoi/training/diagnostics/maps.py
@@ -0,0 +1,99 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+import copy
+import json
+import logging
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.collections import LineCollection
+
+from anemoi.training import diagnostics
+
+LOGGER = logging.getLogger(__name__)
+
+
+class EquirectangularProjection:
+    """Class to convert lat/lon coordinates to Equirectangular coordinates."""
+
+    def __init__(self) -> None:
+        """Initialise the EquirectangularProjection object with offset."""
+        self.x_offset = 0.0
+        self.y_offset = 0.0
+
+    def __call__(self, lon: np.ndarray, lat: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        lon_rad = np.radians(lon)
+        lat_rad = np.radians(lat)
+        x = [v - 2 * np.pi if v > np.pi else v for v in lon_rad]
+        y = lat_rad
+        return x, y
+
+    @staticmethod
+    def inverse(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        lon = np.degrees(x)
+        lat = np.degrees(y)
+        return lon, lat
+
+
+class Coastlines:
+    """Class to plot coastlines from a GeoJSON file."""
+
+    def __init__(self, projection: EquirectangularProjection = None) -> None:
+        """Initialise the Coastlines object.
+
+        Parameters
+        ----------
+        projection : Any, optional
+            Projection Object, by default None
+
+        Raises
+        ------
+        ModuleNotFoundError
+            Whether the importlib_resources or importlib.resources module is not found.
+
+        """
+        try:
+            # this requires python 3.9 or newer
+            from importlib.resources import files
+        except ImportError:
+            try:
+                from importlib_resources import files
+            except ModuleNotFoundError as e:
+                msg = "Please install importlib_resources on Python <=3.8."
+                raise ModuleNotFoundError(msg) from e
+
+        # Get the path to "continents.json" within your library
+        self.continents_file = files(diagnostics) / "continents.json"
+
+        # Load GeoJSON data from the file
+        with self.continents_file.open("rt") as file:
+            self.data = json.load(file)
+
+        if projection is None:
+            self.projection = EquirectangularProjection()
+
+        self.process_data()
+
+    # Function to extract LineString coordinates
+    @staticmethod
+    def extract_coordinates(feature: dict) -> list:
+        return feature["geometry"]["coordinates"]
+
+    def process_data(self) -> None:
+        lines = []
+        for feature in self.data["features"]:
+            coordinates = self.extract_coordinates(feature)
+            x, y = zip(*coordinates)  # Unzip the coordinates into separate x and y lists
+
+            lines.append(list(zip(*self.projection(x, y))))  # Convert lat/lon to Cartesian coordinates
+        self.lines = LineCollection(lines, linewidth=0.5, color="black")
+
+    def plot_continents(self, ax: plt.Axes) -> None:
+        # Add the lines to the axis as a collection
+        # Note that we have to provide a copy of the lines, because of Matplotlib
+        ax.add_collection(copy.copy(self.lines))
diff --git a/src/anemoi/training/diagnostics/mlflow/__init__.py b/src/anemoi/training/diagnostics/mlflow/__init__.py
new file mode 100644
index 00000000..282d6a69
--- /dev/null
+++ b/src/anemoi/training/diagnostics/mlflow/__init__.py
@@ -0,0 +1,6 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
diff --git a/src/anemoi/training/diagnostics/mlflow/auth.py b/src/anemoi/training/diagnostics/mlflow/auth.py
new file mode 100644
index 00000000..7189ea21
--- /dev/null
+++ b/src/anemoi/training/diagnostics/mlflow/auth.py
@@ -0,0 +1,224 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 __future__ import annotations
+
+import logging
+import os
+import time
+from datetime import datetime
+from datetime import timezone
+from functools import wraps
+from getpass import getpass
+from typing import Callable
+
+import requests
+from anemoi.utils.config import load_config
+from anemoi.utils.config import save_config
+from anemoi.utils.timer import Timer
+from requests.exceptions import HTTPError
+
+REFRESH_EXPIRE_DAYS = 29
+
+
+class TokenAuth:
+    """Manage authentication with a keycloak token server."""
+
+    config_file = "mlflow-token.json"
+
+    def __init__(self, url: str, enabled: bool = True, target_env_var: str = "MLFLOW_TRACKING_TOKEN") -> None:
+        """Initialise the token authentication object.
+
+        Parameters
+        ----------
+        url : str
+            URL of the authentication server.
+        enabled : bool, optional
+            Set this to False to turn off authentication, by default True
+        target_env_var : str, optional
+            The environment variable to store the access token in after authenticating,
+            by default `MLFLOW_TRACKING_TOKEN`
+
+        """
+        self.url = url
+        self.target_env_var = target_env_var
+        self._enabled = enabled
+
+        config = self.load_config()
+
+        self._refresh_token = config.get("refresh_token")
+        self.refresh_expires = config.get("refresh_expires", 0)
+        self.access_token = None
+        self.access_expires = 0
+
+        # the command line tool adds a default handler to the root logger on runtime,
+        # so we init our logger here (on runtime, not on import) to avoid duplicate handlers
+        self.log = logging.getLogger(__name__)
+
+    def __call__(self) -> None:
+        self.authenticate()
+
+    @property
+    def refresh_token(self) -> str:
+        return self._refresh_token
+
+    @refresh_token.setter
+    def refresh_token(self, value: str) -> None:
+        self._refresh_token = value
+        self.refresh_expires = time.time() + (REFRESH_EXPIRE_DAYS * 86400)  # 86400 seconds in a day
+
+    @staticmethod
+    def load_config() -> dict:
+        return load_config(TokenAuth.config_file)
+
+    def enabled(fn: Callable) -> Callable:  # noqa: N805
+        """Decorator to call or ignore a function based on the `enabled` flag."""
+
+        @wraps(fn)
+        def _wrapper(self: TokenAuth, *args, **kwargs) -> Callable | None:
+            if self._enabled:
+                return fn(self, *args, **kwargs)
+            return None
+
+        return _wrapper
+
+    @enabled
+    def login(self, force_credentials: bool = False, **kwargs: dict) -> None:
+        """Acquire a new refresh token and save it to disk.
+
+        If an existing valid refresh token is already on disk it will be used.
+        If not, or the token has expired, the user will be asked to obtain one from the API.
+
+        Refresh token expiry time is set in the `REFRESH_EXPIRE_DAYS` constant (default 29 days).
+
+        This function should be called once, interactively, right before starting a training run.
+
+        Parameters
+        ----------
+        force_credentials : bool, optional
+            Force a credential login even if a refreh token is available, by default False.
+        kwargs : dict
+            Additional keyword arguments.
+
+        Raises
+        ------
+        RuntimeError
+            A new refresh token could not be acquired.
+
+        """
+        del kwargs  # unused
+        self.log.info("🌐 Logging in to %s", self.url)
+        new_refresh_token = None
+
+        if not force_credentials and self.refresh_token and self.refresh_expires > time.time():
+            new_refresh_token = self._token_request(ignore_exc=True).get("refresh_token")
+
+        if not new_refresh_token:
+            self.log.info("📝 Please obtain a seed refresh token from %s/seed", self.url)
+            self.log.info("📝 and paste it here (you will not see the output, just press enter after pasting):")
+            self.refresh_token = getpass("Refresh Token: ")
+
+            # perform a new refresh token request to check if the seed refresh token is valid
+            new_refresh_token = self._token_request().get("refresh_token")
+
+        if not new_refresh_token:
+            msg = "❌ Failed to log in. Please try again."
+            raise RuntimeError(msg)
+
+        self.refresh_token = new_refresh_token
+        self.save()
+
+        self.log.info("✅ Successfully logged in to MLflow. Happy logging!")
+
+    @enabled
+    def authenticate(self, **kwargs: dict) -> None:
+        """Check the access token and refresh it if necessary. A new refresh token will also be acquired upon refresh.
+
+        This requires a valid refresh token to be available, obtained from the `login` method.
+
+        The access token is stored in memory and in an environment variable.
+        If the access token is still valid, this function does nothing.
+
+        This function should be called before every MLflow API request.
+
+        Raises
+        ------
+        RuntimeError
+            No refresh token is available or the token request failed.
+
+        """
+        del kwargs  # unused
+        if self.access_expires > time.time():
+            return
+
+        if not self.refresh_token or self.refresh_expires < time.time():
+            msg = "You are not logged in to MLflow. Please log in first."
+            raise RuntimeError(msg)
+
+        with Timer("Access token refreshed", self.log):
+            response = self._token_request()
+
+        self.access_token = response.get("access_token")
+        self.access_expires = time.time() + (response.get("expires_in") * 0.7)  # bit of buffer
+        self.refresh_token = response.get("refresh_token")
+
+        os.environ[self.target_env_var] = self.access_token
+
+    @enabled
+    def save(self, **kwargs: dict) -> None:
+        """Save the latest refresh token to disk."""
+        del kwargs  # unused
+        if not self.refresh_token:
+            self.log.warning("No refresh token to save.")
+            return
+
+        config = {
+            "url": self.url,
+            "refresh_token": self.refresh_token,
+            "refresh_expires": self.refresh_expires,
+        }
+        save_config(self.config_file, config)
+
+        expire_date = datetime.fromtimestamp(self.refresh_expires, tz=timezone.utc)
+        self.log.info("Your MLflow login token is valid until %s UTC", expire_date.strftime("%Y-%m-%d %H:%M:%S"))
+
+    def _token_request(
+        self,
+        ignore_exc: bool = False,
+    ) -> dict:
+        path = "refreshtoken"
+        payload = {"refresh_token": self.refresh_token}
+
+        try:
+            response = self._request(path, payload)
+        except Exception:
+            if ignore_exc:
+                return {}
+            raise
+
+        return response
+
+    def _request(self, path: str, payload: dict) -> dict:
+
+        headers = {
+            "Content-Type": "application/x-www-form-urlencoded",
+        }
+
+        try:
+            response = requests.post(f"{self.url}/{path}", headers=headers, json=payload, timeout=30)
+            response.raise_for_status()
+            response_json = response.json()
+
+            if response_json.get("status", "") != "OK":
+                error_description = response_json.get("response", "Error acquiring token.")
+                msg = f"❌ {error_description}"
+                raise RuntimeError(msg)
+
+            return response_json["response"]
+        except HTTPError:
+            self.log.exception("HTTP error occurred")
+            raise
diff --git a/src/anemoi/training/diagnostics/mlflow/client.py b/src/anemoi/training/diagnostics/mlflow/client.py
new file mode 100644
index 00000000..5c49f929
--- /dev/null
+++ b/src/anemoi/training/diagnostics/mlflow/client.py
@@ -0,0 +1,56 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 __future__ import annotations
+
+from typing import Any
+
+from mlflow import MlflowClient
+
+from anemoi.training.diagnostics.mlflow.auth import TokenAuth
+from anemoi.training.diagnostics.mlflow.utils import health_check
+
+
+class AnemoiMlflowClient(MlflowClient):
+    """Anemoi extension of the MLflow client with token authentication support."""
+
+    def __init__(
+        self,
+        tracking_uri: str,
+        *args,
+        authentication: bool = False,
+        check_health: bool = True,
+        **kwargs,
+    ) -> None:
+        """Behaves like a normal `mlflow.MlflowClient` but with token authentication injected on every call.
+
+        Parameters
+        ----------
+        tracking_uri : str
+            The URI of the MLflow tracking server.
+        authentication : bool, optional
+            Enable token authentication, by default False
+        check_health : bool, optional
+            Check the health of the MLflow server on init, by default True
+        *args : Any
+            Additional arguments to pass to the MLflow client.
+        **kwargs : Any
+            Additional keyword arguments to pass to the MLflow client.
+
+        """
+        self.anemoi_auth = TokenAuth(tracking_uri, enabled=authentication)
+        if check_health:
+            super().__getattribute__("anemoi_auth").authenticate()
+            health_check(tracking_uri)
+        super().__init__(tracking_uri, *args, **kwargs)
+
+    def __getattribute__(self, name: str) -> Any:
+        """Intercept attribute access and inject authentication."""
+        attr = super().__getattribute__(name)
+        if callable(attr) and name != "anemoi_auth":
+            super().__getattribute__("anemoi_auth").authenticate()
+        return attr
diff --git a/src/anemoi/training/diagnostics/mlflow/logger.py b/src/anemoi/training/diagnostics/mlflow/logger.py
new file mode 100644
index 00000000..2abebb1c
--- /dev/null
+++ b/src/anemoi/training/diagnostics/mlflow/logger.py
@@ -0,0 +1,467 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 __future__ import annotations
+
+import io
+import logging
+import os
+import re
+import sys
+import time
+from pathlib import Path
+from threading import Thread
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import Literal
+from weakref import WeakValueDictionary
+
+from packaging.version import Version
+from pytorch_lightning.loggers.mlflow import MLFlowLogger
+from pytorch_lightning.loggers.mlflow import _convert_params
+from pytorch_lightning.loggers.mlflow import _flatten_dict
+from pytorch_lightning.utilities.rank_zero import rank_zero_only
+
+from anemoi.training.diagnostics.mlflow.auth import TokenAuth
+from anemoi.training.diagnostics.mlflow.utils import health_check
+from anemoi.training.utils.jsonify import map_config_to_primitives
+
+if TYPE_CHECKING:
+    from argparse import Namespace
+
+    from omegaconf import OmegaConf
+
+LOGGER = logging.getLogger(__name__)
+
+
+def get_mlflow_run_params(config: OmegaConf, tracking_uri: str) -> tuple[str | None, str, dict[str, Any]]:
+    run_id = None
+    tags = {"projectName": config.diagnostics.log.mlflow.project_name}
+    # create a tag with the command used to run the script
+    command = os.environ.get("ANEMOI_TRAINING_CMD", sys.argv[0])
+    tags["command"] = command.split("/")[-1]  # get the python script name
+    tags["mlflow.source.name"] = command
+    if len(sys.argv) > 1:
+        # add the arguments to the command tag
+        tags["command"] = tags["command"] + " " + " ".join(sys.argv[1:])
+    if config.training.run_id or config.training.fork_run_id:
+        "Either run_id or fork_run_id must be provided to resume a run."
+
+        import mlflow
+
+        if config.diagnostics.log.mlflow.authentication and not config.diagnostics.log.mlflow.offline:
+            TokenAuth(tracking_uri).authenticate()
+
+        mlflow_client = mlflow.MlflowClient(tracking_uri)
+
+        if config.training.run_id and config.diagnostics.log.mlflow.on_resume_create_child:
+            parent_run_id = config.training.run_id  # parent_run_id
+            run_name = mlflow_client.get_run(parent_run_id).info.run_name
+            tags["mlflow.parentRunId"] = parent_run_id
+            tags["resumedRun"] = "True"  # tags can't take boolean values
+        elif config.training.run_id and not config.diagnostics.log.mlflow.on_resume_create_child:
+            run_id = config.training.run_id
+            run_name = mlflow_client.get_run(run_id).info.run_name
+            mlflow_client.update_run(run_id=run_id, status="RUNNING")
+            tags["resumedRun"] = "True"
+        else:
+            parent_run_id = config.training.fork_run_id
+            tags["forkedRun"] = "True"
+            tags["forkedRunId"] = parent_run_id
+
+    if config.diagnostics.log.mlflow.run_name:
+        run_name = config.diagnostics.log.mlflow.run_name
+    else:
+        import uuid
+
+        run_name = f"{uuid.uuid4()!s}"
+    return run_id, run_name, tags
+
+
+class LogsMonitor:
+    """Class for logging terminal output.
+
+    Inspired by the class for logging terminal output in aim.
+    Aim-Code: https://github.com/aimhubio/aim/blob/94646d2d317ec7a43303a16530f7963e4e652921/aim/ext/resource/tracker.py#L20
+
+    Note: If there is an error, the terminal output logging ends before the error message is printed into the log file.
+    In order for the user to see the error message, the user must look at the slurm output file.
+    We provide the SLRM job id in the very beginning of the log file and print the final status of the run in the end.
+
+    Parameters
+    ----------
+    artifact_save_dir : str
+        Directory to save the terminal logs.
+    experiment : MLflow experiment object.
+        MLflow experiment object.
+    run_id: str
+        MLflow run ID.
+    log_time_interval : int
+        Interval (in seconds) at which to write buffered terminal outputs, default 30
+
+    """
+
+    _buffer_registry = WeakValueDictionary()
+    _old_out_write = None
+    _old_err_write = None
+
+    def __init__(
+        self,
+        artifact_save_dir: str | Path,
+        experiment: MLFlowLogger.experiment,
+        run_id: str,
+        log_time_interval: float = 30.0,
+    ) -> None:
+        """Initialize the LogsMonitor.
+
+        Parameters
+        ----------
+        artifact_save_dir : str | Path
+            Directory for artifact saves.
+        experiment : MLFlowLogger.experiment
+            Experiment from MLFlow
+        run_id : str
+            Run ID
+        log_time_interval : float, optional
+            Logging time interval in seconds, by default 30.0
+
+        """
+        # active run
+        self.experiment = experiment
+        self.run_id = run_id
+
+        # terminal log capturing
+        self._log_capture_interval = 1
+        self._log_time_interval = log_time_interval
+        self._old_out = None
+        self._old_err = None
+        self._io_buffer = io.BytesIO()
+
+        # Start thread to collect stats and logs at intervals
+        self._th_collector = Thread(target=self._log_collector, daemon=True)
+        self._shutdown = False
+        self._started = False
+
+        # open your files here
+        self.artifact_save_dir = artifact_save_dir
+        self.file_save_path = Path(artifact_save_dir, "terminal_log.txt")
+        self.file_save_path.parent.mkdir(parents=True, exist_ok=True)
+
+    @classmethod
+    def _install_stream_patches(cls) -> None:
+        cls._old_out_write = sys.stdout.write
+        cls._old_err_write = sys.stderr.write
+
+        def new_out_write(data: str | bytes) -> None:
+            # out to buffer
+            cls._old_out_write(data)
+            if isinstance(data, str):
+                data = data.encode()
+            for buffer in cls._buffer_registry.values():
+                buffer.write(data)
+
+        def new_err_write(data: str | bytes) -> None:
+            # err to buffer
+            cls._old_err_write(data)
+            if isinstance(data, str):
+                data = data.encode()
+            for buffer in cls._buffer_registry.values():
+                buffer.write(data)
+
+        sys.stdout.write = new_out_write
+        sys.stderr.write = new_err_write
+
+    @classmethod
+    def _uninstall_stream_patches(cls) -> None:
+        sys.stdout.write = cls._old_out_write
+        sys.stderr.write = cls._old_err_write
+
+    def start(self) -> None:
+        """Start collection."""
+        if self._started:
+            return
+        self._started = True
+        # install the stream patches if not done yet
+        if not self._buffer_registry:
+            self._install_stream_patches()
+        self._buffer_registry[id(self)] = self._io_buffer
+        # Start thread to asynchronously collect logs
+        self._th_collector.start()
+        LOGGER.info("Termial Log Path: %s", self.file_save_path)
+        if os.getenv("SLURM_JOB_ID"):
+            LOGGER.info("SLURM job id: %s", os.getenv("SLURM_JOB_ID"))
+
+    def finish(self, status: str) -> None:
+        """Stop the monitoring and close the log file."""
+        if not self._started:
+            return
+        LOGGER.info(
+            ("Stopping terminal log monitoring and saving buffered terminal outputs. Final status: %s"),
+            status.upper(),
+        )
+        self._shutdown = True
+        # read and store remaining buffered logs
+        self._store_buffered_logs()
+        # unregister the buffer
+        del self._buffer_registry[id(self)]
+        # uninstall stream patching if no buffer is left in the registry
+        if not self._buffer_registry:
+            self._uninstall_stream_patches()
+
+        with self.file_save_path.open("a") as logfile:
+            logfile.write("\n\n")
+            logfile.flush()
+            logfile.close()
+
+    def _log_collector(self) -> None:
+        """Log collecting thread body.
+
+        Main monitoring loop, which consistently collect and log outputs.
+        """
+        log_capture_time_counter = 0
+
+        while True:
+            if self._shutdown:
+                break
+
+            time.sleep(self._log_time_interval)  # in seconds
+            log_capture_time_counter += self._log_time_interval
+
+            if log_capture_time_counter > self._log_capture_interval:
+                self._store_buffered_logs()
+                log_capture_time_counter = 0
+
+    def _store_buffered_logs(self) -> None:
+        _buffer_size = self._io_buffer.tell()
+        if not _buffer_size:
+            return
+        self._io_buffer.seek(0)
+        # read and reset the buffer
+        data = self._io_buffer.read(_buffer_size)
+        self._io_buffer.seek(0)
+        # handle the buffered data and store
+        # split lines and keep \n at the end of each line
+        lines = [e + b"\n" for e in data.split(b"\n") if e]
+
+        ansi_csi_re = re.compile(b"\001?\033\\[((?:\\d|;)*)([a-dA-D])\002?")
+
+        def _handle_csi(line: bytes) -> bytes:
+            # removes the cursor up and down symbols from the line
+            # skip tqdm status bar updates ending with "curser up" but the last one in buffer to save space
+            def _remove_csi(line: bytes) -> bytes:
+                # replacing the leftmost non-overlapping occurrences of
+                # pattern ansi_csi_re in string line by the replacement ""
+                return re.sub(ansi_csi_re, b"", line)
+
+            for match in ansi_csi_re.finditer(line):
+                arg, command = match.groups()
+                arg = int(arg.decode()) if arg else 1
+                if command == b"A" and (
+                    b"0%" not in line and not self._shutdown and b"[INFO]" not in line and b"[DEBUG]" not in line
+                ):  # cursor up
+                    # only keep x*10% status updates from tqmd status bars that end with a cursor up
+                    # always keep shutdown commands
+                    # always keep logger info and debug prints
+                    line = b""
+            return _remove_csi(line)
+
+        line = None
+        with self.file_save_path.open("a") as logfile:
+            for line in lines:
+                # handle cursor up and down symbols
+                cleaned_line = _handle_csi(line)
+                # handle each line for carriage returns
+                cleaned_line = cleaned_line.rsplit(b"\r")[-1]
+                logfile.write(cleaned_line.decode())
+
+            logfile.flush()
+        self.experiment.log_artifact(self.run_id, str(self.file_save_path))
+
+
+class AnemoiMLflowLogger(MLFlowLogger):
+    """A custom MLflow logger that logs terminal output."""
+
+    def __init__(
+        self,
+        experiment_name: str = "lightning_logs",
+        run_name: str | None = None,
+        tracking_uri: str | None = os.getenv("MLFLOW_TRACKING_URI"),
+        tags: dict[str, Any] | None = None,
+        save_dir: str | None = "./mlruns",
+        log_model: Literal[True, False, "all"] = False,
+        prefix: str = "",
+        resumed: bool | None = False,
+        forked: bool | None = False,
+        run_id: str | None = None,
+        offline: bool | None = False,
+        authentication: bool | None = None,
+        log_hyperparams: bool | None = True,
+    ) -> None:
+        """Initialize the AnemoiMLflowLogger.
+
+        Parameters
+        ----------
+        experiment_name : str, optional
+            Name of experiment, by default "lightning_logs"
+        run_name : str | None, optional
+            Name of run, by default None
+        tracking_uri : str | None, optional
+            Tracking URI of server, by default os.getenv("MLFLOW_TRACKING_URI")
+        tags : dict[str, Any] | None, optional
+            Tags to apply, by default None
+        save_dir : str | None, optional
+            Directory to save logs to, by default "./mlruns"
+        log_model : Literal[True, False, "all"], optional
+            Log model checkpoints to server (expensive), by default False
+        prefix : str, optional
+            Prefix for experiments, by default ""
+        resumed : bool | None, optional
+            Whether the run was resumed or not, by default False
+        forked : bool | None, optional
+            Whether the run was forked or not, by default False
+        run_id : str | None, optional
+            Run id of current run, by default None
+        offline : bool | None, optional
+            Whether to run offline or not, by default False
+        authentication : bool | None, optional
+            Whether to authenticate with server or not, by default None
+        log_hyperparams : bool | None, optional
+            Whether to log hyperparameters, by default True
+
+        """
+        if offline:
+            # OFFLINE - When we run offline we can pass a save_dir pointing to a local path
+            tracking_uri = None
+
+        else:
+            # ONLINE - When we pass a tracking_uri to mlflow then it will ignore the
+            # saving dir and save all artifacts/metrics to the remote server database
+            save_dir = None
+
+        self._resumed = resumed
+        self._forked = forked
+        self._flag_log_hparams = log_hyperparams
+
+        if rank_zero_only.rank == 0:
+            enabled = authentication and not offline
+            self.auth = TokenAuth(tracking_uri, enabled=enabled)
+
+            if offline:
+                LOGGER.info("MLflow is logging offline.")
+            else:
+                LOGGER.info("MLflow token authentication %s for %s", "enabled" if enabled else "disabled", tracking_uri)
+                self.auth.authenticate()
+                health_check(tracking_uri)
+
+        super().__init__(
+            experiment_name=experiment_name,
+            run_name=run_name,
+            tracking_uri=tracking_uri,
+            tags=tags,
+            save_dir=save_dir,
+            log_model=log_model,
+            prefix=prefix,
+            run_id=run_id,
+        )
+
+    @property
+    def experiment(self) -> MLFlowLogger.experiment:
+        if rank_zero_only.rank == 0:
+            self.auth.authenticate()
+        return super().experiment
+
+    @rank_zero_only
+    def log_system_metrics(self) -> None:
+        """Log system metrics (CPU, GPU, etc)."""
+        import mlflow
+        from mlflow.system_metrics.system_metrics_monitor import SystemMetricsMonitor
+
+        mlflow.enable_system_metrics_logging()
+        system_monitor = SystemMetricsMonitor(
+            self.run_id,
+            resume_logging=self.run_id is not None,
+        )
+        self.run_id_to_system_metrics_monitor = {}
+        self.run_id_to_system_metrics_monitor[self.run_id] = system_monitor
+        system_monitor.start()
+
+    @rank_zero_only
+    def log_terminal_output(self, artifact_save_dir: str | Path = "") -> None:
+        """Log terminal logs to MLflow."""
+        # path for logging terminal logs
+        # for now the 'terminal_logs' file is kept in the same folder as the plots
+        artifact_save_dir = Path(artifact_save_dir, self.run_id, "plots")
+
+        log_monitor = LogsMonitor(
+            artifact_save_dir,
+            self.experiment,
+            self.run_id,
+        )
+        self.run_id_to_log_monitor = {}
+        self.run_id_to_log_monitor[self.run_id] = log_monitor
+        log_monitor.start()
+
+    @staticmethod
+    def _clean_params(params: dict[str, Any]) -> dict[str, Any]:
+        """Clean up params to avoid issues with mlflow.
+
+        Too many logged params will make the server take longer to render the
+        experiment.
+
+        Parameters
+        ----------
+        params : dict[str, Any]
+            Parameters to clean up.
+
+        Returns
+        -------
+        dict[str, Any]
+            Cleaned up params ready for MlFlow.
+        """
+        prefixes_to_remove = ["hardware", "data", "dataloader", "model", "training", "diagnostics", "metadata.config"]
+        keys_to_remove = [key for key in params if any(key.startswith(prefix) for prefix in prefixes_to_remove)]
+        for key in keys_to_remove:
+            del params[key]
+        return params
+
+    @rank_zero_only
+    def log_hyperparams(self, params: dict[str, Any] | Namespace) -> None:
+        """Overwrite the log_hyperparams method to flatten config params using '.'."""
+        if self._flag_log_hparams:
+            params = _convert_params(params)
+
+            # this is needed to resolve optional missing config values to a string, instead of raising a missing error
+            if config := params.get("config"):
+                params["config"] = map_config_to_primitives(config)
+
+            params = _flatten_dict(params, delimiter=".")  # Flatten dict with '.' to not break API queries
+            params = self._clean_params(params)
+
+            import mlflow
+            from mlflow.entities import Param
+
+            # Truncate parameter values.
+            truncation_length = 250
+            if Version(mlflow.VERSION) >= Version("1.28.0"):
+                truncation_length = 500
+            params_list = [Param(key=k, value=str(v)[:truncation_length]) for k, v in params.items()]
+
+            for idx in range(0, len(params_list), 100):
+                self.experiment.log_batch(run_id=self.run_id, params=params_list[idx : idx + 100])
+
+    @rank_zero_only
+    def finalize(self, status: str = "success") -> None:
+        # save the last obtained refresh token to disk
+        self.auth.save()
+
+        # finalize logging and system metrics monitor
+        if getattr(self, "run_id_to_system_metrics_monitor", None):
+            self.run_id_to_system_metrics_monitor[self.run_id].finish()
+        if getattr(self, "run_id_to_log_monitor", None):
+            self.run_id_to_log_monitor[self.run_id].finish(status)
+
+        super().finalize(status)
diff --git a/src/anemoi/training/diagnostics/mlflow/utils.py b/src/anemoi/training/diagnostics/mlflow/utils.py
new file mode 100644
index 00000000..9e0ae4e7
--- /dev/null
+++ b/src/anemoi/training/diagnostics/mlflow/utils.py
@@ -0,0 +1,35 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+import os
+
+import requests
+
+
+def health_check(tracking_uri: str) -> None:
+    """Query the health endpoint of an MLflow server.
+
+    If the server is not reachable, raise an error and remind the user that authentication may be required.
+
+    Raises
+    ------
+    ConnectionError
+        If the server is not reachable.
+
+    """
+    token = os.getenv("MLFLOW_TRACKING_TOKEN")
+
+    headers = {"Authorization": f"Bearer {token}"}
+    response = requests.get(f"{tracking_uri}/health", headers=headers, timeout=60)
+
+    if response.text == "OK":
+        return
+
+    error_msg = f"Could not connect to MLflow server at {tracking_uri}. "
+    if not token:
+        error_msg += "The server may require authentication, did you forget to turn it on?"
+    raise ConnectionError(error_msg)
diff --git a/src/anemoi/training/diagnostics/plots.py b/src/anemoi/training/diagnostics/plots.py
new file mode 100644
index 00000000..b2004cf4
--- /dev/null
+++ b/src/anemoi/training/diagnostics/plots.py
@@ -0,0 +1,660 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import matplotlib.pyplot as plt
+import matplotlib.style as mplstyle
+import numpy as np
+from matplotlib.colors import BoundaryNorm
+from matplotlib.colors import ListedColormap
+from matplotlib.colors import TwoSlopeNorm
+from pyshtools.expand import SHGLQ
+from pyshtools.expand import SHExpandGLQ
+from scipy.interpolate import griddata
+
+from anemoi.training.diagnostics.maps import Coastlines
+from anemoi.training.diagnostics.maps import EquirectangularProjection
+
+if TYPE_CHECKING:
+    from matplotlib.figure import Figure
+
+from dataclasses import dataclass
+
+LOGGER = logging.getLogger(__name__)
+
+continents = Coastlines()
+
+
+@dataclass
+class LatLonData:
+    latitudes: np.ndarray
+    longitudes: np.ndarray
+    data: np.ndarray
+
+
+def init_plot_settings() -> None:
+    """Initialize matplotlib plot settings."""
+    small_font_size = 8
+    medium_font_size = 10
+
+    mplstyle.use("fast")
+    plt.rcParams["path.simplify_threshold"] = 0.9
+
+    plt.rc("font", size=small_font_size)  # controls default text sizes
+    plt.rc("axes", titlesize=small_font_size)  # fontsize of the axes title
+    plt.rc("axes", labelsize=medium_font_size)  # fontsize of the x and y labels
+    plt.rc("xtick", labelsize=small_font_size)  # fontsize of the tick labels
+    plt.rc("ytick", labelsize=small_font_size)  # fontsize of the tick labels
+    plt.rc("legend", fontsize=small_font_size)  # legend fontsize
+    plt.rc("figure", titlesize=small_font_size)  # fontsize of the figure title
+
+
+def _hide_axes_ticks(ax: plt.Axes) -> None:
+    """Hide x/y-axis ticks.
+
+    Parameters
+    ----------
+    ax : matplotlib.axes
+        Axes object handle
+
+    """
+    plt.setp(ax.get_xticklabels(), visible=False)
+    plt.setp(ax.get_yticklabels(), visible=False)
+    ax.tick_params(axis="both", which="both", length=0)
+
+
+def plot_loss(
+    x: np.ndarray,
+    colors: np.ndarray,
+    xticks: dict[str, int] | None = None,
+    legend_patches: list | None = None,
+) -> Figure:
+    """Plots data for one multilevel sample.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        Data for Plotting of shape (npred,)
+    colors : np.ndarray
+        Colors for the bars.
+    xticks : dict, optional
+        Dictionary of xticks, by default None
+    legend_patches : list, optional
+        List of legend patches, by default None
+
+    Returns
+    -------
+    Figure
+        The figure object handle.
+
+    """
+    # create plot
+    # more space for legend
+    figsize = (8, 3) if legend_patches else (4, 3)
+    fig, ax = plt.subplots(1, 1, figsize=figsize)
+    # histogram plot
+    ax.bar(np.arange(x.size), x, color=colors, log=1)
+
+    # add xticks and legend if given
+    if xticks:
+        ax.set_xticks(list(xticks.values()), list(xticks.keys()), rotation=60)
+    if legend_patches:
+        # legend outside and to the right of the plot
+        plt.legend(handles=legend_patches, bbox_to_anchor=(1.01, 1), loc="upper left")
+    plt.tight_layout()
+
+    return fig
+
+
+def plot_power_spectrum(
+    parameters: dict[str, int],
+    latlons: np.ndarray,
+    x: np.ndarray,
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+) -> Figure:
+    """Plots power spectrum.
+
+    NB: this can be very slow for large data arrays
+    call it as infrequently as possible!
+
+    Parameters
+    ----------
+    parameters : dict[str, int]
+        Dictionary of variable names and indices
+    latlons : np.ndarray
+        lat/lon coordinates array, shape (lat*lon, 2)
+    x : np.ndarray
+        Input data of shape (lat*lon, nvar*level)
+    y_true : np.ndarray
+        Expected data of shape (lat*lon, nvar*level)
+    y_pred : np.ndarray
+        Predicted data of shape (lat*lon, nvar*level)
+
+    Returns
+    -------
+    Figure
+        The figure object handle.
+
+    """
+    n_plots_x, n_plots_y = len(parameters), 1
+
+    figsize = (n_plots_y * 4, n_plots_x * 3)
+    fig, ax = plt.subplots(n_plots_x, n_plots_y, figsize=figsize)
+
+    pc = EquirectangularProjection()
+    lat, lon = latlons[:, 0], latlons[:, 1]
+    pc_lon, pc_lat = pc(lon, lat)
+    pc_lon = np.array(pc_lon)
+    pc_lat = np.array(pc_lat)
+    # Calculate delta_lon and delta_lat on the projected grid
+    delta_lon = abs(np.diff(pc_lon))
+    non_zero_delta_lon = delta_lon[delta_lon != 0]
+    delta_lat = abs(np.diff(pc_lat))
+    non_zero_delta_lat = delta_lat[delta_lat != 0]
+
+    # Define a regular grid for interpolation
+    n_pix_lon = int(np.floor(abs(pc_lon.max() - pc_lon.min()) / abs(np.min(non_zero_delta_lon))))  # around 400 for O96
+    n_pix_lat = int(np.floor(abs(pc_lat.max() - pc_lat.min()) / abs(np.min(non_zero_delta_lat))))  # around 192 for O96
+    regular_pc_lon = np.linspace(pc_lon.min(), pc_lon.max(), n_pix_lon)
+    regular_pc_lat = np.linspace(pc_lat.min(), pc_lat.max(), n_pix_lat)
+    grid_pc_lon, grid_pc_lat = np.meshgrid(regular_pc_lon, regular_pc_lat)
+
+    for plot_idx, (variable_idx, (variable_name, output_only)) in enumerate(parameters.items()):
+        yt = y_true[..., variable_idx].squeeze()
+        yp = y_pred[..., variable_idx].squeeze()
+
+        # check for any nan in yt
+        nan_flag = np.isnan(yt).any()
+
+        method = "linear" if nan_flag else "cubic"
+        if output_only:
+            xt = x[..., variable_idx].squeeze()
+            yt_i = griddata((pc_lon, pc_lat), (yt - xt), (grid_pc_lon, grid_pc_lat), method=method, fill_value=0.0)
+            yp_i = griddata((pc_lon, pc_lat), (yp - xt), (grid_pc_lon, grid_pc_lat), method=method, fill_value=0.0)
+        else:
+            yt_i = griddata((pc_lon, pc_lat), yt, (grid_pc_lon, grid_pc_lat), method=method, fill_value=0.0)
+            yp_i = griddata((pc_lon, pc_lat), yp, (grid_pc_lon, grid_pc_lat), method=method, fill_value=0.0)
+
+        # Masking NaN values
+        if nan_flag:
+            mask = np.isnan(yt_i)
+            if mask.any():
+                yt_i = np.where(mask, 0.0, yt_i)
+                yp_i = np.where(mask, 0.0, yp_i)
+
+        amplitude_t = np.array(compute_spectra(yt_i))
+        amplitude_p = np.array(compute_spectra(yp_i))
+
+        ax[plot_idx].loglog(
+            np.arange(1, amplitude_t.shape[0]),
+            amplitude_t[1 : (amplitude_t.shape[0])],
+            label="Truth (data)",
+        )
+        ax[plot_idx].loglog(
+            np.arange(1, amplitude_p.shape[0]),
+            amplitude_p[1 : (amplitude_p.shape[0])],
+            label="Predicted",
+        )
+
+        ax[plot_idx].legend()
+        ax[plot_idx].set_title(variable_name)
+
+        ax[plot_idx].set_xlabel("$k$")
+        ax[plot_idx].set_ylabel("$P(k)$")
+        ax[plot_idx].set_aspect("auto", adjustable=None)
+    fig.tight_layout()
+    return fig
+
+
+def compute_spectra(field: np.ndarray) -> np.ndarray:
+    """Compute spectral variability of a field by wavenumber.
+
+    Parameters
+    ----------
+    field : np.ndarray
+        lat lon field to calculate the spectra of
+
+    Returns
+    -------
+    np.ndarray
+        spectra of field by wavenumber
+
+    """
+    field = np.array(field)
+
+    # compute real and imaginary parts of power spectra of field
+    lmax = field.shape[0] - 1  # maximum degree of expansion
+    zero_w = SHGLQ(lmax)
+    coeffs_field = SHExpandGLQ(field, w=zero_w[1], zero=zero_w[0])
+
+    # Re**2 + Im**2
+    coeff_amp = coeffs_field[0, :, :] ** 2 + coeffs_field[1, :, :] ** 2
+
+    # sum over meridional direction
+    return np.sum(coeff_amp, axis=0)
+
+
+def plot_histogram(
+    parameters: dict[str, int],
+    x: np.ndarray,
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    precip_and_related_fields: list | None = None,
+) -> Figure:
+    """Plots histogram.
+
+    NB: this can be very slow for large data arrays
+    call it as infrequently as possible!
+
+    Parameters
+    ----------
+    parameters : dict[str, int]
+        Dictionary of variable names and indices
+    x : np.ndarray
+        Input data of shape (lat*lon, nvar*level)
+    y_true : np.ndarray
+        Expected data of shape (lat*lon, nvar*level)
+    y_pred : np.ndarray
+        Predicted data of shape (lat*lon, nvar*level)
+    precip_and_related_fields : list, optional
+        List of precipitation-like variables, by default []
+
+    Returns
+    -------
+    Figure
+        The figure object handle.
+
+    """
+    precip_and_related_fields = precip_and_related_fields or []
+
+    n_plots_x, n_plots_y = len(parameters), 1
+
+    figsize = (n_plots_y * 4, n_plots_x * 3)
+    fig, ax = plt.subplots(n_plots_x, n_plots_y, figsize=figsize)
+
+    for plot_idx, (variable_idx, (variable_name, output_only)) in enumerate(parameters.items()):
+        yt = y_true[..., variable_idx].squeeze()
+        yp = y_pred[..., variable_idx].squeeze()
+        # postprocessed outputs so we need to handle possible NaNs
+
+        # Calculate the histogram and handle NaNs
+        if output_only:
+            # histogram of true increment and predicted increment
+            xt = x[..., variable_idx].squeeze() * int(output_only)
+            yt_xt = yt - xt
+            yp_xt = yp - xt
+            # enforce the same binning for both histograms
+            bin_min = min(np.nanmin(yt_xt), np.nanmin(yp_xt))
+            bin_max = max(np.nanmax(yt_xt), np.nanmax(yp_xt))
+            hist_yt, bins_yt = np.histogram(yt_xt[~np.isnan(yt_xt)], bins=100, range=[bin_min, bin_max])
+            hist_yp, bins_yp = np.histogram(yp_xt[~np.isnan(yp_xt)], bins=100, range=[bin_min, bin_max])
+        else:
+            # enforce the same binning for both histograms
+            bin_min = min(np.nanmin(yt), np.nanmin(yp))
+            bin_max = max(np.nanmax(yt), np.nanmax(yp))
+            hist_yt, bins_yt = np.histogram(yt[~np.isnan(yt)], bins=100, range=[bin_min, bin_max])
+            hist_yp, bins_yp = np.histogram(yp[~np.isnan(yp)], bins=100, range=[bin_min, bin_max])
+
+        # Visualization trick for tp
+        if variable_name in precip_and_related_fields:
+            # in-place multiplication does not work here because variables are different numpy types
+            hist_yt = hist_yt * bins_yt[:-1]
+            hist_yp = hist_yp * bins_yp[:-1]
+        # Plot the modified histogram
+        ax[plot_idx].bar(bins_yt[:-1], hist_yt, width=np.diff(bins_yt), color="blue", alpha=0.7, label="Truth (data)")
+        ax[plot_idx].bar(bins_yp[:-1], hist_yp, width=np.diff(bins_yp), color="red", alpha=0.7, label="Predicted")
+
+        ax[plot_idx].set_title(variable_name)
+        ax[plot_idx].set_xlabel(variable_name)
+        ax[plot_idx].set_ylabel("Density")
+        ax[plot_idx].legend()
+        ax[plot_idx].set_aspect("auto", adjustable=None)
+
+    fig.tight_layout()
+    return fig
+
+
+def plot_predicted_multilevel_flat_sample(
+    parameters: dict[str, int],
+    n_plots_per_sample: int,
+    latlons: np.ndarray,
+    clevels: float,
+    cmap_precip: str,
+    x: np.ndarray,
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    precip_and_related_fields: list | None = None,
+) -> Figure:
+    """Plots data for one multilevel latlon-"flat" sample.
+
+    NB: this can be very slow for large data arrays
+    call it as infrequently as possible!
+
+    Parameters
+    ----------
+    parameters : dict[str, int]
+        Dictionary of variable names and indices
+    n_plots_per_sample : int
+        Number of plots per sample
+    latlons : np.ndarray
+        lat/lon coordinates array, shape (lat*lon, 2)
+    clevels : float
+        Accumulation levels used for precipitation related plots
+    cmap_precip: str
+        Colors used for each accumulation level
+    x : np.ndarray
+        Input data of shape (lat*lon, nvar*level)
+    y_true : np.ndarray
+        Expected data of shape (lat*lon, nvar*level)
+    y_pred : np.ndarray
+        Predicted data of shape (lat*lon, nvar*level)
+    precip_and_related_fields : list, optional
+        List of precipitation-like variables, by default []
+
+    Returns
+    -------
+    Figure
+        The figure object handle.
+
+    """
+    n_plots_x, n_plots_y = len(parameters), n_plots_per_sample
+
+    figsize = (n_plots_y * 4, n_plots_x * 3)
+    fig, ax = plt.subplots(n_plots_x, n_plots_y, figsize=figsize)
+
+    pc = EquirectangularProjection()
+    lat, lon = latlons[:, 0], latlons[:, 1]
+    pc_lon, pc_lat = pc(lon, lat)
+
+    for plot_idx, (variable_idx, (variable_name, output_only)) in enumerate(parameters.items()):
+        xt = x[..., variable_idx].squeeze() * int(output_only)
+        yt = y_true[..., variable_idx].squeeze()
+        yp = y_pred[..., variable_idx].squeeze()
+        if n_plots_x > 1:
+            plot_flat_sample(
+                fig,
+                ax[plot_idx, :],
+                pc_lon,
+                pc_lat,
+                xt,
+                yt,
+                yp,
+                variable_name,
+                clevels,
+                cmap_precip,
+                precip_and_related_fields,
+            )
+        else:
+            plot_flat_sample(
+                fig,
+                ax,
+                pc_lon,
+                pc_lat,
+                xt,
+                yt,
+                yp,
+                variable_name,
+                clevels,
+                cmap_precip,
+                precip_and_related_fields,
+            )
+
+    return fig
+
+
+def plot_flat_sample(
+    fig: Figure,
+    ax: plt.Axes,
+    lon: np.ndarray,
+    lat: np.ndarray,
+    input_: np.ndarray,
+    truth: np.ndarray,
+    pred: np.ndarray,
+    vname: str,
+    clevels: float,
+    cmap_precip: str,
+    precip_and_related_fields: list | None = None,
+) -> None:
+    """Plot a "flat" 1D sample.
+
+    Data on non-rectangular (reduced Gaussian) grids.
+
+    Parameters
+    ----------
+    fig : _type_
+        Figure object handle
+    ax : matplotlib.axes
+        Axis object handle
+    lon : np.ndarray
+        longitude coordinates array, shape (lon,)
+    lat : np.ndarray
+        latitude coordinates array, shape (lat,)
+    input_ : np.ndarray
+        Input data of shape (lat*lon,)
+    truth : np.ndarray
+        Expected data of shape (lat*lon,)
+    pred : np.ndarray
+        Predicted data of shape (lat*lon,)
+    vname : str
+        Variable name
+    clevels : float
+        Accumulation levels used for precipitation related plots
+    cmap_precip: str
+        Colors used for each accumulation level
+    precip_and_related_fields : list, optional
+        List of precipitation-like variables, by default []
+
+    """
+    precip_and_related_fields = precip_and_related_fields or []
+    if vname in precip_and_related_fields:
+        # Create a custom colormap for precipitation
+        nws_precip_colors = cmap_precip
+        precip_colormap = ListedColormap(nws_precip_colors)
+
+        # Defining the actual precipitation accumulation levels in mm
+        cummulation_lvls = clevels
+        norm = BoundaryNorm(cummulation_lvls, len(cummulation_lvls) + 1)
+
+        # converting to mm from m
+        truth *= 1000.0
+        pred *= 1000.0
+        scatter_plot(fig, ax[1], lon=lon, lat=lat, data=truth, cmap=precip_colormap, norm=norm, title=f"{vname} target")
+        scatter_plot(fig, ax[2], lon=lon, lat=lat, data=pred, cmap=precip_colormap, norm=norm, title=f"{vname} pred")
+        scatter_plot(
+            fig,
+            ax[3],
+            lon=lon,
+            lat=lat,
+            data=truth - pred,
+            cmap="bwr",
+            norm=TwoSlopeNorm(vcenter=0.0),
+            title=f"{vname} pred err",
+        )
+    elif vname == "mwd":
+        cyclic_colormap = "twilight"
+
+        def error_plot_in_degrees(array1: np.ndarray, array2: np.ndarray) -> np.ndarray:
+            """Calculate error between two arrays in degrees in range [-180, 180]."""
+            tmp = (array1 - array2) % 360
+            return np.where(tmp > 180, tmp - 360, tmp)
+
+        sample_shape = truth.shape
+        pred = np.maximum(np.zeros(sample_shape), np.minimum(360 * np.ones(sample_shape), (pred)))
+        scatter_plot(fig, ax[1], lon=lon, lat=lat, data=truth, cmap=cyclic_colormap, title=f"{vname} target")
+        scatter_plot(fig, ax[2], lon=lon, lat=lat, data=pred, cmap=cyclic_colormap, title=f"capped {vname} pred")
+        err_plot = error_plot_in_degrees(truth, pred)
+        scatter_plot(
+            fig,
+            ax[3],
+            lon=lon,
+            lat=lat,
+            data=err_plot,
+            cmap="bwr",
+            norm=TwoSlopeNorm(vcenter=0.0),
+            title=f"{vname} pred err: {np.nanmean(np.abs(err_plot)):.{4}f} deg.",
+        )
+    else:
+        scatter_plot(fig, ax[1], lon=lon, lat=lat, data=truth, title=f"{vname} target")
+        scatter_plot(fig, ax[2], lon=lon, lat=lat, data=pred, title=f"{vname} pred")
+        scatter_plot(
+            fig,
+            ax[3],
+            lon=lon,
+            lat=lat,
+            data=truth - pred,
+            cmap="bwr",
+            norm=TwoSlopeNorm(vcenter=0.0),
+            title=f"{vname} pred err",
+        )
+
+    if sum(input_) != 0:
+        if vname == "mwd":
+            scatter_plot(fig, ax[0], lon=lon, lat=lat, data=input_, cmap=cyclic_colormap, title=f"{vname} input")
+            err_plot = error_plot_in_degrees(pred, input_)
+            scatter_plot(
+                fig,
+                ax[4],
+                lon=lon,
+                lat=lat,
+                data=err_plot,
+                cmap="bwr",
+                norm=TwoSlopeNorm(vcenter=0.0),
+                title=f"{vname} increment [pred - input] % 360",
+            )
+            err_plot = error_plot_in_degrees(truth, input_)
+            scatter_plot(
+                fig,
+                ax[5],
+                lon=lon,
+                lat=lat,
+                data=err_plot,
+                cmap="bwr",
+                norm=TwoSlopeNorm(vcenter=0.0),
+                title=f"{vname} persist err: {np.nanmean(np.abs(err_plot)):.{4}f} deg.",
+            )
+        else:
+            scatter_plot(fig, ax[0], lon=lon, lat=lat, data=input_, title=f"{vname} input")
+            scatter_plot(
+                fig,
+                ax[4],
+                lon=lon,
+                lat=lat,
+                data=pred - input_,
+                cmap="bwr",
+                norm=TwoSlopeNorm(vcenter=0.0),
+                title=f"{vname} increment [pred - input]",
+            )
+            scatter_plot(
+                fig,
+                ax[5],
+                lon=lon,
+                lat=lat,
+                data=truth - input_,
+                cmap="bwr",
+                norm=TwoSlopeNorm(vcenter=0.0),
+                title=f"{vname} persist err",
+            )
+    else:
+        ax[0].axis("off")
+        ax[4].axis("off")
+        ax[5].axis("off")
+
+
+def scatter_plot(
+    fig: Figure,
+    ax: plt.Axes,
+    *,
+    lon: np.array,
+    lat: np.array,
+    data: np.array,
+    cmap: str = "viridis",
+    norm: str | None = None,
+    title: str | None = None,
+) -> None:
+    """Lat-lon scatter plot: can work with arbitrary grids.
+
+    Parameters
+    ----------
+    fig : _type_
+        Figure object handle
+    ax : matplotlib.axes
+        Axis object handle
+    lon : np.ndarray
+        longitude coordinates array, shape (lon,)
+    lat : np.ndarray
+        latitude coordinates array, shape (lat,)
+    data : _type_
+        Data to plot
+    cmap : str, optional
+        Colormap string from matplotlib, by default "viridis"
+    norm : str, optional
+        Normalization string from matplotlib, by default None
+    title : str, optional
+        Title for plot, by default None
+
+    """
+    psc = ax.scatter(
+        lon,
+        lat,
+        c=data,
+        cmap=cmap,
+        s=1,
+        alpha=1.0,
+        norm=norm,
+        rasterized=True,
+    )
+    ax.set_xlim((-np.pi, np.pi))
+    ax.set_ylim((-np.pi / 2, np.pi / 2))
+
+    continents.plot_continents(ax)
+
+    if title is not None:
+        ax.set_title(title)
+
+    ax.set_aspect("auto", adjustable=None)
+    _hide_axes_ticks(ax)
+    fig.colorbar(psc, ax=ax)
+
+
+def plot_graph_features(
+    latlons: np.ndarray,
+    features: np.ndarray,
+) -> Figure:
+    """Plot trainable graph features.
+
+    Parameters
+    ----------
+    latlons : np.ndarray
+        Latitudes and longitudes
+    features : np.ndarray
+        Trainable Features
+
+    Returns
+    -------
+    Figure
+        Figure object handle
+
+    """
+    nplots = features.shape[-1]
+    figsize = (nplots * 4, 3)
+    fig, ax = plt.subplots(1, nplots, figsize=figsize)
+
+    lat, lon = latlons[:, 0], latlons[:, 1]
+
+    pc = EquirectangularProjection()
+    pc_lon, pc_lat = pc(lon, lat)
+
+    for i in range(nplots):
+        ax_ = ax[i] if nplots > 1 else ax
+        scatter_plot(fig, ax_, lon=pc_lon, lat=pc_lat, data=features[..., i])
+
+    return fig
diff --git a/src/anemoi/training/distributed/__init__.py b/src/anemoi/training/distributed/__init__.py
new file mode 100644
index 00000000..282d6a69
--- /dev/null
+++ b/src/anemoi/training/distributed/__init__.py
@@ -0,0 +1,6 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
diff --git a/src/anemoi/training/distributed/strategy.py b/src/anemoi/training/distributed/strategy.py
new file mode 100644
index 00000000..c15828ca
--- /dev/null
+++ b/src/anemoi/training/distributed/strategy.py
@@ -0,0 +1,162 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+import logging
+import os
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from lightning_fabric.utilities.optimizer import _optimizers_to_device
+from pytorch_lightning.overrides.distributed import _sync_module_states
+from pytorch_lightning.strategies.ddp import DDPStrategy
+from pytorch_lightning.trainer.states import TrainerFn
+
+from anemoi.training.utils.seeding import get_base_seed
+
+LOGGER = logging.getLogger(__name__)
+
+
+class DDPGroupStrategy(DDPStrategy):
+    """Distributed Data Parallel strategy with group communication."""
+
+    def __init__(self, num_gpus_per_model: int, **kwargs: dict) -> None:
+        """Initialize the distributed strategy.
+
+        Parameters
+        ----------
+        num_gpus_per_model : int
+            Number of GPUs per model to shard over.
+        **kwargs : dict
+            Additional keyword arguments.
+
+        """
+        super().__init__(**kwargs)
+        self.model_comm_group_size = num_gpus_per_model
+
+    def setup(self, trainer: pl.Trainer) -> None:
+        assert self.accelerator is not None, "Accelerator is not initialized for distributed strategy"
+        self.accelerator.setup(trainer)
+
+        # determine the model groups that work together:
+
+        assert self.world_size % self.model_comm_group_size == 0, (
+            f"Total number of GPUs ({self.world_size}) must be divisible by the number of GPUs "
+            f"per model ({self.model_comm_group_size})."
+        )
+
+        model_comm_group_ranks = np.split(
+            np.arange(self.world_size, dtype=int),
+            int(self.world_size / self.model_comm_group_size),
+        )
+        model_comm_groups = [
+            torch.distributed.new_group(x) for x in model_comm_group_ranks
+        ]  # every rank has to create all of these
+
+        model_comm_group_id, model_comm_group_nr, model_comm_group_rank = self.get_my_model_comm_group(
+            self.model_comm_group_size,
+        )
+        model_comm_group = model_comm_groups[model_comm_group_id]
+        self.model.set_model_comm_group(model_comm_group)
+        LOGGER.debug(
+            "Rank %d model_comm_group is %s, group number %d, with local group rank %d and comms_group_ranks %s",
+            self.global_rank,
+            str(model_comm_group_nr),
+            model_comm_group_id,
+            model_comm_group_rank,
+            str(model_comm_group_ranks[model_comm_group_id]),
+        )
+
+        # register hooks for correct gradient reduction
+        self.register_parameter_hooks()
+
+        # move the model to the correct device
+        self.model_to_device()
+
+        # skip wrapping the model if we are not fitting as no gradients need to be exchanged
+        trainer_fn = trainer.state.fn
+
+        if trainer_fn == TrainerFn.FITTING and self._layer_sync:
+            assert self.model is not None, "Model is not initialized for distributed strategy"
+            self.model = self._layer_sync.apply(self.model)
+
+        self.setup_precision_plugin()
+
+        if trainer_fn == TrainerFn.FITTING:
+            # do not wrap with DDP if not fitting as there's no gradients to reduce
+            self.configure_ddp()
+
+            # set up optimizers after the wrapped module has been moved to the device
+            self.setup_optimizers(trainer)
+            _optimizers_to_device(self.optimizers, self.root_device)
+
+            import torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook as post_localSGD
+
+            if isinstance(self._ddp_comm_state, post_localSGD.PostLocalSGDState):
+                self._enable_model_averaging()
+        else:
+            # we need to manually synchronize the module's states since we aren't using the DDP wrapper
+            assert self.model is not None, "Model is not initialized for distributed strategy"
+            _sync_module_states(self.model)
+
+        # seed ranks
+        self.seed_rnd(model_comm_group_id)
+
+    def get_my_model_comm_group(self, num_gpus_per_model: int) -> tuple[int, np.ndarray, int]:
+        """Determine tasks that work together and from a model group.
+
+        Parameters
+        ----------
+        num_gpus_per_model : int
+            Number of GPUs per model to shard over.
+
+        Returns
+        -------
+        tuple[int, np.ndarray, int]
+            Model_comm_group id, Model_comm_group Nr, Model_comm_group rank
+        """
+        model_comm_groups = np.arange(0, self.world_size, dtype=np.int32)
+        model_comm_groups = np.split(model_comm_groups, self.world_size / num_gpus_per_model)
+
+        model_comm_group_id = None
+        for i, model_comm_group in enumerate(model_comm_groups):
+            if self.global_rank in model_comm_group:
+                model_comm_group_id = i
+                model_comm_group_nr = model_comm_group
+                model_comm_group_rank = np.ravel(np.asarray(model_comm_group == self.global_rank).nonzero())[0]
+        return model_comm_group_id, model_comm_group_nr, model_comm_group_rank
+
+    def seed_rnd(self, model_comm_group_id: int) -> None:
+        """Seed the random number generators for the rank."""
+        base_seed = get_base_seed()
+        initial_seed = base_seed * (model_comm_group_id + 1)
+        rnd_seed = pl.seed_everything(initial_seed)  # note: workers are seeded independently in dataloader
+        np_rng = np.random.default_rng(rnd_seed)
+        sanity_rnd = (torch.rand(1), np_rng.random())
+        LOGGER.debug(
+            (
+                "Strategy: Rank %d, model comm group id %d, base seed %d, seeded with %d, "
+                "running with random seed: %d, sanity rnd: %s"
+            ),
+            int(os.environ.get("SLURM_PROCID", "0")),
+            model_comm_group_id,
+            base_seed,
+            initial_seed,
+            rnd_seed,
+            sanity_rnd,
+        )
+
+    def register_parameter_hooks(self) -> None:
+        """Register parameter hooks for gradient reduction.
+
+        Here, we rescale parameters that only see a subset of the input on each rank
+        -> these are still divided by the total number of GPUs in DDP as if each rank would see a full set of inputs
+        note: the trainable parameters are added before the split across GPUs and are therefore not rescaled.
+        """
+        for name, param in self.model.named_parameters():
+            if param.requires_grad is True and "trainable" not in name:
+                param.register_hook(lambda grad: grad * float(self.model_comm_group_size))
diff --git a/src/anemoi/training/losses/__init__.py b/src/anemoi/training/losses/__init__.py
new file mode 100644
index 00000000..33d7fa0a
--- /dev/null
+++ b/src/anemoi/training/losses/__init__.py
@@ -0,0 +1,8 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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.
+#
diff --git a/src/anemoi/training/losses/mse.py b/src/anemoi/training/losses/mse.py
new file mode 100644
index 00000000..b72a7aea
--- /dev/null
+++ b/src/anemoi/training/losses/mse.py
@@ -0,0 +1,91 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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 __future__ import annotations
+
+import logging
+
+import torch
+from torch import nn
+
+LOGGER = logging.getLogger(__name__)
+
+
+class WeightedMSELoss(nn.Module):
+    """Latitude-weighted MSE loss."""
+
+    def __init__(
+        self,
+        node_weights: torch.Tensor,
+        data_variances: torch.Tensor | None = None,
+        ignore_nans: bool | None = False,
+    ) -> None:
+        """Latitude- and (inverse-)variance-weighted MSE Loss.
+
+        Parameters
+        ----------
+        node_weights : torch.Tensor of shape (N, )
+            Weight of each node in the loss function
+        data_variances : Optional[torch.Tensor], optional
+            precomputed, per-variable stepwise variance estimate, by default None
+        ignore_nans : bool, optional
+            Allow nans in the loss and apply methods ignoring nans for measuring the loss, by default False
+
+        """
+        super().__init__()
+
+        self.avg_function = torch.nanmean if ignore_nans else torch.mean
+        self.sum_function = torch.nansum if ignore_nans else torch.sum
+
+        self.register_buffer("weights", node_weights, persistent=True)
+        if data_variances is not None:
+            self.register_buffer("ivar", data_variances, persistent=True)
+
+    def forward(
+        self,
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        squash: bool = True,
+    ) -> torch.Tensor:
+        """Calculates the lat-weighted MSE loss.
+
+        Parameters
+        ----------
+        pred : torch.Tensor
+            Prediction tensor, shape (bs, lat*lon, n_outputs)
+        target : torch.Tensor
+            Target tensor, shape (bs, lat*lon, n_outputs)
+        squash : bool, optional
+            Average last dimension, by default True
+
+        Returns
+        -------
+        torch.Tensor
+            Weighted MSE loss
+
+        """
+        out = torch.square(pred - target)
+
+        # Use variances if available
+        if hasattr(self, "ivar"):
+            out *= self.ivar
+
+        # Squash by last dimension
+        if squash:
+            out = self.avg_function(out, dim=-1)
+            # Weight by area
+            out *= self.weights.expand_as(out)
+            out /= self.sum_function(self.weights.expand_as(out))
+            return self.sum_function(out)
+
+        # Weight by area
+        out *= self.weights[..., None].expand_as(out)
+        # keep last dimension (variables) when summing weights
+        out /= self.sum_function(self.weights[..., None].expand_as(out), axis=(0, 1, 2))
+        return self.sum_function(out, axis=(0, 1, 2))
diff --git a/src/anemoi/training/losses/utils.py b/src/anemoi/training/losses/utils.py
new file mode 100644
index 00000000..9a866a0a
--- /dev/null
+++ b/src/anemoi/training/losses/utils.py
@@ -0,0 +1,53 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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 __future__ import annotations
+
+import logging
+
+import torch
+from torch import nn
+
+LOGGER = logging.getLogger(__name__)
+
+
+def grad_scaler(
+    module: nn.Module,
+    grad_in: tuple[torch.Tensor, ...],
+    grad_out: tuple[torch.Tensor, ...],
+) -> tuple[torch.Tensor, ...] | None:
+    """Scales the loss gradients.
+
+    Uses the formula in https://arxiv.org/pdf/2306.06079.pdf, section 4.3.2
+
+    Use <module>.register_full_backward_hook(grad_scaler, prepend=False) to register this hook.
+
+    Parameters
+    ----------
+    module : nn.Module
+        Loss object (not used)
+    grad_in : tuple[torch.Tensor, ...]
+        Loss gradients
+    grad_out : tuple[torch.Tensor, ...]
+        Output gradients (not used)
+
+    Returns
+    -------
+    tuple[torch.Tensor, ...]
+        Re-scaled input gradients
+
+    """
+    del module, grad_out
+    # first grad_input is that of the predicted state and the second is that of the "ground truth" (== zero)
+    channels = grad_in[0].shape[-1]  # number of channels
+    channel_weights = torch.reciprocal(torch.sum(torch.abs(grad_in[0]), dim=1, keepdim=True))  # channel-wise weights
+    new_grad_in = (
+        (channels * channel_weights) / torch.sum(channel_weights, dim=-1, keepdim=True) * grad_in[0]
+    )  # rescaled gradient
+    return new_grad_in, grad_in[1]
diff --git a/src/anemoi/training/train/__init__.py b/src/anemoi/training/train/__init__.py
new file mode 100644
index 00000000..282d6a69
--- /dev/null
+++ b/src/anemoi/training/train/__init__.py
@@ -0,0 +1,6 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
diff --git a/src/anemoi/training/train/forecaster.py b/src/anemoi/training/train/forecaster.py
new file mode 100644
index 00000000..ff1acfd7
--- /dev/null
+++ b/src/anemoi/training/train/forecaster.py
@@ -0,0 +1,370 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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.
+#
+
+import logging
+import math
+import os
+from collections import defaultdict
+from collections.abc import Mapping
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from anemoi.models.data_indices.collection import IndexCollection
+from anemoi.models.interface import AnemoiModelInterface
+from anemoi.utils.config import DotDict
+from hydra.utils import instantiate
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
+from timm.scheduler import CosineLRScheduler
+from torch.distributed.distributed_c10d import ProcessGroup
+from torch.distributed.optim import ZeroRedundancyOptimizer
+from torch.utils.checkpoint import checkpoint
+from torch_geometric.data import HeteroData
+
+from anemoi.training.losses.mse import WeightedMSELoss
+from anemoi.training.losses.utils import grad_scaler
+from anemoi.training.utils.jsonify import map_config_to_primitives
+
+LOGGER = logging.getLogger(__name__)
+
+
+class GraphForecaster(pl.LightningModule):
+    """Graph neural network forecaster for PyTorch Lightning."""
+
+    def __init__(
+        self,
+        *,
+        config: DictConfig,
+        graph_data: HeteroData,
+        statistics: dict,
+        data_indices: IndexCollection,
+        metadata: dict,
+    ) -> None:
+        """Initialize graph neural network forecaster.
+
+        Parameters
+        ----------
+        config : DictConfig
+            Job configuration
+        graph_data : HeteroData
+            Graph object
+        statistics : dict
+            Statistics of the training data
+        data_indices : IndexCollection
+            Indices of the training data,
+        metadata : dict
+            Provenance information
+
+        """
+        super().__init__()
+
+        graph_data = graph_data.to(self.device)
+
+        self.model = AnemoiModelInterface(
+            statistics=statistics,
+            data_indices=data_indices,
+            metadata=metadata,
+            graph_data=graph_data,
+            config=DotDict(map_config_to_primitives(OmegaConf.to_container(config, resolve=True))),
+        )
+
+        self.data_indices = data_indices
+
+        self.save_hyperparameters()
+
+        self.latlons_data = graph_data[config.graph.data].x
+        self.loss_weights = graph_data[config.graph.data][config.model.node_loss_weight].squeeze()
+
+        self.logger_enabled = config.diagnostics.log.wandb.enabled or config.diagnostics.log.mlflow.enabled
+
+        self.metric_ranges, self.metric_ranges_validation, loss_scaling = self.metrics_loss_scaling(
+            config,
+            data_indices,
+        )
+        self.loss = WeightedMSELoss(node_weights=self.loss_weights, data_variances=loss_scaling)
+        self.metrics = WeightedMSELoss(node_weights=self.loss_weights, ignore_nans=True)
+
+        if config.training.loss_gradient_scaling:
+            self.loss.register_full_backward_hook(grad_scaler, prepend=False)
+
+        self.multi_step = config.training.multistep_input
+        self.lr = (
+            config.hardware.num_nodes
+            * config.hardware.num_gpus_per_node
+            * config.training.lr.rate
+            / config.hardware.num_gpus_per_model
+        )
+        self.lr_iterations = config.training.lr.iterations
+        self.lr_min = config.training.lr.min
+        self.rollout = config.training.rollout.start
+        self.rollout_epoch_increment = config.training.rollout.epoch_increment
+        self.rollout_max = config.training.rollout.max
+
+        self.use_zero_optimizer = config.training.zero_optimizer
+
+        self.model_comm_group = None
+
+        LOGGER.debug("Rollout window length: %d", self.rollout)
+        LOGGER.debug("Rollout increase every : %d epochs", self.rollout_epoch_increment)
+        LOGGER.debug("Rollout max : %d", self.rollout_max)
+        LOGGER.debug("Multistep: %d", self.multi_step)
+
+        self.enable_plot = config.diagnostics.plot.enabled
+
+        self.model_comm_group_id = int(os.environ.get("SLURM_PROCID", "0")) // config.hardware.num_gpus_per_model
+        self.model_comm_group_rank = int(os.environ.get("SLURM_PROCID", "0")) % config.hardware.num_gpus_per_model
+        self.model_comm_num_groups = math.ceil(
+            config.hardware.num_gpus_per_node * config.hardware.num_nodes / config.hardware.num_gpus_per_model,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x, self.model_comm_group)
+
+    @staticmethod
+    def metrics_loss_scaling(config: DictConfig, data_indices: IndexCollection) -> tuple[dict, torch.Tensor]:
+        metric_ranges = defaultdict(list)
+        metric_ranges_validation = defaultdict(list)
+        loss_scaling = (
+            np.ones((len(data_indices.internal_data.output.full),), dtype=np.float32)
+            * config.training.loss_scaling.default
+        )
+
+        pressure_level = instantiate(config.training.pressure_level_scaler)
+
+        LOGGER.info(
+            "Pressure level scaling: use scaler %s with slope %.4f and minimum %.2f",
+            type(pressure_level).__name__,
+            pressure_level.slope,
+            pressure_level.minimum,
+        )
+
+        for key, idx in data_indices.internal_model.output.name_to_index.items():
+            # Split pressure levels on "_" separator
+            split = key.split("_")
+            if len(split) > 1 and split[-1].isdigit():
+                # Create grouped metrics for pressure levels (e.g. Q, T, U, V, etc.) for logger
+                metric_ranges[f"pl_{split[0]}"].append(idx)
+                # Create pressure levels in loss scaling vector
+                if split[0] in config.training.loss_scaling.pl:
+                    loss_scaling[idx] = config.training.loss_scaling.pl[split[0]] * pressure_level.scaler(
+                        int(split[-1]),
+                    )
+                else:
+                    LOGGER.debug("Parameter %s was not scaled.", key)
+            else:
+                metric_ranges[f"sfc_{key}"].append(idx)
+                # Create surface variables in loss scaling vector
+                if key in config.training.loss_scaling.sfc:
+                    loss_scaling[idx] = config.training.loss_scaling.sfc[key]
+                else:
+                    LOGGER.debug("Parameter %s was not scaled.", key)
+            # Create specific metrics from hydra to log in logger
+            if key in config.training.metrics:
+                metric_ranges[key] = [idx]
+        loss_scaling = torch.from_numpy(loss_scaling)
+        # metric for validation, after postprocessing
+        for key, idx in data_indices.model.output.name_to_index.items():
+            # Split pressure levels on "_" separator
+            split = key.split("_")
+            if len(split) > 1 and split[1].isdigit():
+                # Create grouped metrics for pressure levels (e.g. Q, T, U, V, etc.) for logger
+                metric_ranges_validation[f"pl_{split[0]}"].append(idx)
+            else:
+                metric_ranges_validation[f"sfc_{key}"].append(idx)
+            # Create specific metrics from hydra to log in logger
+            if key in config.training.metrics:
+                metric_ranges_validation[key] = [idx]
+        return metric_ranges, metric_ranges_validation, loss_scaling
+
+    def set_model_comm_group(self, model_comm_group: ProcessGroup) -> None:
+        LOGGER.debug("set_model_comm_group: %s", model_comm_group)
+        self.model_comm_group = model_comm_group
+
+    def advance_input(
+        self,
+        x: torch.Tensor,
+        y_pred: torch.Tensor,
+        batch: torch.Tensor,
+        rollout_step: int,
+    ) -> torch.Tensor:
+        x = x.roll(-1, dims=1)
+
+        # Get prognostic variables
+        x[:, -1, :, :, self.data_indices.internal_model.input.prognostic] = y_pred[
+            ...,
+            self.data_indices.internal_model.output.prognostic,
+        ]
+
+        # get new "constants" needed for time-varying fields
+        x[:, -1, :, :, self.data_indices.internal_model.input.forcing] = batch[
+            :,
+            self.multi_step + rollout_step,
+            :,
+            :,
+            self.data_indices.internal_data.input.forcing,
+        ]
+        return x
+
+    def _step(
+        self,
+        batch: torch.Tensor,
+        batch_idx: int,
+        validation_mode: bool = False,
+    ) -> tuple[torch.Tensor, Mapping[str, torch.Tensor]]:
+        del batch_idx
+        loss = torch.zeros(1, dtype=batch.dtype, device=self.device, requires_grad=False)
+        # for validation not normalized in-place because remappers cannot be applied in-place
+        batch = self.model.pre_processors(batch, in_place=not validation_mode)
+        metrics = {}
+
+        # start rollout of preprocessed batch
+        x = batch[
+            :,
+            0 : self.multi_step,
+            ...,
+            self.data_indices.internal_data.input.full,
+        ]  # (bs, multi_step, latlon, nvar)
+
+        y_preds = []
+        for rollout_step in range(self.rollout):
+            # prediction at rollout step rollout_step, shape = (bs, latlon, nvar)
+            y_pred = self(x)
+
+            y = batch[:, self.multi_step + rollout_step, ..., self.data_indices.internal_data.output.full]
+            # y includes the auxiliary variables, so we must leave those out when computing the loss
+            loss += checkpoint(self.loss, y_pred, y, use_reentrant=False)
+
+            x = self.advance_input(x, y_pred, batch, rollout_step)
+
+            if validation_mode:
+                metrics_next, y_preds_next = self.calculate_val_metrics(
+                    y_pred,
+                    y,
+                    rollout_step,
+                    enable_plot=self.enable_plot,
+                )
+                metrics.update(metrics_next)
+                y_preds.extend(y_preds_next)
+
+        # scale loss
+        loss *= 1.0 / self.rollout
+        return loss, metrics, y_preds
+
+    def calculate_val_metrics(
+        self,
+        y_pred: torch.Tensor,
+        y: torch.Tensor,
+        rollout_step: int,
+        enable_plot: bool = False,
+    ) -> tuple[dict, list]:
+        metrics = {}
+        y_preds = []
+        y_postprocessed = self.model.post_processors(y, in_place=False)
+        y_pred_postprocessed = self.model.post_processors(y_pred, in_place=False)
+        for mkey, indices in self.metric_ranges_validation.items():
+            metrics[f"{mkey}_{rollout_step + 1}"] = self.metrics(
+                y_pred_postprocessed[..., indices],
+                y_postprocessed[..., indices],
+            )
+
+        if enable_plot:
+            y_preds.append(y_pred)
+        return metrics, y_preds
+
+    def training_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor:
+        train_loss, _, _ = self._step(batch, batch_idx)
+        self.log(
+            "train_wmse",
+            train_loss,
+            on_epoch=True,
+            on_step=True,
+            prog_bar=True,
+            logger=self.logger_enabled,
+            batch_size=batch.shape[0],
+            sync_dist=True,
+        )
+        self.log(
+            "rollout",
+            float(self.rollout),
+            on_step=True,
+            logger=self.logger_enabled,
+            rank_zero_only=True,
+            sync_dist=False,
+        )
+        return train_loss
+
+    def lr_scheduler_step(self, scheduler: CosineLRScheduler, metric: None = None) -> None:
+        """Step the learning rate scheduler by Pytorch Lightning.
+
+        Parameters
+        ----------
+        scheduler : CosineLRScheduler
+            Learning rate scheduler object.
+        metric : Optional[Any]
+            Metric object for e.g. ReduceLRonPlateau. Default is None.
+
+        """
+        del metric
+        scheduler.step(epoch=self.trainer.global_step)
+
+    def on_train_epoch_end(self) -> None:
+        if self.rollout_epoch_increment > 0 and self.current_epoch % self.rollout_epoch_increment == 0:
+            self.rollout += 1
+            LOGGER.debug("Rollout window length: %d", self.rollout)
+        self.rollout = min(self.rollout, self.rollout_max)
+
+    def validation_step(self, batch: torch.Tensor, batch_idx: int) -> None:
+        with torch.no_grad():
+            val_loss, metrics, y_preds = self._step(batch, batch_idx, validation_mode=True)
+        self.log(
+            "val_wmse",
+            val_loss,
+            on_epoch=True,
+            on_step=True,
+            prog_bar=True,
+            logger=self.logger_enabled,
+            batch_size=batch.shape[0],
+            sync_dist=True,
+        )
+        for mname, mvalue in metrics.items():
+            self.log(
+                "val_" + mname,
+                mvalue,
+                on_epoch=True,
+                on_step=False,
+                prog_bar=False,
+                logger=self.logger_enabled,
+                batch_size=batch.shape[0],
+                sync_dist=True,
+            )
+        return val_loss, y_preds
+
+    def configure_optimizers(self) -> tuple[list[torch.optim.Optimizer], list[dict]]:
+        if self.use_zero_optimizer:
+            optimizer = ZeroRedundancyOptimizer(
+                self.trainer.model.parameters(),
+                optimizer_class=torch.optim.AdamW,
+                betas=(0.9, 0.95),
+                lr=self.lr,
+            )
+        else:
+            optimizer = torch.optim.AdamW(
+                self.trainer.model.parameters(),
+                betas=(0.9, 0.95),
+                lr=self.lr,
+            )  # , fused=True)
+
+        scheduler = CosineLRScheduler(
+            optimizer,
+            lr_min=self.lr_min,
+            t_initial=self.lr_iterations,
+            warmup_t=1000,
+        )
+        return [optimizer], [{"scheduler": scheduler, "interval": "step"}]
diff --git a/src/anemoi/training/train/train.py b/src/anemoi/training/train/train.py
new file mode 100644
index 00000000..f48b9467
--- /dev/null
+++ b/src/anemoi/training/train/train.py
@@ -0,0 +1,359 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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 __future__ import annotations
+
+import datetime
+import logging
+import os
+from functools import cached_property
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import hydra
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from anemoi.utils.provenance import gather_provenance_info
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
+from pytorch_lightning.profilers import PyTorchProfiler
+from pytorch_lightning.utilities.rank_zero import rank_zero_only
+
+from anemoi.training.data.datamodule import AnemoiDatasetsDataModule
+from anemoi.training.diagnostics.callbacks import get_callbacks
+from anemoi.training.diagnostics.logger import get_mlflow_logger
+from anemoi.training.diagnostics.logger import get_tensorboard_logger
+from anemoi.training.diagnostics.logger import get_wandb_logger
+from anemoi.training.distributed.strategy import DDPGroupStrategy
+from anemoi.training.train.forecaster import GraphForecaster
+from anemoi.training.utils.jsonify import map_config_to_primitives
+from anemoi.training.utils.seeding import get_base_seed
+
+if TYPE_CHECKING:
+    from torch_geometric.data import HeteroData
+
+LOGGER = logging.getLogger(__name__)
+
+
+class AnemoiTrainer:
+    """Utility class for training the model."""
+
+    def __init__(self, config: DictConfig) -> None:
+        """Initialize the Anemoi trainer.
+
+        Parameters
+        ----------
+        config : DictConfig
+            Config object from Hydra.
+
+        """
+        # Allow for lower internal precision of float32 matrix multiplications.
+        # This can increase performance (and TensorCore usage, where available).
+        torch.set_float32_matmul_precision("high")
+        # Resolve the config to avoid shenanigans with lazy loading
+        OmegaConf.resolve(config)
+        self.config = config
+
+        # Default to not warm-starting from a checkpoint
+        self.start_from_checkpoint = bool(self.config.training.run_id) or bool(self.config.training.fork_run_id)
+        self.load_weights_only = config.training.load_weights_only
+        self.parent_uuid = None
+
+        self.config.training.run_id = self.run_id
+        LOGGER.info("Run id: %s", self.config.training.run_id)
+        # Update paths to contain the run ID
+        self._update_paths()
+
+        self._log_information()
+
+    @cached_property
+    def datamodule(self) -> AnemoiDatasetsDataModule:
+        """DataModule instance and DataSets."""
+        datamodule = AnemoiDatasetsDataModule(self.config)
+        self.config.data.num_features = len(datamodule.ds_train.data.variables)
+        return datamodule
+
+    @cached_property
+    def data_indices(self) -> dict:
+        """Returns a dictionary of data indices.
+
+        This is used to slice the data.
+        """
+        return self.datamodule.data_indices
+
+    @cached_property
+    def initial_seed(self) -> int:
+        """Initial seed for the RNG.
+
+        This sets the same initial seed for all ranks. Ranks are re-seeded in the
+        strategy to account for model communication groups.
+        """
+        initial_seed = get_base_seed()
+        rnd_seed = pl.seed_everything(initial_seed, workers=True)
+        np_rng = np.random.default_rng(rnd_seed)
+        (torch.rand(1), np_rng.random())
+        LOGGER.debug(
+            "Initial seed: Rank %d, initial seed %d, running with random seed: %d",
+            int(os.environ.get("SLURM_PROCID", "0")),
+            initial_seed,
+            rnd_seed,
+        )
+        return initial_seed
+
+    @cached_property
+    def graph_data(self) -> HeteroData:
+        """Graph data.
+
+        Creates the graph in all workers.
+        """
+        graph_filename = Path(
+            self.config.hardware.paths.graph,
+            self.config.hardware.files.graph,
+        )
+
+        if graph_filename.exists() and not self.config.graph.overwrite:
+            LOGGER.info("Loading graph data from %s", graph_filename)
+            return torch.load(graph_filename)
+
+        from anemoi.graphs.create import GraphCreator
+
+        return GraphCreator(config=self.config.graph).create(
+            save_path=graph_filename,
+            overwrite=self.config.graph.overwrite,
+        )
+
+    @cached_property
+    def model(self) -> GraphForecaster:
+        """Provide the model instance."""
+        kwargs = {
+            "config": self.config,
+            "data_indices": self.data_indices,
+            "graph_data": self.graph_data,
+            "metadata": self.metadata,
+            "statistics": self.datamodule.statistics,
+        }
+        if self.load_weights_only:
+            LOGGER.info("Restoring only model weights from %s", self.last_checkpoint)
+            return GraphForecaster.load_from_checkpoint(self.last_checkpoint, **kwargs)
+        return GraphForecaster(**kwargs)
+
+    @rank_zero_only
+    def _get_mlflow_run_id(self) -> str:
+        run_id = self.mlflow_logger.run_id
+        # for resumed runs or offline runs logging this can be uesful
+        LOGGER.info("Mlflow Run id: %s", run_id)
+        return run_id
+
+    @cached_property
+    def run_id(self) -> str:
+        """Unique identifier for the current run."""
+        if self.config.training.run_id and not self.config.training.fork_run_id:
+            # Return the provided run ID - reuse run_id if resuming run
+            return self.config.training.run_id
+
+        if self.config.diagnostics.log.mlflow.enabled:
+            # if using mlflow with a new run get the run_id from mlflow
+            return self._get_mlflow_run_id()
+
+        # Generate a random UUID
+        import uuid
+
+        return str(uuid.uuid4())
+
+    @cached_property
+    def wandb_logger(self) -> pl.loggers.WandbLogger:
+        """WandB logger."""
+        return get_wandb_logger(self.config, self.model)
+
+    @cached_property
+    def mlflow_logger(self) -> pl.loggers.MLFlowLogger:
+        """Mlflow logger."""
+        return get_mlflow_logger(self.config)
+
+    @cached_property
+    def tensorboard_logger(self) -> pl.loggers.TensorBoardLogger:
+        """TensorBoard logger."""
+        return get_tensorboard_logger(self.config)
+
+    @cached_property
+    def last_checkpoint(self) -> str | None:
+        """Path to the last checkpoint."""
+        if not self.start_from_checkpoint:
+            return None
+
+        checkpoint = Path(
+            self.config.hardware.paths.checkpoints.parent,
+            self.config.training.fork_run_id or self.run_id,
+            self.config.hardware.files.warm_start or "last.ckpt",
+        )
+
+        # Check if the last checkpoint exists
+        if Path(checkpoint).exists():
+            LOGGER.info("Resuming training from last checkpoint: %s", checkpoint)
+            return checkpoint
+
+        LOGGER.warning("Could not find last checkpoint: %s", checkpoint)
+        return None
+
+    @cached_property
+    def callbacks(self) -> list[pl.callbacks.Callback]:
+        return get_callbacks(self.config)
+
+    @cached_property
+    def metadata(self) -> dict:
+        """Metadata and provenance information."""
+        return map_config_to_primitives(
+            {
+                "version": "1.0",
+                "config": self.config,
+                "seed": self.initial_seed,
+                "run_id": self.run_id,
+                "dataset": self.datamodule.metadata,
+                "data_indices": self.datamodule.data_indices,
+                "provenance_training": gather_provenance_info(),
+                "timestamp": datetime.datetime.now(tz=datetime.timezone.utc),
+            },
+        )
+
+    @cached_property
+    def profiler(self) -> PyTorchProfiler | None:
+        """Returns a pytorch profiler object, if profiling is enabled."""
+        if self.config.diagnostics.profiler:
+            assert (
+                self.config.diagnostics.log.tensorboard.enabled
+            ), "Tensorboard logging must be enabled when profiling! Check your job config."
+            return PyTorchProfiler(
+                dirpath=self.config.hardware.paths.logs.tensorboard,
+                filename="anemoi-profiler",
+                export_to_chrome=False,
+                # profiler-specific keywords
+                activities=[
+                    # torch.profiler.ProfilerActivity.CPU,  # this is memory-hungry
+                    torch.profiler.ProfilerActivity.CUDA,
+                ],
+                schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
+                on_trace_ready=torch.profiler.tensorboard_trace_handler(
+                    dir_name=self.config.hardware.paths.logs.tensorboard,
+                ),
+                profile_memory=True,
+                record_shapes=True,
+                with_stack=True,
+            )
+        return None
+
+    @cached_property
+    def loggers(self) -> list:
+        loggers = []
+        if self.config.diagnostics.log.wandb.enabled:
+            loggers.append(self.wandb_logger)
+        if self.config.diagnostics.log.tensorboard.enabled:
+            loggers.append(self.tensorboard_logger)
+        if self.config.diagnostics.log.mlflow.enabled:
+            loggers.append(self.mlflow_logger)
+        return loggers
+
+    @cached_property
+    def accelerator(self) -> str:
+        assert self.config.hardware.accelerator in {
+            "auto",
+            "cpu",
+            "gpu",
+            "cuda",
+            "tpu",
+        }, f"Invalid accelerator ({self.config.hardware.accelerator}) in hardware config."
+        if self.config.hardware.accelerator == "cpu":
+            LOGGER.info("WARNING: Accelerator set to CPU, this should only be used for debugging.")
+        return self.config.hardware.accelerator
+
+    def _log_information(self) -> None:
+        # Log number of variables (features)
+        num_fc_features = len(self.datamodule.ds_train.data.variables) - len(self.config.data.forcing)
+        LOGGER.debug("Total number of prognostic variables: %d", num_fc_features)
+        LOGGER.debug("Total number of auxiliary variables: %d", len(self.config.data.forcing))
+
+        # Log learning rate multiplier when running single-node, multi-GPU and/or multi-node
+        total_number_of_model_instances = (
+            self.config.hardware.num_nodes
+            * self.config.hardware.num_gpus_per_node
+            / self.config.hardware.num_gpus_per_model
+        )
+        LOGGER.debug(
+            "Total GPU count / model group size: %d - NB: the learning rate will be scaled by this factor!",
+            total_number_of_model_instances,
+        )
+        LOGGER.debug("Effective learning rate: %.3e", total_number_of_model_instances * self.config.training.lr.rate)
+        LOGGER.debug("Rollout window length: %d", self.config.training.rollout.start)
+
+    def _update_paths(self) -> None:
+        """Update the paths in the configuration."""
+        if self.run_id:  # when using mlflow only rank0 will have a run_id except when resuming runs
+            # Multi-gpu new runs or forked runs - only rank 0
+            # Multi-gpu resumed runs - all ranks
+            self.config.hardware.paths.checkpoints = Path(self.config.hardware.paths.checkpoints, self.run_id)
+            self.config.hardware.paths.plots = Path(self.config.hardware.paths.plots, self.run_id)
+        elif self.config.training.fork_run_id:
+            # Only rank non zero in the forked run will go here
+            parent_run = self.config.training.fork_run_id
+            self.config.hardware.paths.checkpoints = Path(self.config.hardware.paths.checkpoints, parent_run)
+
+    @cached_property
+    def strategy(self) -> DDPGroupStrategy:
+        """Training strategy."""
+        return DDPGroupStrategy(
+            self.config.hardware.num_gpus_per_model,
+            static_graph=not self.config.training.accum_grad_batches > 1,
+        )
+
+    def train(self) -> None:
+        """Training entry point."""
+        trainer = pl.Trainer(
+            accelerator=self.accelerator,
+            callbacks=self.callbacks,
+            deterministic=self.config.training.deterministic,
+            detect_anomaly=self.config.diagnostics.debug.anomaly_detection,
+            strategy=self.strategy,
+            devices=self.config.hardware.num_gpus_per_node,
+            num_nodes=self.config.hardware.num_nodes,
+            precision=self.config.training.precision,
+            max_epochs=self.config.training.max_epochs,
+            logger=self.loggers,
+            log_every_n_steps=self.config.diagnostics.log.interval,
+            # run a fixed no of batches per epoch (helpful when debugging)
+            limit_train_batches=self.config.dataloader.limit_batches.training,
+            limit_val_batches=self.config.dataloader.limit_batches.validation,
+            num_sanity_val_steps=4,
+            accumulate_grad_batches=self.config.training.accum_grad_batches,
+            gradient_clip_val=self.config.training.gradient_clip.val,
+            gradient_clip_algorithm=self.config.training.gradient_clip.algorithm,
+            # we have our own DDP-compliant sampler logic baked into the dataset
+            use_distributed_sampler=False,
+            profiler=self.profiler,
+            enable_progress_bar=self.config.diagnostics.enable_progress_bar,
+        )
+
+        trainer.fit(
+            self.model,
+            datamodule=self.datamodule,
+            ckpt_path=None if self.load_weights_only else self.last_checkpoint,
+        )
+
+        if self.config.diagnostics.print_memory_summary:
+            LOGGER.debug("memory summary: %s", torch.cuda.memory_summary())
+
+        LOGGER.debug("---- DONE. ----")
+
+
+@hydra.main(version_base=None, config_path="../config", config_name="config")
+def main(config: DictConfig) -> None:
+    AnemoiTrainer(config).train()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/anemoi/training/utils/__init__.py b/src/anemoi/training/utils/__init__.py
new file mode 100644
index 00000000..33d7fa0a
--- /dev/null
+++ b/src/anemoi/training/utils/__init__.py
@@ -0,0 +1,8 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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.
+#
diff --git a/src/anemoi/training/utils/checkpoint.py b/src/anemoi/training/utils/checkpoint.py
new file mode 100644
index 00000000..697a92de
--- /dev/null
+++ b/src/anemoi/training/utils/checkpoint.py
@@ -0,0 +1,64 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from anemoi.utils.checkpoints import save_metadata
+
+from anemoi.training.train.forecaster import GraphForecaster
+
+
+def load_and_prepare_model(lightning_checkpoint_path: str) -> tuple[torch.nn.Module, dict]:
+    """Load the lightning checkpoint and extract the pytorch model and its metadata.
+
+    Parameters
+    ----------
+    lightning_checkpoint_path : str
+        path to lightning checkpoint
+
+    Returns
+    -------
+    tuple[torch.nn.Module, dict]
+        pytorch model, metadata
+
+    """
+    module = GraphForecaster.load_from_checkpoint(lightning_checkpoint_path)
+    model = module.model
+
+    metadata = dict(**model.metadata)
+    model.metadata = None
+    model.config = None
+
+    return model, metadata
+
+
+def save_inference_checkpoint(model: torch.nn.Module, metadata: dict, save_path: Path | str) -> Path:
+    """Save a pytorch checkpoint for inference with the model metadata.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+        Pytorch model
+    metadata : dict
+        Anemoi Metadata to inject into checkpoint
+    save_path : Path | str
+        Directory to save anemoi checkpoint
+
+    Returns
+    -------
+    Path
+        Path to saved checkpoint
+    """
+    save_path = Path(save_path)
+    inference_filepath = save_path.parent / f"inference-{save_path.name}"
+
+    torch.save(model, inference_filepath)
+    save_metadata(inference_filepath, metadata)
+    return inference_filepath
diff --git a/src/anemoi/training/utils/jsonify.py b/src/anemoi/training/utils/jsonify.py
new file mode 100644
index 00000000..ddf9b86c
--- /dev/null
+++ b/src/anemoi/training/utils/jsonify.py
@@ -0,0 +1,64 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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.
+
+import datetime
+from pathlib import Path
+
+import torch
+from anemoi.models.data_indices.collection import BaseIndex
+from anemoi.models.data_indices.collection import IndexCollection
+from anemoi.models.data_indices.tensor import BaseTensorIndex
+from anemoi.utils.dates import frequency_to_string
+from omegaconf import DictConfig
+from omegaconf import ListConfig
+from omegaconf import OmegaConf
+
+
+def map_config_to_primitives(config: OmegaConf) -> dict:
+    """Ensure that the metadata information is JSON-serializable.
+
+    Parameters
+    ----------
+    config : OmegaConf
+        config object to be mapped to primitives.
+
+    Returns
+    -------
+    dict
+        JSON serializable dictionary.
+
+    Raises
+    ------
+    TypeError
+        Cannot map object to primitives.
+
+    """
+    if config is None or isinstance(config, (int, float, str, bool)):
+        return config
+
+    if isinstance(config, Path):
+        config = str(config)
+    elif isinstance(config, datetime.date):
+        config = config.isoformat()
+    elif isinstance(config, datetime.timedelta):
+        config = frequency_to_string(config)
+    elif isinstance(config, (list, tuple)):
+        config = [map_config_to_primitives(v) for v in config]
+    elif isinstance(config, dict):
+        config = {k: map_config_to_primitives(v) for k, v in config.items()}
+    elif isinstance(config, (DictConfig, ListConfig)):
+        config = map_config_to_primitives(OmegaConf.to_container(config, resolve=True))
+    elif isinstance(config, torch.Tensor):
+        config = map_config_to_primitives(config.tolist())
+    elif isinstance(config, (IndexCollection, BaseTensorIndex, BaseIndex)):
+        config = map_config_to_primitives(config.todict())
+    else:
+        msg = f"Cannot serialize object of type {type(config)}"
+        raise TypeError(msg)
+
+    return config
diff --git a/src/anemoi/training/utils/mlflow_sync.py b/src/anemoi/training/utils/mlflow_sync.py
new file mode 100644
index 00000000..52ceee2f
--- /dev/null
+++ b/src/anemoi/training/utils/mlflow_sync.py
@@ -0,0 +1,273 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+import logging
+import os
+import tempfile
+from itertools import starmap
+from pathlib import Path
+
+
+def export_log_output_file_path() -> tempfile._TemporaryFileWrapper:
+    # to set up the file where to send the output logs
+    if not os.getenv("TMPDIR") and not os.getenv("SCRATCH"):
+        error_msg = "Please set one of those variables TMPDIR:{} or SCRATCH:{} to proceed.".format(
+            os.environ["SCRATCH"],
+            os.environ["TMPDIR"],
+        )
+        raise ValueError(error_msg)
+
+    tmpdir = os.environ["TMPDIR"] if os.getenv("TMPDIR") else os.environ["SCRATCH"]
+    user = os.environ["USER"]
+    Path(tmpdir).mkdir(parents=True, exist_ok=True)
+    temp = tempfile.NamedTemporaryFile(dir=tmpdir, prefix=f"{user}_")  # noqa: SIM115
+    os.environ["MLFLOW_EXPORT_IMPORT_LOG_OUTPUT_FILE"] = temp.name
+    return temp
+
+
+temp = export_log_output_file_path()
+
+import mlflow  # noqa: E402
+from mlflow.entities import RunStatus  # noqa: E402
+from mlflow.entities import RunTag  # noqa: E402
+from mlflow.tracking.context.default_context import _get_user  # noqa: E402
+from mlflow.utils.mlflow_tags import MLFLOW_USER  # noqa: E402
+from mlflow.utils.validation import MAX_METRICS_PER_BATCH  # noqa: E402
+from mlflow.utils.validation import MAX_PARAMS_TAGS_PER_BATCH  # noqa: E402
+
+try:
+    import mlflow_export_import.common.utils as mlflow_utils
+    from mlflow_export_import.client.client_utils import create_http_client
+    from mlflow_export_import.run.export_run import _get_metrics_with_steps
+    from mlflow_export_import.run.export_run import _inputs_to_dict
+    from mlflow_export_import.run.import_run import _import_inputs
+    from mlflow_export_import.run.run_data_importer import _log_data
+    from mlflow_export_import.run.run_data_importer import _log_metrics
+    from mlflow_export_import.run.run_data_importer import _log_params
+except ImportError:
+    msg = "The 'mlflow-export-import' package is not installed. Please install it from https://github.com/mlflow/mlflow-export-import"
+    raise ImportError(msg) from None
+
+LOGGER = logging.getLogger(__name__)
+
+
+# # This functions are based on the existing functions in mlflow_export_import.run.run_data_importer.py
+def _log_tags(client: mlflow.MlflowClient, run_dct: dict, run_id: str, batch_size: int, src_user_id: str) -> None:
+    def get_data(run_dct: dict, *args) -> list:
+        del args  # unused
+        tags = run_dct["tags"]
+        tags = list(starmap(RunTag, tags.items()))
+        user_id = _get_user()
+        tags.append(RunTag(MLFLOW_USER, user_id))
+
+        return tags
+
+    def log_data(run_id: str, tags: list) -> None:
+        client.log_batch(run_id, tags=tags)
+
+    args_get = {
+        "src_user_id": src_user_id,
+    }
+
+    _log_data(run_dct, run_id, batch_size, get_data, log_data, args_get)
+
+
+def import_run_data(mlflow_client: mlflow.MlflowClient, run_dct: dict, run_id: str, src_user_id: str) -> None:
+    _log_params(mlflow_client, run_dct, run_id, MAX_PARAMS_TAGS_PER_BATCH)
+    _log_metrics(mlflow_client, run_dct, run_id, MAX_METRICS_PER_BATCH)
+    _log_tags(
+        mlflow_client,
+        run_dct,
+        run_id,
+        MAX_PARAMS_TAGS_PER_BATCH,
+        src_user_id,
+    )
+
+
+class MlFlowSync:
+    """Class to sync an offline run to the destination tracking uri."""
+
+    def __init__(
+        self,
+        source_tracking_uri: str,
+        dest_tracking_uri: str,
+        run_id: str,
+        experiment_name: str = "anemoi_debug",
+        export_deleted_runs: bool = False,
+        log_level: str = "INFO",
+    ) -> None:
+        self.source_tracking_uri = source_tracking_uri
+        self.dest_tracking_uri = dest_tracking_uri
+        self.run_id = run_id
+        self.experiment_name = experiment_name
+        self.export_deleted_runs = export_deleted_runs
+        self.log_level = log_level
+
+        LOGGER.setLevel(self.log_level)
+
+    @staticmethod
+    def update_run_id(params: dict, key: str, new_run_id: str, offline_run_id: str) -> dict:
+        params[f"config.training.{key}"] = new_run_id
+        params[f"config.training.offline_{key}"] = offline_run_id
+
+        if key == "run_id":
+            params[f"metadata.offline_{key}"] = offline_run_id
+            params[f"metadata.{key}"] = new_run_id
+        return params
+
+    def update_parent_run_info(self, tags: dict, tag_key: str, tag_dest: str, dst_run_id: str) -> dict:
+        mlflow.set_tracking_uri(self.dest_tracking_uri)
+
+        # Check if there is already a parent run in the destination tracking uri
+        runs = mlflow.search_runs(
+            experiment_ids=mlflow.get_experiment_by_name(self.experiment_name).experiment_id,
+            filter_string=f"params.metadata.offline_run_id = '{tags[tag_key]}'",
+        )
+
+        if not runs.empty:
+            if (runs.shape[0] > 1) and ("tags.resumedRun" in runs.columns):
+                new_parent_run_id = runs[runs["tags.resumedRun"] != "True"].run_id.iloc[0]
+            else:
+                new_parent_run_id = runs.iloc[0].run_id  # read the parent run online run_id
+        else:
+            new_parent_run_id = dst_run_id
+        tags[tag_dest] = tags[tag_key]  # keep offline parent run_id
+        tags[tag_key] = new_parent_run_id  # update new online parent run_id
+        return tags
+
+    def check_run_is_logged(self, status: str = "FINISHED") -> bool:
+        """Blocks sync if top-level parent run or single runs are unavailable."""
+        run_logged = False
+        if status == "FINISHED":
+            mlflow.set_tracking_uri(self.dest_tracking_uri)
+            synced_runs = mlflow.search_runs(
+                experiment_ids=mlflow.get_experiment_by_name(self.experiment_name).experiment_id,
+                filter_string=f"params.metadata.offline_run_id = '{self.run_id}'",
+            )
+            if not synced_runs.empty:  # single run (no child) already logged
+                run_logged = True
+        return run_logged
+
+    def sync(
+        self,
+    ) -> None:
+        """Sync an offline run to the destination tracking uri."""
+        src_mlflow_client = mlflow.MlflowClient(self.source_tracking_uri)
+        dest_mlflow_client = mlflow.MlflowClient(self.dest_tracking_uri)
+        http_client = create_http_client(dest_mlflow_client)
+        # GET SOURCE RUN ##
+        run = src_mlflow_client.get_run(self.run_id)
+        run_logged = self.check_run_is_logged(status=run.info.status)
+        if run_logged:
+            LOGGER.info("Run already imported %s into experiment %s", self.run_id, self.experiment_name)
+            return
+
+        if run.info.lifecycle_stage == "deleted" and not self.export_deleted_runs:
+            LOGGER.warning(
+                "Not exporting run %s because its lifecycle_stage is  %s",
+                run.info.run_id,
+                run.info.lifecycle_stage,
+            )
+
+            return
+        msg = {
+            "run_id": run.info.run_id,
+            "lifecycle_stage": run.info.lifecycle_stage,
+            "experiment_id": run.info.experiment_id,
+        }
+        LOGGER.info("Exporting run: %s", msg)
+
+        run_info = mlflow_utils.strip_underscores(run.info)
+        src_user_id = run_info["user_id"]
+
+        exp = dest_mlflow_client.get_experiment_by_name(self.experiment_name)
+        dst_run = dest_mlflow_client.create_run(exp.experiment_id)
+        dst_run_id = dst_run.info.run_id
+
+        tags = dict(sorted(run.data.tags.items()))
+
+        params = run.data.params
+        # So far there is no easy way to force mlflow to use a specific run_id, that means
+        # that when we online sync the offline runs those will have run run_ids. To keep
+        # track of online and offline governance in that case we update run_ids info
+
+        if (params["config.training.fork_run_id"] == "None") and (params["metadata.run_id"] == run.info.run_id):
+            params = self.update_run_id(params, "run_id", new_run_id=dst_run_id, offline_run_id=run.info.run_id)
+
+        elif "forkedRun" in tags:
+            try:
+                tags = self.update_parent_run_info(
+                    tags=tags,
+                    tag_key="forkedRunId",
+                    tag_dest="offline.forkedRunId",
+                    dst_run_id=dst_run_id,
+                )
+                params = self.update_run_id(
+                    params,
+                    "fork_run_id",
+                    new_run_id=tags["forkedRunId"],
+                    offline_run_id=tags["offline.forkedRunId"],
+                )
+                params = self.update_run_id(params, "run_id", new_run_id=dst_run_id, offline_run_id=run.info.run_id)
+
+            except AttributeError:
+                LOGGER.warning("No forked run parent found")
+
+        elif "resumedRun" in tags:
+            try:
+                tags = self.update_parent_run_info(
+                    tags=tags,
+                    tag_key="mlflow.parentRunId",
+                    tag_dest="mlflow.offline.parentRunId",
+                    dst_run_id=dst_run_id,
+                )
+                params = self.update_run_id(
+                    params,
+                    "run_id",
+                    new_run_id=tags["mlflow.parentRunId"],
+                    offline_run_id=tags["mlflow.offline.parentRunId"],
+                )
+
+                params["config.training.offline_run_id_folder"] = run.info.run_id
+
+            except AttributeError:
+                LOGGER.warning("No parent run found")
+
+        tags["offlineRun"] = "True"
+
+        src_run_dct = {
+            "params": run.data.params,
+            "metrics": _get_metrics_with_steps(src_mlflow_client, run),
+            "tags": tags,
+            "inputs": _inputs_to_dict(run.inputs),
+        }
+
+        try:
+            import_run_data(
+                dest_mlflow_client,
+                src_run_dct,
+                dst_run_id,
+                src_user_id,
+            )
+            _import_inputs(http_client, src_run_dct, dst_run_id)
+
+            path = Path(self.source_tracking_uri, run.info.experiment_id, self.run_id, "artifacts")
+            if path.exists():
+                mlflow.set_tracking_uri(self.dest_tracking_uri)
+                dest_mlflow_client.log_artifacts(dst_run_id, path)
+            dest_mlflow_client.set_terminated(dst_run_id, RunStatus.to_string(RunStatus.FINISHED))
+
+        except Exception as e:
+            dest_mlflow_client.set_terminated(dst_run_id, RunStatus.to_string(RunStatus.FAILED))
+            import traceback
+
+            traceback.print_exc()
+            raise Exception(e, "Importing run %s of experiment %s failed", dst_run_id, exp.name) from e  # noqa: TRY002
+
+        LOGGER.info("Imported run %s into experiment %s", dst_run_id, self.experiment_name)
+
+        temp.close()
diff --git a/src/anemoi/training/utils/seeding.py b/src/anemoi/training/utils/seeding.py
new file mode 100644
index 00000000..d766bd1a
--- /dev/null
+++ b/src/anemoi/training/utils/seeding.py
@@ -0,0 +1,47 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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 __future__ import annotations
+
+import os
+
+
+def get_base_seed(base_seed_env: str | None = None) -> int:
+    """Gets the base seed from the environment variables.
+
+    Option to manually set a seed via export ANEMOI_BASE_SEED=xxx in job script
+
+    Parameters
+    ----------
+    base_seed_env : str, optional
+        Environment variable to use for the base seed, by default None
+
+    Returns
+    -------
+    int
+        Base seed.
+
+    """
+    env_var_list = ["ANEMOI_BASE_SEED", "SLURM_JOB_ID"]
+    if base_seed_env is not None:
+        env_var_list = [base_seed_env, *env_var_list]
+
+    base_seed = None
+    for env_var in env_var_list:
+        if env_var in os.environ:
+            base_seed = int(os.environ.get(env_var))
+            break
+
+    assert base_seed is not None, f"Base seed not found in environment variables {env_var_list}"
+
+    base_seed_threshold = 1000
+    if base_seed < base_seed_threshold:
+        base_seed *= base_seed_threshold  # make it (hopefully) big enough
+
+    return base_seed
diff --git a/src/anemoi/training/utils/usable_indices.py b/src/anemoi/training/utils/usable_indices.py
new file mode 100644
index 00000000..7bdd5cbd
--- /dev/null
+++ b/src/anemoi/training/utils/usable_indices.py
@@ -0,0 +1,56 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 __future__ import annotations
+
+import numpy as np
+
+
+def get_usable_indices(
+    missing_indices: set[int] | None,
+    series_length: int,
+    rollout: int,
+    multistep: int,
+    timeincrement: int = 1,
+) -> np.ndarray:
+    """Get the usable indices of a series whit missing indices.
+
+    Parameters
+    ----------
+    missing_indices : set[int]
+        Dataset to be used.
+    series_length : int
+        Length of the series.
+    rollout : int
+        Number of steps to roll out.
+    multistep : int
+        Number of previous indices to include as predictors.
+    timeincrement : int
+        Time increment, by default 1.
+
+    Returns
+    -------
+    usable_indices : np.array
+        Array of usable indices.
+    """
+    prev_invalid_dates = (multistep - 1) * timeincrement
+    next_invalid_dates = rollout * timeincrement
+
+    usable_indices = np.arange(series_length)  # set of all indices
+
+    if missing_indices is None:
+        missing_indices = set()
+
+    missing_indices |= {-1, series_length}  # to filter initial and final indices
+
+    # Missing indices
+    for i in missing_indices:
+        usable_indices = usable_indices[
+            (usable_indices < i - next_invalid_dates) + (usable_indices > i + prev_invalid_dates)
+        ]
+
+    return usable_indices
diff --git a/src/hydra_plugins/anemoi_searchpath/__init__.py b/src/hydra_plugins/anemoi_searchpath/__init__.py
new file mode 100644
index 00000000..33d7fa0a
--- /dev/null
+++ b/src/hydra_plugins/anemoi_searchpath/__init__.py
@@ -0,0 +1,8 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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.
+#
diff --git a/src/hydra_plugins/anemoi_searchpath/anemoi_searchpath_plugin.py b/src/hydra_plugins/anemoi_searchpath/anemoi_searchpath_plugin.py
new file mode 100644
index 00000000..75efae9a
--- /dev/null
+++ b/src/hydra_plugins/anemoi_searchpath/anemoi_searchpath_plugin.py
@@ -0,0 +1,69 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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.
+
+import logging
+import os
+from pathlib import Path
+
+from hydra.core.config_search_path import ConfigSearchPath
+from hydra.plugins.search_path_plugin import SearchPathPlugin
+
+LOGGER = logging.getLogger(__name__)
+
+
+class AnemoiSearchPathPlugin(SearchPathPlugin):
+    """Prepend the Anemoi home directory to the hydra searchpath."""
+
+    def manipulate_search_path(self, search_path: ConfigSearchPath) -> None:
+        """Prepend the Anemoi directories to the hydra searchpath.
+
+        This builds the hierarchy of the search path by prepending the Anemoi
+        directories to the search path. The hierarchy is as follows in decreasing priority:
+        - Current working directory
+        - Anemoi Config Env directory
+        - Anemoi Home directory
+
+        Parameters
+        ----------
+        search_path : ConfigSearchPath
+            Hydra ConfigSearchPath object.
+
+        """
+        for suffix in ("", "config"):
+            anemoi_home_path = Path(Path.home(), ".config", "anemoi", "training", suffix)
+            if anemoi_home_path.exists() and not Path(anemoi_home_path, "config").exists():
+                search_path.prepend(
+                    provider="anemoi-home-searchpath-plugin",
+                    path=str(anemoi_home_path),
+                )
+                LOGGER.info("Prepending Anemoi Home (%s) to the search path.", anemoi_home_path)
+                LOGGER.debug("Search path is now: %s", search_path)
+
+        for suffix in ("", "config"):
+            env_anemoi_config_path = os.getenv("ANEMOI_CONFIG_PATH")
+            if env_anemoi_config_path is None:
+                continue
+            anemoi_config_path = Path(env_anemoi_config_path, suffix)
+            if anemoi_config_path.exists() and not Path(anemoi_config_path, "config").exists():
+                search_path.prepend(
+                    provider="anemoi-env-searchpath-plugin",
+                    path=str(anemoi_config_path),
+                )
+                LOGGER.info("Prepending Anemoi Config Env (%s) to the search path.", anemoi_config_path)
+                LOGGER.debug("Search path is now: %s", search_path)
+
+        for suffix in ("", "config"):
+            cwd_path = Path.cwd() / suffix
+            if cwd_path.exists() and not Path(cwd_path, "config").exists():
+                search_path.prepend(
+                    provider="anemoi-cwd-searchpath-plugin",
+                    path=str(cwd_path),
+                )
+                LOGGER.info("Prepending current user directory (%s) to the search path.", cwd_path)
+                LOGGER.debug("Search path is now: %s", search_path)
+        LOGGER.info("Search path is now: %s", search_path)


Test patch:
diff --git a/docs/dev/testing.rst b/docs/dev/testing.rst
new file mode 100644
index 00000000..47559615
--- /dev/null
+++ b/docs/dev/testing.rst
@@ -0,0 +1,190 @@
+#########
+ Testing
+#########
+
+Comprehensive testing is crucial for maintaining the reliability and
+stability of Anemoi Training. This guide outlines our testing strategy
+and best practices for contributing tests.
+
+*******************
+ Testing Framework
+*******************
+
+We use pytest as our primary testing framework. Pytest offers a simple
+and powerful way to write and run tests.
+
+***************
+ Writing Tests
+***************
+
+General Guidelines
+==================
+
+#. Write tests for all new features and bug fixes.
+#. Aim for high test coverage, especially for critical components.
+#. Keep tests simple, focused, and independent of each other.
+#. Use descriptive names for test functions, following the pattern
+   `test_<functionality>_<scenario>`.
+
+Example Test Structure
+======================
+
+.. code:: python
+
+   import pytest
+   from anemoi.training import SomeFeature
+
+
+   def test_some_feature_normal_input():
+       feature = SomeFeature()
+       result = feature.process(normal_input)
+       assert result == expected_output
+
+
+   def test_some_feature_edge_case():
+       feature = SomeFeature()
+       with pytest.raises(ValueError):
+           feature.process(invalid_input)
+
+****************
+ Types of Tests
+****************
+
+1. Unit Tests
+=============
+
+Test individual components in isolation. These should be the majority of
+your tests.
+
+2. Integration Tests
+====================
+
+Test how different components work together. These are particularly
+important for data processing pipelines and model training workflows.
+
+3. Functional Tests
+===================
+
+Test entire features or workflows from start to finish. These ensure
+that the system works as expected from a user's perspective.
+
+4. Parametrized Tests
+=====================
+
+Use pytest's parametrize decorator to run the same test with different
+inputs:
+
+.. code:: python
+
+   @pytest.mark.parametrize(
+       "input,expected",
+       [
+           (2, 4),
+           (3, 9),
+           (4, 16),
+       ],
+   )
+   def test_square(input, expected):
+       assert square(input) == expected
+
+You can also consider ``hypothesis`` for property-based testing.
+
+5. Fixtures
+===========
+
+Use fixtures to set up common test data or objects:
+
+.. code:: python
+
+   @pytest.fixture
+   def sample_dataset():
+       # Create and return a sample dataset
+       pass
+
+
+   def test_data_loading(sample_dataset):
+       # Use the sample_dataset fixture in your test
+       pass
+
+***************
+ Running Tests
+***************
+
+To run all tests:
+
+.. code:: bash
+
+   pytest
+
+To run tests in a specific file:
+
+.. code:: bash
+
+   pytest tests/test_specific_feature.py
+
+To run tests with a specific mark:
+
+.. code:: bash
+
+   pytest -m slow
+
+***************
+ Test Coverage
+***************
+
+We use pytest-cov to measure test coverage. To run tests with coverage:
+
+.. code:: bash
+
+   pytest --cov=anemoi_training
+
+Aim for at least 80% coverage for new features, and strive to maintain
+or improve overall project coverage.
+
+************************
+ Continuous Integration
+************************
+
+All tests are run automatically on our CI/CD pipeline for every pull
+request. Ensure all tests pass before submitting your PR.
+
+*********************
+ Performance Testing
+*********************
+
+For performance-critical components:
+
+#. Write benchmarks.
+#. Compare performance before and after changes.
+#. Set up performance regression tests in CI.
+
+**********************
+ Mocking and Patching
+**********************
+
+Use unittest.mock or pytest-mock for mocking external dependencies or
+complex objects:
+
+.. code:: python
+
+   def test_api_call(mocker):
+       mock_response = mocker.Mock()
+       mock_response.json.return_value = {"data": "mocked"}
+       mocker.patch("requests.get", return_value=mock_response)
+
+       result = my_api_function()
+       assert result == "mocked"
+
+****************
+ Best Practices
+****************
+
+#. Keep tests fast: Optimize slow tests or mark them for separate
+   execution.
+#. Use appropriate assertions: pytest provides a rich set of assertions.
+#. Test edge cases and error conditions, not just the happy path.
+#. Regularly review and update tests as the codebase evolves.
+#. Document complex test setups or scenarios.
+
+By following these guidelines and continuously improving our test suite,
+we can ensure the reliability and maintainability of Anemoi Training.
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 00000000..f77701dd
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,7 @@
+[pytest]
+markers =
+    data_dependent: marks tests depending on data (deselect with '-m "not data_dependent"')
+    auth: marks tests that require authentication (deselect with '-m "not auth"')
+    gpu: marks tests that require a GPU (deselect with '-m "not gpu"')
+
+tmp_path_retention_policy = none
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 00000000..711b1e2d
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,32 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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.
+
+
+import pytest
+from _pytest.fixtures import SubRequest
+from hydra import compose
+from hydra import initialize
+from omegaconf import DictConfig
+
+from anemoi.training.data.datamodule import AnemoiDatasetsDataModule
+
+
+@pytest.fixture
+def config(request: SubRequest) -> DictConfig:
+    overrides = request.param
+    with initialize(version_base=None, config_path="../src/anemoi/training/config"):
+        # config is relative to a module
+        return compose(config_name="debug", overrides=overrides)
+
+
+@pytest.fixture
+def datamodule() -> AnemoiDatasetsDataModule:
+    with initialize(version_base=None, config_path="../src/anemoi/training/config"):
+        # config is relative to a module
+        cfg = compose(config_name="config")
+    return AnemoiDatasetsDataModule(cfg)
diff --git a/tests/diagnostics/mlflow/test_auth.py b/tests/diagnostics/mlflow/test_auth.py
new file mode 100644
index 00000000..9f3a709b
--- /dev/null
+++ b/tests/diagnostics/mlflow/test_auth.py
@@ -0,0 +1,159 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 __future__ import annotations
+
+import os
+import time
+
+import pytest
+
+from anemoi.training.diagnostics.mlflow.auth import TokenAuth
+
+
+def mocks(
+    mocker: pytest.MockerFixture,
+    token_request: dict | None = None,
+    load_config: dict | None = None,
+) -> pytest.Mock:
+    if load_config is None:
+        load_config = {}
+    if token_request is None:
+        token_request = {}
+    response = {
+        "access_token": "access_token",
+        "expires_in": 3600,
+        "refresh_token": "new_refresh_token",
+    }
+    response.update(token_request)
+
+    config = {
+        "refresh_token": "old_refresh_token",
+        "refresh_expires": time.time() + 3600,
+    }
+    config.update(load_config)
+
+    mock_token_request = mocker.patch(
+        "anemoi.training.diagnostics.mlflow.auth.TokenAuth._token_request",
+        return_value=response,
+    )
+    mocker.patch(
+        "anemoi.training.diagnostics.mlflow.auth.load_config",
+        return_value=config,
+    )
+    mocker.patch(
+        "anemoi.training.diagnostics.mlflow.auth.save_config",
+    )
+    mocker.patch(
+        "anemoi.training.diagnostics.mlflow.auth.getpass",
+        return_value="seed_refresh_token",
+    )
+    mocker.patch("os.environ")
+
+    return mock_token_request
+
+
+def test_auth(mocker: pytest.MockerFixture) -> None:
+    mock_token_request = mocks(mocker)
+
+    auth = TokenAuth("https://test.url")
+
+    assert auth.access_token is None
+    assert auth.refresh_token == "old_refresh_token"  # noqa: S105
+
+    auth.authenticate()
+    # test that no new token is requested the second time
+    auth.authenticate()
+
+    mock_token_request.assert_called_once()
+
+    assert auth.access_token == "access_token"  # noqa: S105
+    assert auth.access_expires > time.time()
+    assert auth.refresh_token == "new_refresh_token"  # noqa: S105
+
+
+def test_not_logged_in(mocker: pytest.MockerFixture) -> None:
+    # no refresh token
+    mocks(mocker, load_config={"refresh_token": None})
+    auth = TokenAuth("https://test.url")
+    pytest.raises(RuntimeError, auth.authenticate)
+
+    # expired refresh token
+    mocks(mocker, load_config={"refresh_expires": time.time() - 1})
+    auth = TokenAuth("https://test.url")
+    pytest.raises(RuntimeError, auth.authenticate)
+
+
+def test_login(mocker: pytest.MockerFixture) -> None:
+    # normal login
+    mock_token_request = mocks(mocker)
+    auth = TokenAuth("https://test.url")
+    auth.login()
+
+    mock_token_request.assert_called_once()
+
+    # normal credential login
+    mock_token_request = mocks(mocker, load_config={"refresh_token": None})
+    auth = TokenAuth("https://test.url")
+    auth.login()
+
+    mock_token_request.assert_called_once()
+
+    # forced credential login
+    mock_token_request = mocks(mocker)
+    auth = TokenAuth("https://test.url")
+    auth.login(force_credentials=True)
+
+    mock_token_request.assert_called_once()
+
+    # failed login
+    mock_token_request = mocks(mocker, token_request={"refresh_token": None})
+    auth = TokenAuth("https://test.url")
+    pytest.raises(RuntimeError, auth.login)
+
+    assert mock_token_request.call_count == 2
+
+
+def test_enabled(mocker: pytest.MockerFixture) -> None:
+    mock_token_request = mocks(mocker)
+    auth = TokenAuth("https://test.url", enabled=False)
+    auth.authenticate()
+
+    mock_token_request.assert_not_called()
+
+
+def test_api(mocker: pytest.MockerFixture) -> None:
+    mocks(mocker)
+    auth = TokenAuth("https://test.url")
+    mock_post = mocker.patch("requests.post")
+
+    # successful request
+    response_json = {"status": "OK", "response": {}}
+    mocker.patch("requests.post.return_value.json", return_value=response_json)
+    response = auth._request("path", {"key": "value"})
+
+    assert response == response_json["response"]
+    mock_post.assert_called_once_with(
+        "https://test.url/path",
+        json={"key": "value"},
+        headers=mocker.ANY,
+        timeout=mocker.ANY,
+    )
+
+    # api error
+    error_response = {"status": "ERROR", "response": {}}
+    mocker.patch("requests.post.return_value.json", return_value=error_response)
+
+    with pytest.raises(RuntimeError):
+        auth._request("path", {"key": "value"})
+
+
+def test_target_env_var(mocker: pytest.MockerFixture) -> None:
+    mocks(mocker)
+    auth = TokenAuth("https://test.url", target_env_var="MLFLOW_TEST_ENV_VAR")
+    auth.authenticate()
+
+    os.environ.__setitem__.assert_called_once_with("MLFLOW_TEST_ENV_VAR", "access_token")
diff --git a/tests/diagnostics/mlflow/test_client.py b/tests/diagnostics/mlflow/test_client.py
new file mode 100644
index 00000000..93c0aeeb
--- /dev/null
+++ b/tests/diagnostics/mlflow/test_client.py
@@ -0,0 +1,38 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+if TYPE_CHECKING:
+    import pytest_mock
+
+from anemoi.training.diagnostics.mlflow.client import AnemoiMlflowClient
+
+
+@pytest.fixture(autouse=True)
+def mocks(mocker: pytest_mock.MockerFixture) -> None:
+    mocker.patch("anemoi.training.diagnostics.mlflow.client.TokenAuth")
+    mocker.patch("anemoi.training.diagnostics.mlflow.client.health_check")
+    mocker.patch("anemoi.training.diagnostics.mlflow.client.AnemoiMlflowClient.search_experiments")
+
+
+def test_auth_injected() -> None:
+    client = AnemoiMlflowClient("http://localhost:5000", authentication=True, check_health=False)
+    client.search_experiments()
+    client.search_experiments()
+
+    assert client.anemoi_auth.authenticate.call_count == 2
+
+
+def test_health_check() -> None:
+    # the internal health check will trigger an authenticate call
+    client = AnemoiMlflowClient("http://localhost:5000", authentication=True, check_health=True)
+
+    client.anemoi_auth.authenticate.assert_called_once()
diff --git a/tests/diagnostics/test_checkpoint.py b/tests/diagnostics/test_checkpoint.py
new file mode 100644
index 00000000..e2dce376
--- /dev/null
+++ b/tests/diagnostics/test_checkpoint.py
@@ -0,0 +1,139 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 __future__ import annotations
+
+import datetime
+import shutil
+from pathlib import Path
+
+import pytest
+import torch
+from anemoi.utils.checkpoints import load_metadata
+from anemoi.utils.config import DotDict
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
+from pytorch_lightning import Trainer
+from pytorch_lightning.demos.boring_classes import BoringModel
+from torch import nn
+
+from anemoi.training.diagnostics.callbacks import AnemoiCheckpoint
+from anemoi.training.utils.jsonify import map_config_to_primitives
+
+
+class DummyModel(torch.nn.Module):
+    """Dummy pytorch model for testing."""
+
+    def __init__(self, *, config: DotDict, metadata: dict):
+        super().__init__()
+
+        self.config = config
+        self.metadata = metadata
+        self.fc1 = nn.Linear(32, 5)
+        self.fc2 = nn.Linear(5, 1)
+        self.relu = nn.ReLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.relu(x)
+        return self.fc2(x)
+
+
+class DummyModule(BoringModel):
+    """Dummy lightning module for testing."""
+
+    def __init__(self, *, config: DotDict, metadata: dict) -> None:
+        super().__init__()
+        self.model = DummyModel(config=config, metadata=metadata)
+        self.save_hyperparameters()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+
+
+@pytest.fixture
+def tmp_path(tmp_path_factory: pytest.TempPathFactory) -> str:
+    # returns a session-scoped temporary directory
+    return str(tmp_path_factory.mktemp("test_checkpoint"))
+
+
+@pytest.fixture
+def config() -> DictConfig:
+    config = DictConfig({"diagnostics": {"log": {"wandb": {"enabled": False}, "mlflow": {"enabled": False}}}})
+    return DotDict(map_config_to_primitives(OmegaConf.to_container(config, resolve=True)))
+
+
+@pytest.fixture
+def checkpoint_settings(tmp_path: str) -> dict[str : str | bool]:
+    return {
+        "dirpath": tmp_path,
+        "verbose": False,
+        "save_weights_only": False,
+        "auto_insert_metric_name": False,
+        "save_on_train_epoch_end": False,
+        "enable_version_counter": False,
+    }
+
+
+@pytest.fixture
+def callback(tmp_path: str, config: DictConfig, checkpoint_settings: dict) -> AnemoiCheckpoint:
+    callback = AnemoiCheckpoint(
+        config=config,
+        filename="{step}",
+        save_last=True,
+        train_time_interval=datetime.timedelta(seconds=0.1),
+        save_top_k=3,
+        monitor="step",
+        mode="max",
+        **checkpoint_settings,
+    )
+    callback.dirpath = tmp_path
+    return callback
+
+
+@pytest.fixture
+def metadata(config: DictConfig) -> dict:
+    return map_config_to_primitives(
+        {
+            "config": config,
+        },
+    )
+
+
+@pytest.fixture
+def model(metadata: dict, config: DictConfig) -> DummyModule:
+    kwargs = {
+        "metadata": metadata,
+        "config": config,
+    }
+    return DummyModule(**kwargs)
+
+
+@pytest.mark.data_dependent
+def test_same_uuid(tmp_path: str, callback: AnemoiCheckpoint, model: DummyModule) -> None:
+    """Test if the inference checkpoints and lightning checkpoints store same uuid.
+
+    Args:
+        tmp_path (str): path to checkpoint dir
+        callback (AnemoiCheckpoint): callback to store checkpoints
+        model (DummyModule): dummy lightning module
+    """
+    trainer = Trainer(default_root_dir=tmp_path, accelerator="cpu", callbacks=[callback], max_epochs=2, logger=False)
+    trainer.fit(model)
+
+    for ckpt_path in Path(tmp_path).iterdir():
+        if str(ckpt_path.name).startswith("inference"):
+            pl_ckpt_name = (ckpt_path.name).replace("inference-", "")
+
+            if Path(tmp_path + "/" + pl_ckpt_name).exists():
+                uuid = load_metadata(ckpt_path)["uuid"]
+
+                pl_model = DummyModule.load_from_checkpoint(tmp_path + "/" + pl_ckpt_name)
+
+                assert uuid == pl_model.hparams["metadata"]["uuid"]
+
+    shutil.rmtree(tmp_path)
diff --git a/tests/hydra/test_search_path_plugins.py b/tests/hydra/test_search_path_plugins.py
new file mode 100644
index 00000000..dd981b66
--- /dev/null
+++ b/tests/hydra/test_search_path_plugins.py
@@ -0,0 +1,25 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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 hydra import initialize
+from hydra.core.global_hydra import GlobalHydra
+from hydra.core.plugins import Plugins
+from hydra.plugins.search_path_plugin import SearchPathPlugin
+
+from hydra_plugins.anemoi_searchpath.anemoi_searchpath_plugin import AnemoiSearchPathPlugin
+
+
+def test_anemoi_home_searchpath_discovery() -> None:
+    # Tests that this plugin can be discovered via the plugins subsystem when looking at all Plugins
+    assert AnemoiSearchPathPlugin.__name__ in [x.__name__ for x in Plugins.instance().discover(SearchPathPlugin)]
+
+
+def test_config_installed() -> None:
+    with initialize(version_base=None):
+        config_loader = GlobalHydra.instance().config_loader()
+        assert "default" in config_loader.get_group_options("hydra/output")
diff --git a/tests/train/test_loss_scaling.py b/tests/train/test_loss_scaling.py
new file mode 100644
index 00000000..84ca2189
--- /dev/null
+++ b/tests/train/test_loss_scaling.py
@@ -0,0 +1,162 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+import pytest
+import torch
+from _pytest.fixtures import SubRequest
+from anemoi.models.data_indices.collection import IndexCollection
+from omegaconf import DictConfig
+
+from anemoi.training.train.forecaster import GraphForecaster
+
+
+@pytest.fixture
+def fake_data(request: SubRequest) -> tuple[DictConfig, IndexCollection]:
+    config = DictConfig(
+        {
+            "data": {
+                "forcing": ["x"],
+                "diagnostic": ["z", "q"],
+                "remapped": {
+                    "d": ["cos_d", "sin_d"],
+                },
+            },
+            "training": {
+                "loss_scaling": {
+                    "default": 1,
+                    "sfc": {
+                        "z": 0.1,
+                        "other": 100,
+                    },
+                    "pl": {"y": 0.5},
+                },
+                "metrics": ["other", "y_850"],
+                "pressure_level_scaler": request.param,
+            },
+        },
+    )
+    name_to_index = {"x": 0, "y_50": 1, "y_500": 2, "y_850": 3, "z": 5, "q": 4, "other": 6, "d": 7}
+    data_indices = IndexCollection(config=config, name_to_index=name_to_index)
+    return config, data_indices
+
+
+linear_scaler = {
+    "_target_": "anemoi.training.data.scaling.LinearPressureLevelScaler",
+    "minimum": 0.0,
+    "slope": 0.001,
+}
+relu_scaler = {
+    "_target_": "anemoi.training.data.scaling.ReluPressureLevelScaler",
+    "minimum": 0.2,
+    "slope": 0.001,
+}
+constant_scaler = {
+    "_target_": "anemoi.training.data.scaling.NoPressureLevelScaler",
+    "minimum": 1.0,
+    "slope": 0.0,
+}
+polynomial_scaler = {
+    "_target_": "anemoi.training.data.scaling.PolynomialPressureLevelScaler",
+    "minimum": 0.2,
+    "slope": 0.001,
+}
+
+expected_linear_scaling = torch.Tensor(
+    [
+        50 / 1000 * 0.5,  # y_50
+        500 / 1000 * 0.5,  # y_500
+        850 / 1000 * 0.5,  # y_850
+        1,  # q
+        0.1,  # z
+        100,  # other
+        1,  # cos_d
+        1,  # sin_d
+    ],
+)
+expected_relu_scaling = torch.Tensor(
+    [
+        0.2 * 0.5,  # y_50
+        500 / 1000 * 0.5,  # y_500
+        850 / 1000 * 0.5,  # y_850
+        1,  # q
+        0.1,  # z
+        100,  # other
+        1,  # cos_d
+        1,  # sin_d
+    ],
+)
+expected_constant_scaling = torch.Tensor(
+    [
+        1 * 0.5,  # y_50
+        1 * 0.5,  # y_500
+        1 * 0.5,  # y_850
+        1,  # q
+        0.1,  # z
+        100,  # other
+        1,  # cos_d
+        1,  # sin_d
+    ],
+)
+expected_polynomial_scaling = torch.Tensor(
+    [
+        ((50 / 1000) ** 2 + 0.2) * 0.5,  # y_50
+        ((500 / 1000) ** 2 + 0.2) * 0.5,  # y_500
+        ((850 / 1000) ** 2 + 0.2) * 0.5,  # y_850
+        1,  # q
+        0.1,  # z
+        100,  # other
+        1,  # cos_d
+        1,  # sin_d
+    ],
+)
+
+
+@pytest.mark.parametrize(
+    ("fake_data", "expected_scaling"),
+    [
+        (linear_scaler, expected_linear_scaling),
+        (relu_scaler, expected_relu_scaling),
+        (constant_scaler, expected_constant_scaling),
+        (polynomial_scaler, expected_polynomial_scaling),
+    ],
+    indirect=["fake_data"],
+)
+def test_loss_scaling_vals(fake_data: tuple[DictConfig, IndexCollection], expected_scaling: torch.Tensor) -> None:
+    config, data_indices = fake_data
+
+    _, _, loss_scaling = GraphForecaster.metrics_loss_scaling(config, data_indices)
+
+    assert torch.allclose(loss_scaling, expected_scaling)
+
+
+@pytest.mark.parametrize("fake_data", [linear_scaler], indirect=["fake_data"])
+def test_metric_range(fake_data: tuple[DictConfig, IndexCollection]) -> None:
+    config, data_indices = fake_data
+
+    metric_range, metric_ranges_validation, _ = GraphForecaster.metrics_loss_scaling(config, data_indices)
+
+    expected_metric_range_validation = {
+        "pl_y": [
+            data_indices.model.output.name_to_index["y_50"],
+            data_indices.model.output.name_to_index["y_500"],
+            data_indices.model.output.name_to_index["y_850"],
+        ],
+        "sfc_other": [data_indices.model.output.name_to_index["other"]],
+        "sfc_q": [data_indices.model.output.name_to_index["q"]],
+        "sfc_z": [data_indices.model.output.name_to_index["z"]],
+        "other": [data_indices.model.output.name_to_index["other"]],
+        "sfc_d": [data_indices.model.output.name_to_index["d"]],
+        "y_850": [data_indices.model.output.name_to_index["y_850"]],
+    }
+
+    expected_metric_range = expected_metric_range_validation.copy()
+    del expected_metric_range["sfc_d"]
+    expected_metric_range["sfc_cos_d"] = [data_indices.internal_model.output.name_to_index["cos_d"]]
+    expected_metric_range["sfc_sin_d"] = [data_indices.internal_model.output.name_to_index["sin_d"]]
+
+    assert dict(metric_ranges_validation) == expected_metric_range_validation
+    assert dict(metric_range) == expected_metric_range
diff --git a/tests/utils/test_usable_indices.py b/tests/utils/test_usable_indices.py
new file mode 100644
index 00000000..6bc5c83f
--- /dev/null
+++ b/tests/utils/test_usable_indices.py
@@ -0,0 +1,45 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+import numpy as np
+
+from anemoi.training.utils.usable_indices import get_usable_indices
+
+
+def test_get_usable_indices() -> None:
+    """Test get_usable_indices function."""
+    # Test base case
+    valid_indices = get_usable_indices(missing_indices=None, series_length=10, rollout=1, multistep=1, timeincrement=1)
+    expected_values = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8])
+    assert np.allclose(valid_indices, expected_values)
+
+    # Test multiple steps inputs
+    valid_indices = get_usable_indices(missing_indices=None, series_length=10, rollout=1, multistep=2, timeincrement=1)
+    expected_values = np.array([1, 2, 3, 4, 5, 6, 7, 8])
+    assert np.allclose(valid_indices, expected_values)
+
+    # Test roll out
+    valid_indices = get_usable_indices(missing_indices=None, series_length=10, rollout=2, multistep=1, timeincrement=1)
+    expected_values = np.array([0, 1, 2, 3, 4, 5, 6, 7])
+    assert np.allclose(valid_indices, expected_values)
+
+    # Test longer time increments
+    valid_indices = get_usable_indices(missing_indices=None, series_length=10, rollout=1, multistep=2, timeincrement=2)
+    expected_values = np.array([2, 3, 4, 5, 6, 7])
+    assert np.allclose(valid_indices, expected_values)
+
+    # Test missing indices
+    missing_indices = {7, 5}
+    valid_indices = get_usable_indices(
+        missing_indices=missing_indices,
+        series_length=10,
+        rollout=1,
+        multistep=2,
+        timeincrement=1,
+    )
+    expected_values = np.array([1, 2, 3])
+    assert np.allclose(valid_indices, expected_values)