diff --git a/training/pyproject.toml b/training/pyproject.toml index e3712e92fd..8723b8d105 100644 --- a/training/pyproject.toml +++ b/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", diff --git a/training/src/anemoi/training/config/lora.yaml b/training/src/anemoi/training/config/lora.yaml new file mode 100644 index 0000000000..03c0f98a7a --- /dev/null +++ b/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'] diff --git a/training/src/anemoi/training/config/training/lora.yaml b/training/src/anemoi/training/config/training/lora.yaml new file mode 100644 index 0000000000..d1707a8ff8 --- /dev/null +++ b/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: [] diff --git a/training/src/anemoi/training/schemas/training.py b/training/src/anemoi/training/schemas/training.py index e5469e8700..6399932828 100644 --- a/training/src/anemoi/training/schemas/training.py +++ b/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"), ] diff --git a/training/src/anemoi/training/train/methods/__init__.py b/training/src/anemoi/training/train/methods/__init__.py index 7bd465df38..84a6166ee3 100644 --- a/training/src/anemoi/training/train/methods/__init__.py +++ b/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", ] diff --git a/training/src/anemoi/training/train/methods/base.py b/training/src/anemoi/training/train/methods/base.py index 18bdf0adcb..826d784cdb 100644 --- a/training/src/anemoi/training/train/methods/base.py +++ b/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, diff --git a/training/src/anemoi/training/train/methods/lora_single.py b/training/src/anemoi/training/train/methods/lora_single.py new file mode 100644 index 0000000000..5952ecc757 --- /dev/null +++ b/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") diff --git a/training/src/anemoi/training/train/train.py b/training/src/anemoi/training/train/train.py index e48abe1029..5ba5b235f2 100644 --- a/training/src/anemoi/training/train/train.py +++ b/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 diff --git a/training/tests/integration/config/test_lora.yaml b/training/tests/integration/config/test_lora.yaml new file mode 100644 index 0000000000..9b44cac73a --- /dev/null +++ b/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 diff --git a/training/tests/integration/conftest.py b/training/tests/integration/conftest.py index 2a692aa956..137caa7b49 100644 --- a/training/tests/integration/conftest.py +++ b/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( diff --git a/training/tests/integration/test_training_cycle.py b/training/tests/integration/test_training_cycle.py index c288188a9a..dc5b85d0c5 100644 --- a/training/tests/integration/test_training_cycle.py +++ b/training/tests/integration/test_training_cycle.py @@ -348,6 +348,46 @@ def test_config_validation_diffusion(diffusion_config: tuple[DictConfig, str]) - @skip_if_offline @pytest.mark.slow +def test_training_cycle_lora(lora_config: tuple[DictConfig, list[str]], get_test_archive: GetTestArchive) -> None: + cfg, url = lora_config + get_test_archive(url) + trainer = AnemoiTrainer(cfg) + has_lora = False + for name, _param in trainer.model.named_parameters(): + if "lora_A" in name: + has_lora = True + break + assert has_lora, "Expected model to have LoRA embeddings for LoRA training" + trainer.train() + + output_dir = Path(cfg.system.output.root + "/" + cfg.system.output.checkpoints.root) + assert output_dir.exists(), f"Checkpoint directory not found at: {output_dir}" + + run_dirs = [item for item in output_dir.iterdir() if item.is_dir() and item.name != "dummy_id"] + assert ( + len(run_dirs) == 1 + ), f"Expected exactly one run_id directory, found {len(run_dirs)}: {[d.name for d in run_dirs]}" + + checkpoint_dir = run_dirs[0] + assert len(list(checkpoint_dir.glob("anemoi-by_epoch-*.ckpt"))) == 2, "Expected 2 checkpoints after first run" + + cfg.training.run_id = checkpoint_dir.name + cfg.training.max_epochs = cfg.training.max_epochs + 3 + trainer = AnemoiTrainer(cfg) + has_lora = False + for name, _param in trainer.model.named_parameters(): + if "lora_A" in name: + has_lora = True + break + assert has_lora, "Expected model to have LoRA embeddings for LoRA training" + trainer.train() + + +def test_config_validation_lora(lora_config: tuple[DictConfig, str]) -> None: + cfg, _ = lora_config + BaseSchema(**cfg) + + @pytest.mark.mlflow def test_training_cycle_mlflow_dry_run( mlflow_dry_run_config: tuple[DictConfig, str],