1 of 27

Transfer learning

Breakout session,

MLPP meeting – Offenbach

Teresa García (AEMET),

Jasper Wijnands (KNMI),

Håvard Haugen (MET Norway)

2 of 27

Royal Netherlands Meteorological Institute (KNMI)

2

April 1, 2025

Agenda

  • Development of base models (KNMI)
  • Transfer learning experiments using base models (MET Norway, AEMET)
  • Group discussion

3 of 27

Development of base stretched grid models for transfer learning using CERRA

Jasper Wijnands (KNMI),

with contributions from Bastien François (KNMI)

4 of 27

Royal Netherlands Meteorological Institute (KNMI)

4

April 1, 2025

Introduction

  • Set of common variables defined in MLPP-47 to align the process of transfer learning across organisations
    • Surface (sfc): 10u, 10v, 2d, 2t, msl, sp, skt, tp
    • Pressure levels (pl): q, t, u, v, z (no w)
    • Pressure levels used: 50, 100, 150, 200, 250, 300, 400, 500, 700, 850, 925, 1000 (no 600)
    • Forcing variables: lsm, insolation, z, cos_julian_day, cos_latitude, cos_local_time, cos_longitude, sin_julian_day, sin_latitude, sin_local_time, sin_longitude
  • Global n320 model for transfer learning developed by Mario Santa Cruz (ECMWF)
  • CERRA model presented here is an alternative starting point for transfer learning
  • Rationale:
    • 36 years of CERRA data available for training a stretched grid model, much more than typically available in regional datasets
    • Starting from a stretched grid model may better align with final task than starting from a global model

​

5 of 27

Royal Netherlands Meteorological Institute (KNMI)

5

April 1, 2025

Specifications CERRA-1024 base model

  • First version uses resolutions of o96 + 5.5km with hidden mesh resolution of 5 + 9
  • 1024 channels (i.e., 246 million trainable parameters)
  • Regional domain weighting of 25%
    • Further increasing regional weighting negatively impacts global performance
  • Normalisation based on statistics of global dataset to align with transfer learning task
    • For fine-tuning, also use global statistics
  • Rescaled learning rate of 0.002 (lr-AdamW), local LR of 1.25e-4
  • Trained from scratch for 260k iterations, 16 nodes with 4 A100 GPUs each, model sharding across 4 GPUs, effective batch size 16
  • Total training time: about two weeks

​

​

​

​

6 of 27

Royal Netherlands Meteorological Institute (KNMI)

6

April 1, 2025

MLflow training statistics

​

​

​

​

​

​

​

​

​

​

  • Checkpoint at 260k iterations used by MET Norway for transfer learning

​

​

​

​

​

7 of 27

Royal Netherlands Meteorological Institute (KNMI)

7

April 1, 2025

CERRA-1024 with rollout

  • Rollout training with 10 steps, 5k iterations
  • Training directly on maximum rollout (i.e., without increments)
  • Model could potentially replace stage A + B + C (Nipen et al., 2024)
  • Large simplification for Member States to train a stretched grid model, as only stage D would be required (with more iterations)
    • Directly use maximum rollout
  • Initial experiments ongoing together with AEMET

8 of 27

Royal Netherlands Meteorological Institute (KNMI)

8

April 1, 2025

Verification against CERRA

9 of 27

Royal Netherlands Meteorological Institute (KNMI)

9

April 1, 2025

Latest developments

  • More data, more compute, more trainable parameters 🡪 better performance ?
  • Combine many elements explored over the past year in a single model
    • Using 1B+ trainable parameters (num_channels: 2176), with just enough GPU memory left for rollout
    • Model sharding across 4 full nodes, as it does not affect training speed as badly as expected
    • Increase global resolution from o96 to n320 to align with transfer learning task (hidden mesh resolutions [7, 9])
    • Highest learning rate possible without divergence/NaN, to explore loss landscape with larger steps at first. Typically overtakes lower LR experiments towards the end of training
    • Effective batch size of 8. Experiments by ECMWF concluded batch size 16 is more efficient than 32. The intention here is to save some compute by scaling down further
  • Thanks to ECMWF for compute support

​

​

10 of 27

Preliminary results

11 of 27

Questions?

Royal Netherlands Meteorological Institute

11

April 1, 2025

12 of 27

Fine-tuning experiments

Håvard H. Haugen�

MET Norway: Olav Ersland • Lars Falk-Petersen • Thomas Nipen • Magnus Sikora Ingstad • Jørn Kristiansen • Máté Mile • Even Nordhagen • Aram Farhad Salihi • Ivar Seierstad • Roel Stappers • Paulina Tedesco

​

​

12

13 of 27

Fine-tuning experiments

Experiments:

  • Low res
    • Spatial resolution: o96 + 2.5km
    • Graph resolution: [5,9]
    • Scheduler: 5k iterations, bs 16, lr 6e-06
  • High res
    • Spatial resolution: n320 + 2.5km
    • Graph resolution: [7,10]
    • Scheduler: 5k iterations, bs 16, lr 6e-06

13

14 of 27

Headline scores

14

15 of 27

Headline scores

​

15

16 of 27

Example forecasts - precipitation

​

Cerra tl high res

Cloudy skies (r4)

Cerra tl low res

16

17 of 27

Example forecasts - wind speed

Cerra tl high res

Cloudy skies (r4)

Cerra tl low res

17

18 of 27

18

Multi-domain training*

CERRA

MetCoOp

DANRA

DOWA

ARA

CARRA East

  • Alternate between different domains across different batches
  • Allows for a higher learning rate for the high resolution part of the training
  • Stage A -> Multi-domain -> stage D (?)

*Ongoing work in Desination Earth by Sophie Buurman, Even Nordhagen, Aram Salihi, Caglar Kücük, Mathias Schreiner, Mario Santa Cruz and Thomas Nipen

19 of 27

Adventures of a newcomer to transfer learning

Teresa García (Aemet)

20 of 27

Motivation

  • We don’t have very long datasets available. Two 6h frequency datasets and one 3h frequency dataset from september 2021 to 2024 and resolution 2.5km, which have been uploaded to the catalogue:
    • aemet-an-harm-2p5km-2021-2024-6h-v0-iberia
    • aemet-an-harm-2p5km-2021-2024-6h-v0-canarias
    • aemet-an-harm-2p5km-2021-2024-3h-v0-canarias (tp 3h accumulated, iberia coming soon)
  • Other colleagues in the MLPP have observed that in that case better results are obtained pretraining with ERA5. Probably better if the starting point is CERRA?
  • Trying to avoid redoing steps A+B+C in (Nipen et al., 2024) and saving computational resources.

​

21 of 27

Preparing the config file

  • Started last week with the test, we don’t have results to share yet. But we can share some experiences preparing the experiments, which might be useful for other people trying to do transfer learning.
  • Configuration:
    • ERA5 as forcing datasets
    • 1024 channels
    • Resolution 10 in the local mesh (trained with resolution 9)
    • Resolution 5 in the global mesh
    • 12 KNN in the encoder
  • Use the same version of the model, go to previous commit:
    • git checkout eae93c62d0bcb3eef7c044664ecc898426491e3e

​

​

22 of 27

First problem encountered

  • In theory there are two ways to start transfer learning from the config file:
    • Forking / resuming run id
    • Warm start
  • Both fork run id and resume run id continue your own experiments given the mlflow id, but the starting point was not my experiment and was not in ATOS ☹
  • Warm start didn’t work, it has been adressed in a later branch ☹:
    • https://github.com/ecmwf/anemoi-core/tree/feat/support-ckpt-file

​

​

​

23 of 27

Solutions

  • First try using warm_start:
    • In the config file:

warm_start: KNMI_CERRA_o96_5.5km_5_9_ch1024-epoch_082-step_260000.ckpt

training:

fork_run_id: 'cerra_o96_1024'

    • Then you create a new folder named ‘cerra_o96_1024’ in your checkpoint path and copy the checkpoint there.
    • Change line 76 in anemoi-core/training/src/anemoi/training/utils/checkpoint.py to:
      • checkpoint = torch.load(ckpt_path, map_location=model.device, weights_only=False)
    • But then MLflow logging needs to be disabled (otherwise it assigns a new run id and looks for the check point there):

log:

mlflow:

enabled: False

    • Then there is no way of plotting the loss ☹

​

​

24 of 27

Solutions

  • Second try using fork run id:

​

    • Make a folder in your checkpoints path with the name of the original run id

fa1d632ee25f40478f6594d5b762bb4d.

    • Paste the checkpoint there as last.ckpt.
    • Fork the run id.
    • No need to disable mlflow ☺.

​

25 of 27

Second problem encountered

  • We have a limit of 4 nodes we can use. Tried to use model sharding and the model with 2.5km resolution didn’t fit in 4 nodes ☹
  • Not sure why, maybe because of the 10 resolution in the local area?
  • With a thinning of 2 it’s ok (5km). Attention: if you comment the thinning but upload the graph constructed with thinning it works anyway.
  • Food for thought: when going to higher resolutions, should we have a look at what people in LLMs are doing to reduce memory footprint in their models?

26 of 27

Third problem encountered

  • Killed this weekend by an mlflow duplicate key error:
    • 2025/03/28 08:19:00 WARNING mlflow.system_metrics.system_metrics_monitor: Failed to log system metrics: BAD_REQUEST: (raised as a result of Query-invoked autoflush; consider using a session.no_autoflush block if this flush is occurring prematurely)(psycopg2.errors.UniqueViolation) duplicate key value violates unique constraint "metric_pk"DETAIL: Key (key, "timestamp", step, run_uuid, value, is_nan)=(system/network_receive_megabytes, 1743149833782, 3475, 6b845c7d689f4dd29e3f23331f9b06d5, 1283.76378199994, f) already exists.
  • To be continued at a ‘Bring your problem session’?

​

​

​

27 of 27

Royal Netherlands Meteorological Institute (KNMI)

27

March 27, 2025

Group discussion

  • What can you change / not change between different stages?
  • What are some of the pitfalls?
  • Multi-domain training and potential impact on transfer learning

​