diff --git a/.github/workflows/python-pull-request.yml b/.github/workflows/python-pull-request.yml index ee4284cdea..86feda713f 100644 --- a/.github/workflows/python-pull-request.yml +++ b/.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 diff --git a/training/pyproject.toml b/training/pyproject.toml index e3712e92fd..d2cd49685a 100644 --- a/training/pyproject.toml +++ b/training/pyproject.toml @@ -46,10 +46,11 @@ dependencies = [ "numba", "numpy", "nvidia-ml-py>=13.580.82", + "peft>=0.17.1", "pydantic>=2.9", "pytorch-lightning>=2.1", "timm>=0.9.2", - "torch>=2.2,<2.11", + "torch>=2.2,<2.12", "torch-geometric>=2.3.1", "torchinfo>=1.8", "torchvision>=0.18", @@ -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" diff --git a/training/src/anemoi/training/config/peft.yaml b/training/src/anemoi/training/config/peft.yaml new file mode 100644 index 0000000000..102b3b2dac --- /dev/null +++ b/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: ??? diff --git a/training/src/anemoi/training/config/training/peft.yaml b/training/src/anemoi/training/config/training/peft.yaml new file mode 100644 index 0000000000..6a5cf52a03 --- /dev/null +++ b/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: [] diff --git a/training/src/anemoi/training/schemas/training.py b/training/src/anemoi/training/schemas/training.py index e5469e8700..4f83db542a 100644 --- a/training/src/anemoi/training/schemas/training.py +++ b/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"), ] diff --git a/training/src/anemoi/training/train/methods/__init__.py b/training/src/anemoi/training/train/methods/__init__.py index 7bd465df38..825095ebb5 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 .peft_single import PeftSingleTraining from .single import SingleTraining __all__ = [ "DiffusionTendencyTraining", "DiffusionTraining", "EnsembleTraining", + "PeftSingleTraining", "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/peft_single.py b/training/src/anemoi/training/train/methods/peft_single.py new file mode 100644 index 0000000000..3c76b7d70b --- /dev/null +++ b/training/src/anemoi/training/train/methods/peft_single.py @@ -0,0 +1,79 @@ +# (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 hydra.utils import get_class +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 PeftSingleTraining(SingleTraining): + """PEFT training method.""" + + 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, + ) + + peft_config_class = get_class(config.training.peft_config._target_) + self.peft_config = peft_config_class(**config.training.peft_config.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 "PeftSingleTraining" in checkpoint["hyper_parameters"]["config"].training.training_method: + self._inject_peft_adapters() + + def on_checkpoint_loaded(self) -> None: + if not isinstance(self.model, PeftModel): + self._inject_peft_adapters() + + def _inject_peft_adapters(self) -> None: + get_peft_model(self.model, self.peft_config) + LOGGER.info("PEFT 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_lily.yaml b/training/tests/integration/config/test_lily.yaml new file mode 100644 index 0000000000..c74d81383a --- /dev/null +++ b/training/tests/integration/config/test_lily.yaml @@ -0,0 +1,26 @@ +# 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" + peft_config: + _target_: peft.LilyConfig + config: + r: 32 + stride_A: 2 + num_B: 4 + scaling: 1.0 + target_modules: ['mlp.0', 'mlp.2', 'dummy_layer'] + modules_to_save: ['node_data_extractor.1'] + +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/config/test_lora.yaml b/training/tests/integration/config/test_lora.yaml new file mode 100644 index 0000000000..e1e06d3af7 --- /dev/null +++ b/training/tests/integration/config/test_lora.yaml @@ -0,0 +1,26 @@ +# 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" + peft_config: + _target_: peft.LoraConfig + config: + r: 8 + lora_alpha: 32 + target_modules: ['mlp.0', 'mlp.2', 'dummy_layer'] + modules_to_save: ['node_data_extractor.1'] + use_dora: True + use_rslora: True + +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/config/test_miss.yaml b/training/tests/integration/config/test_miss.yaml new file mode 100644 index 0000000000..9cb3c607d5 --- /dev/null +++ b/training/tests/integration/config/test_miss.yaml @@ -0,0 +1,27 @@ +# 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" + peft_config: + _target_: peft.MissConfig + config: + r: 4 + mini_r: 1 + init_weights: bat + miss_dropout: 0.1 + bias: MiSS_only + target_modules: ['node_mlp.mlp.0', 'node_mlp.mlp.2'] + modules_to_save: ['node_data_extractor.1'] + +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/config/test_oft.yaml b/training/tests/integration/config/test_oft.yaml new file mode 100644 index 0000000000..d5dbd2fa04 --- /dev/null +++ b/training/tests/integration/config/test_oft.yaml @@ -0,0 +1,27 @@ +# 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" + peft_config: + _target_: peft.OFTConfig + config: + r: 0 + oft_block_size: 32 + module_dropout: 0.1 + use_cayley_neumann: True + bias: oft_only + target_modules: ['mlp.0', 'mlp.2', 'dummy_layer'] + modules_to_save: ['node_data_extractor.1'] + +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..78a903a3df 100644 --- a/training/tests/integration/conftest.py +++ b/training/tests/integration/conftest.py @@ -549,6 +549,134 @@ 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="peft") + + 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 +def lily_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_lily"): + template = compose(config_name="peft") + + use_case_modifications = OmegaConf.load(Path.cwd() / "training/tests/integration/config/test_lily.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 +def miss_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_miss"): + template = compose(config_name="peft") + + use_case_modifications = OmegaConf.load(Path.cwd() / "training/tests/integration/config/test_miss.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 +def oft_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_oft"): + template = compose(config_name="peft") + + use_case_modifications = OmegaConf.load(Path.cwd() / "training/tests/integration/config/test_oft.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..555a598610 100644 --- a/training/tests/integration/test_training_cycle.py +++ b/training/tests/integration/test_training_cycle.py @@ -348,6 +348,188 @@ 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" + + # Test resume training + cfg.training.run_id = checkpoint_dir.name + cfg.training.fork_run_id = None + cfg.training.max_epochs = cfg.training.max_epochs + 3 + cfg.training.transfer_learning = False + cfg.training.load_weights_only = False + trainer = AnemoiTrainer(cfg) + has_lora = False + trainer.train() + 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" + + +def test_config_validation_lora(lora_config: tuple[DictConfig, str]) -> None: + cfg, _ = lora_config + BaseSchema(**cfg) + + +@skip_if_offline +@pytest.mark.slow +def test_training_cycle_lily(lily_config: tuple[DictConfig, list[str]], get_test_archive: GetTestArchive) -> None: + cfg, url = lily_config + get_test_archive(url) + trainer = AnemoiTrainer(cfg) + has_lily = False + for name, _param in trainer.model.named_parameters(): + if "lily_A" in name: + has_lily = True + break + assert has_lily, "Expected model to have Lily embeddings for Lily 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" + + # Test resume training + cfg.training.run_id = checkpoint_dir.name + cfg.training.fork_run_id = None + cfg.training.max_epochs = cfg.training.max_epochs + 3 + cfg.training.transfer_learning = False + cfg.training.load_weights_only = False + trainer = AnemoiTrainer(cfg) + has_lily = False + trainer.train() + for name, _param in trainer.model.named_parameters(): + if "lily_A" in name: + has_lily = True + break + assert has_lily, "Expected model to have Lily embeddings for Lily training" + + +def test_config_validation_lily(lily_config: tuple[DictConfig, str]) -> None: + cfg, _ = lily_config + BaseSchema(**cfg) + + +@skip_if_offline +@pytest.mark.slow +def test_training_cycle_miss(miss_config: tuple[DictConfig, list[str]], get_test_archive: GetTestArchive) -> None: + cfg, url = miss_config + get_test_archive(url) + trainer = AnemoiTrainer(cfg) + has_miss = False + for name, _param in trainer.model.named_modules(): + if "miss_block" in name: + has_miss = True + break + assert has_miss, "Expected model to have Miss block for Miss 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" + + # Test resume training + cfg.training.run_id = checkpoint_dir.name + cfg.training.fork_run_id = None + cfg.training.max_epochs = cfg.training.max_epochs + 3 + cfg.training.transfer_learning = False + cfg.training.load_weights_only = False + trainer = AnemoiTrainer(cfg) + has_miss = False + trainer.train() + for name, _param in trainer.model.named_modules(): + if "miss_block" in name: + has_miss = True + break + assert has_miss, "Expected model to have Miss block for Miss training" + + +def test_config_validation_miss(miss_config: tuple[DictConfig, str]) -> None: + cfg, _ = miss_config + BaseSchema(**cfg) + + +@skip_if_offline +@pytest.mark.slow +def test_training_cycle_oft(oft_config: tuple[DictConfig, list[str]], get_test_archive: GetTestArchive) -> None: + cfg, url = oft_config + get_test_archive(url) + trainer = AnemoiTrainer(cfg) + has_oft = False + for name, _param in trainer.model.named_modules(): + if "oft_R" in name: + has_oft = True + break + assert has_oft, "Expected model to have Oft block for Oft 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" + + # Test resume training + cfg.training.run_id = checkpoint_dir.name + cfg.training.fork_run_id = None + cfg.training.max_epochs = cfg.training.max_epochs + 3 + cfg.training.transfer_learning = False + cfg.training.load_weights_only = False + trainer = AnemoiTrainer(cfg) + has_oft = False + trainer.train() + for name, _param in trainer.model.named_modules(): + if "oft_R" in name: + has_oft = True + break + assert has_oft, "Expected model to have Oft block for Oft training" + + +def test_config_validation_oft(oft_config: tuple[DictConfig, str]) -> None: + cfg, _ = oft_config + BaseSchema(**cfg) + + @pytest.mark.mlflow def test_training_cycle_mlflow_dry_run( mlflow_dry_run_config: tuple[DictConfig, str],