feat(training): add lora adapter to Single training method

training/pyproject.toml

@@ -46,6 +46,7 @@ dependencies = [
   "numba",
   "numpy",
   "nvidia-ml-py>=13.580.82",
+  "peft>=0.17.1",
   "pydantic>=2.9",
   "pytorch-lightning>=2.1",
   "timm>=0.9.2",

training/src/anemoi/training/config/lora.yaml

@@ -0,0 +1,21 @@
+defaults:
+- data: zarr
+- dataloader: native_grid
+- diagnostics: evaluation
+- system: example
+- graph: multi_scale
+- model: gnn
+- task: forecaster
+- training: lora
+- _self_
+
+
+# set to true to switch on config validation
+config_validation: True
+
+training:
+  lora_config:
+    r: 8
+    lora_alpha: 32
+    target_modules: ['mlp.0', 'mlp.2', 'dummy_layer']
+    modules_to_save: ['node_data_extractor.1']

training/src/anemoi/training/config/training/lora.yaml

@@ -0,0 +1,130 @@
+---
+defaults:
+  - scalers: global
+  - optimization: default
+  - weight_averaging: null
+
+# resume or fork a training from a checkpoint last.ckpt or specified in hardware.files.warm_start
+run_id: null
+fork_run_id: ???
+transfer_learning: True # activate to perform transfer learning
+load_weights_only: True # only load model weights, do not restore optimiser states etc.
+update_ds_stats_on_ckpt_load:
+  states: False # rebuild state processors from current dataset when loading a checkpoint
+  tendencies: True # rebuild tendency processors from current dataset when loading a checkpoint
+
+# run in deterministic mode ; slows down
+deterministic: False
+
+# miscellaneous
+precision: 16-mixed
+
+# 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
+
+num_sanity_val_steps: 6
+
+# clipp gradients, 0 : don't clip, default algorithm: norm, alternative: value
+gradient_clip:
+  val: 32.
+  algorithm: value
+
+# select model
+training_method: anemoi.training.train.methods.LoRASingleTraining
+
+lora_config:
+  r: 8
+  lora_alpha: 32
+  target_modules: ['mlp.0', 'mlp.2', 'dummy_layer']
+  modules_to_save: ['node_data_extractor.1']
+
+# select strategy
+strategy:
+  _target_: anemoi.training.distributed.strategy.DDPGroupStrategy
+  num_gpus_per_model: ${system.hardware.num_gpus_per_model}
+  read_group_size: ${dataloader.read_group_size}
+
+# loss functions
+
+# 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
+
+# loss function for the model
+training_loss:
+  datasets:
+    data: # user-defined key in data
+      # loss class to initialise
+      _target_: anemoi.training.losses.MSELoss
+      # Scalers to include in loss calculation
+      # A selection of available scalers are listed in training/scalers.
+      # '*' is a valid entry to use all `scalers` given, if a scaler is to be excluded
+      # add `!scaler_name`, i.e. ['*', '!scaler_1'], and `scaler_1` will not be added.
+      scalers: ['pressure_level', 'general_variable', 'node_weights', 'time_steps']
+      ignore_nans: False
+
+# Validation metrics calculation,
+# This may be a list, in which case all metrics will be calculated
+# and logged according to their name.
+# These metrics are calculated in the output model space, and thus
+# have undergone postprocessing.
+validation_metrics:
+  datasets:
+    data: # user-defined key in data
+      # loss class to initialise
+      mse:
+        _target_: anemoi.training.losses.MSELoss
+        # Scalers to include in loss calculation
+        # Cannot scale over the variable dimension due to possible remappings.
+        # Available scalers include:
+        # - 'loss_weights_mask': Giving imputed NaNs a zero weight in the loss function
+        # Use the `scale_validation_metrics` section to variable scale.
+        scalers: ['node_weights', 'time_steps']
+        # other kwargs
+        ignore_nans: True
+
+# Variable groups definition for scaling
+# The variable level scaling methods are defined under training/scalers
+# A default group is required and is appended as prefix to the metric of all variables not assigned to a group.
+# Variables are assigned to a group by their param if contained in the metadata, else by their name.
+
+# If more complex grouping is required, groups can be defined as a dictionary, such that all
+# keys must be evaluate to True.
+# .e.g. to set the variable group based on if the metadata specifies the variable is a pressure level
+# you can write the following:
+#   variable_groups:
+#     default: sfc
+#     pl:
+#        is_pressure_level: True
+# See `anemoi.transform.variables.Variable` for the available metadata.
+# Note that the former formulation of
+# :
+# variable_groups:
+#  default: sfc
+#  pl: [q, t, u, v, w, z]
+#
+# still works
+# param is an alias for the variable name in the case of no metadata.
+
+variable_groups:
+  datasets:
+    data: # user-defined key in data
+      default: sfc
+      pl:
+        param: [q, t, u, v, w, z]
+
+metrics:
+  datasets:
+    data: # user-defined key in data
+      - z_500
+      - t_850
+      - u_850
+      - v_850
+
+# Set max_epochs or max_steps. Training stops at the first limit reached.
+max_epochs: null
+max_steps: 150000
+
+submodules_to_freeze: []

training/src/anemoi/training/schemas/training.py

@@ -114,6 +114,23 @@ class TargetForcing(BaseModel):
     "Use target time as a fraction between input boundary times as input."
 
 
+class LoRAConfig(BaseModel):
+    """LoRA parameters.
+
+    See https://huggingface.co/docs/peft/package_reference/lora#peft.LoraConfig
+    for more information.
+    """
+
+    r: int = 8
+    "Lora attention dimension (the “rank”)."
+    lora_alpha: int = 32
+    "The alpha parameter for Lora scaling."
+    target_modules: list[str] = Field(examples=["mlp.0", "mlp.2"])
+    "The names of the modules to apply the adapter to."
+    modules_to_save: list[str] = Field(examples=["node_data_extractor.1"])
+    "List of modules apart from adapter layers to be set as trainable."
+
+
 class LossScalingSchema(BaseModel):
     default: int = 1
     "Default scaling value applied to the variables loss. Default to 1."
@@ -473,7 +490,18 @@ class DiffusionTendencyTrainingSchema(BaseTrainingSchema):
     "Training objective."
 
 
+class LoRASingleTrainingSchema(BaseTrainingSchema):
+    training_method: Literal["anemoi.training.train.methods.LoRASingleTraining"] = Field(..., alias="training_method")
+    "Training objective."
+    lora_config: LoRAConfig
+    "Configuration for the LoRA adapter."
+
+
 TrainingSchema = Annotated[
-    SingleTrainingSchema | EnsembleTrainingSchema | DiffusionTrainingSchema | DiffusionTendencyTrainingSchema,
+    SingleTrainingSchema
+    | EnsembleTrainingSchema
+    | DiffusionTrainingSchema
+    | DiffusionTendencyTrainingSchema
+    | LoRASingleTrainingSchema,
     Discriminator("training_method"),
 ]

training/src/anemoi/training/train/methods/__init__.py

@@ -10,11 +10,13 @@
 from .diffusion import DiffusionTendencyTraining
 from .diffusion import DiffusionTraining
 from .ensemble import EnsembleTraining
+from .lora_single import LoRASingleTraining
 from .single import SingleTraining
 
 __all__ = [
     "DiffusionTendencyTraining",
     "DiffusionTraining",
     "EnsembleTraining",
+    "LoRASingleTraining",
     "SingleTraining",
 ]

training/src/anemoi/training/train/methods/base.py

@@ -442,6 +442,13 @@ def on_load_checkpoint(self, checkpoint: torch.nn.Module) -> None:
             for dataset_name, data_indices in checkpoint["hyper_parameters"]["data_indices"].items()
         }
 
+    def on_checkpoint_loaded(self) -> None:
+        """Called once a model has been loaded from a checkpoint.
+
+        Override to perform actions on the model. This applies to both transfer
+        learning and weights-only settings.
+        """
+
     def _update_scaler_for_dataset(
         self,
         name: str,

training/src/anemoi/training/train/methods/lora_single.py

@@ -0,0 +1,78 @@
+# (C) Copyright 2024 Anemoi contributors.
+# Copyright (C) Bull S.A.S - 2025
+#
+# 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
+
+from peft import LoraConfig
+from peft import PeftModel
+from peft import get_peft_model
+
+from anemoi.training.train.methods.single import SingleTraining
+
+if TYPE_CHECKING:
+    import torch
+    from torch_geometric.data import HeteroData
+
+    from anemoi.models.data_indices.collection import IndexCollection
+    from anemoi.training.schemas.base_schema import BaseSchema
+    from training.src.anemoi.training.tasks.base import BaseTask
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+class LoRASingleTraining(SingleTraining):
+    """Graph neural network forecaster for PyTorch Lightning."""
+
+    def __init__(
+        self,
+        *,
+        config: BaseSchema,
+        task: BaseTask,
+        graph_data: HeteroData,
+        statistics: dict,
+        statistics_tendencies: dict,
+        data_indices: dict[str, IndexCollection],
+        metadata: dict,
+        supporting_arrays: dict,
+    ) -> None:
+
+        super().__init__(
+            config=config,
+            task=task,
+            graph_data=graph_data,
+            statistics=statistics,
+            statistics_tendencies=statistics_tendencies,
+            data_indices=data_indices,
+            metadata=metadata,
+            supporting_arrays=supporting_arrays,
+        )
+
+        self.lora_config = LoraConfig(**config.training.lora_config)
+
+    def on_load_checkpoint(self, checkpoint: torch.nn.Module) -> None:
+        self._update_checkpoint_state_dict_for_load(checkpoint)
+
+        self._ckpt_model_name_to_index = {
+            dataset_name: data_indices.name_to_index
+            for dataset_name, data_indices in checkpoint["hyper_parameters"]["data_indices"].items()
+        }
+        if "LoRASingleTraining" in checkpoint["hyper_parameters"]["config"].training.training_method:
+            self._inject_lora_adapters()
+
+    def on_checkpoint_loaded(self) -> None:
+        if not isinstance(self.model, PeftModel):
+            self._inject_lora_adapters()
+
+    def _inject_lora_adapters(self) -> None:
+        get_peft_model(self.model, self.lora_config)
+        LOGGER.info("LoRA adapters injected into the model")

training/src/anemoi/training/train/train.py

@@ -297,6 +297,8 @@ def model(self) -> pl.LightningModule:
             if self.config.training.transfer_learning:
                 LOGGER.info("Loading weights with Transfer Learning from %s", self.last_checkpoint)
                 model = transfer_learning_loading(model, self.last_checkpoint)
+                # Added for LoRA
+                model.on_checkpoint_loaded()
             else:
                 LOGGER.info("Restoring only model weights from %s", self.last_checkpoint)
                 # pop data_indices so that the data indices on the checkpoint do not get overwritten
@@ -308,6 +310,8 @@ def model(self) -> pl.LightningModule:
                     strict=False,
                     weights_only=False,  # required for Pytorch Lightning 2.6
                 )
+                # Added for LoRA
+                model.on_checkpoint_loaded()
 
             model.data_indices = self.data_indices
             # Validate data indices between checkpoint and current config

training/tests/integration/config/test_lora.yaml

@@ -0,0 +1,17 @@
+# Modifications for the basic template "config.yaml"
+system:
+  input:
+    dataset: anemoi-integration-tests/training/datasets/aifs-ea-an-oper-0001-mars-o96-2017-2017-6h-v8-testing.zarr
+
+training:
+  fork_run_id: "dummy_id"
+
+dataloader:
+  training:
+    datasets:
+      data:
+        end: 2017-01-08 12:00:00
+  validation:
+    datasets:
+      data:
+        start: 2017-01-08 18:00:00

training/tests/integration/conftest.py

@@ -549,6 +549,38 @@ def diffusion_config(
     return cfg, url_dataset
 
 
+@pytest.fixture
+def lora_config(
+    testing_modifications_with_temp_dir: DictConfig,
+    get_tmp_path: GetTmpPath,
+    get_test_data: GetTestData,
+    migrator: Migrator,
+) -> tuple[DictConfig, str]:
+    with initialize(version_base=None, config_path="../../src/anemoi/training/config", job_name="test_lora"):
+        template = compose(config_name="lora")
+
+    use_case_modifications = OmegaConf.load(Path.cwd() / "training/tests/integration/config/test_lora.yaml")
+    tmp_dir_dataset, url_dataset = get_tmp_path(use_case_modifications.system.input.dataset)
+    use_case_modifications.system.input.dataset = str(tmp_dir_dataset)
+
+    cfg = OmegaConf.merge(template, testing_modifications_with_temp_dir, use_case_modifications)
+    OmegaConf.resolve(cfg)
+
+    existing_ckpt = get_test_data(
+        "anemoi-integration-tests/training/checkpoints/testing-checkpoint-gnn-global-2026-03-06.ckpt",
+    )
+    _, new_ckpt, _ = migrator.sync(existing_ckpt)
+
+    checkpoint_dir = Path(cfg.system.output.root + "/" + cfg.system.output.checkpoints.root + "/dummy_id")
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    torch.save(new_ckpt, checkpoint_dir / "last.ckpt")
+
+    cfg.diagnostics.plot.callbacks = []  # remove plotting callbacks as they are tested in global training cycle test
+    cfg.diagnostics.callbacks = []  # remove RolloutEval callback as it is tested in global training cycle test
+
+    return cfg, url_dataset
+
+
 @pytest.fixture(
     params=[
         pytest.param(