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

.github/workflows/python-pull-request.yml

@@ -2,7 +2,7 @@ name: Code Quality and Testing
 
 on:
   pull_request:
-    types: [opened, synchronize, reopened]
+    types: [opened, synchronize, reopened, edited]
   push:
     branches:
       - main

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",
@@ -206,4 +207,5 @@ lint.ignore = [
   "UP007",
 ]
 lint.flake8-import-conventions.extend-aliases.pytorch_lightning = "pl"
+lint.pep8-naming.ignore-names = [ "num_B", "stride_A" ]  # peft Lily parameters
 lint.pydocstyle.convention = "numpy"

training/src/anemoi/training/config/peft.yaml

@@ -0,0 +1,19 @@
+defaults:
+- data: zarr
+- dataloader: native_grid
+- diagnostics: evaluation
+- system: example
+- graph: multi_scale
+- model: gnn
+- task: forecaster
+- training: peft
+- _self_
+
+
+# set to true to switch on config validation
+config_validation: True
+
+training:
+  peft_config:
+    _target_: ???
+    config: ???

training/src/anemoi/training/config/training/peft.yaml

@@ -0,0 +1,128 @@
+---
+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.PeftSingleTraining
+
+peft_config:
+  _target_: ???
+  config: ???
+
+# 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

@@ -13,6 +13,7 @@
 from typing import Any
 from typing import Literal
 from typing import Self
+from typing import Union
 
 from pydantic import AfterValidator
 from pydantic import ConfigDict
@@ -21,6 +22,8 @@
 from pydantic import NonNegativeFloat
 from pydantic import NonNegativeInt
 from pydantic import PositiveInt
+from pydantic import TypeAdapter
+from pydantic import ValidationError
 from pydantic import field_validator
 from pydantic import model_validator
 
@@ -114,6 +117,133 @@ class TargetForcing(BaseModel):
     "Use target time as a fraction between input boundary times as input."
 
 
+class LoRASchema(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] | 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."
+    lora_dropout: float = 0.0
+    "The dropout probability for Lora layers."
+    use_dora: bool = False
+    "Enable Weight-Decomposed Low-Rank Adaptation (DoRA)."
+    use_rslora: bool = False
+    "Rank-Stabilized LoRA."
+
+
+class LilySchema(BaseModel):
+    """Lily parameters.
+
+    See https://huggingface.co/docs/peft/v0.20.0/en/package_reference/lily
+    for more information.
+    """
+
+    r: int = 32
+    "Lily's rank. Typically 2x, 3x, or 4x the LoRA rank you would normally use."
+    stride_A: int = 2
+    "Stride for the A matrix. if stride_A=4, every 4 adjacent layers share the same A adapter."
+    num_B: int = 4
+    " The number of shared B adapters."
+    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."
+    scaling: float = 1.0
+    "A scalar multiplier applied to the combined adapter output (scaling * A @ combined_B)."
+
+
+class MissSchema(BaseModel):
+    """Miss parameters.
+
+    See https://huggingface.co/docs/peft/v0.20.0/en/package_reference/miss
+    for more information.
+    """
+
+    r: int = 32
+    " The rank of MiSS across different layers. It is best to set r to an even number."
+    mini_r: int = 1
+    "The rank of MiSS corresponds to a low-rank decomposition along the out_features dimension."
+    init_weights: bool | Literal["bat", "mini"] = True
+    "Different initializations correspond to different MiSS variants."
+    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."
+    miss_dropout: float = 0.0
+    "The dropout probability for Miss layers."
+    bias: Literal["none", "all", "MiSS_only"] = "none"
+    "Bias type for MiSS. Can be 'none', 'all' or 'MiSS_only'."
+
+
+class OFTSchema(BaseModel):
+    """OFT parameters.
+
+    See https://huggingface.co/docs/peft/v0.20.0/en/package_reference/oft
+    for more information.
+    """
+
+    r: int = 0
+    "The number of OFT blocks per injected layer."
+    oft_block_size: int = 32
+    "OFT block size across different layers."
+    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."
+    module_dropout: float = 0.0
+    "The multiplicative dropout probability, by setting OFT blocks to identity during training."
+    use_cayley_neumann: bool = False
+    """Use the Cayley-Neumann parameterization (efficient but approximate) or the vanilla Cayley
+    parameterization (exact but computationally expensive because of matrix inverse)."""
+    bias: Literal["none", "all", "oft_only"] = "none"
+    "Bias type for OFT. Can be 'none', 'all' or 'oft_only'."
+
+
+class PeftTarget(StrEnum):
+    lora = "peft.LoraConfig"
+    lily = "peft.LilyConfig"
+    miss = "peft.MissConfig"
+    oft = "peft.OFTConfig"
+
+
+target_to_schema = {
+    PeftTarget.lora: LoRASchema,
+    PeftTarget.lily: LilySchema,
+    PeftTarget.miss: MissSchema,
+    PeftTarget.oft: OFTSchema,
+}
+
+
+class PeftSchema(BaseModel, validate_assignment=False):
+    target_: PeftTarget = Field(..., alias="_target_")
+    "Peft configuration target."
+    config: Union[dict, LoRASchema, LilySchema, MissSchema, OFTSchema]
+    "Target schema containing peft configuration."
+
+    @model_validator(mode="after")
+    def schema_consistent_with_target(self) -> type["PeftSchema"]:
+        schema_cls = target_to_schema.get(self.target_)
+        if schema_cls is None:
+            error_msg = f"Unknown target: {self.target_}"
+            raise ValidationError(error_msg)
+
+        validated = TypeAdapter(schema_cls).validate_python(self.config)
+        # If it's a RootModel (like ConstantImputerSchema), extract the root dict
+        if hasattr(validated, "root"):
+            self.config = validated.root
+
+        return self
+
+
 class LossScalingSchema(BaseModel):
     default: int = 1
     "Default scaling value applied to the variables loss. Default to 1."
@@ -473,7 +603,18 @@ class DiffusionTendencyTrainingSchema(BaseTrainingSchema):
     "Training objective."
 
 
+class PeftSingleTrainingSchema(BaseTrainingSchema):
+    training_method: Literal["anemoi.training.train.methods.PeftSingleTraining"] = Field(..., alias="training_method")
+    "Training objective."
+    peft_config: PeftSchema
+    "Configuration for the Peft adapter."
+
+
 TrainingSchema = Annotated[
-    SingleTrainingSchema | EnsembleTrainingSchema | DiffusionTrainingSchema | DiffusionTendencyTrainingSchema,
+    SingleTrainingSchema
+    | EnsembleTrainingSchema
+    | DiffusionTrainingSchema
+    | DiffusionTendencyTrainingSchema
+    | PeftSingleTrainingSchema,
     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 .peft_single import PeftSingleTraining
 from .single import SingleTraining
 
 __all__ = [
     "DiffusionTendencyTraining",
     "DiffusionTraining",
     "EnsembleTraining",
+    "PeftSingleTraining",
     "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,