Basic Training Configuration

Anemoi training is designed so you can adjust key parts of the models and training process without needing to modify the underlying code.

A basic introduction to the configuration system is provided in Configuration Basics. This section will go into more detail on how to configure the training pipeline.

Default Config Groups

A typical config file will start with specifying the default config settings at the top as follows:

defaults:
- data: zarr
- dataloader: native_grid
- diagnostics: evaluation
- system: example
- graph: multi_scale
- model: gnn
- task: forecaster
- training: single
- _self_

These are group configs for each section. The options after the defaults are then used to override the configs, by assigning new features and keywords.

You can also find these defaults in other configs, like the system, which implements:

defaults:
- hardware: example
- input: example
- output: example

Config files are resolved in decreasing priority order: a path supplied via --config-path always takes precedence, followed by the current working directory, and finally the packaged defaults shipped with anemoi-training. A file or group override found in a higher-priority location shadows any matching file in a lower-priority location.

YAML-based config overrides

The config files are written in YAML format. This allows for easy overrides of the default settings. For example, to change the model from the default GNN to a transformer, you can use the following config in the config groups.:

model: transformer

This will override the default model config with the transformer model.

You can also override individual settings. For example, to change the learning rate from the default value of 0.625e-4 to 1e-3, you can add the following to the config you’re using:

training:
   lr:
      rate: 1e-3

You can also change the GPU count to whatever you have available:

system:
   hardware:
      num_gpus_per_node: 1

This matches the interface of the underlying defaults in Anemoi training.

Dataloader Breaking Change

Starting from the current configuration schema, dataloader dataset reader settings must be provided under dataset_config.

Use:

dataloader:
   training:
      datasets:
         your_dataset_name:
            dataset_config:
               dataset: ${system.input.dataset}
               frequency: ${data.frequency}
               drop: []
            start: 1985
            end: 2020

Do not use the previous dataset/name nesting. Configuration validation now enforces the new layout.

Example Config File

Here is an example of a config file that changes the model to a transformer, the learning rate to 1e-3, and the number of GPUs to 1. We also need to specify the paths to the data, output, and graph data and give the names of the files to use. You can get a dataset from the Anemoi Datasets catalogue or create one using the Anemoi Datasets package.

You can create a graph using Anemoi Graphs or one will be created for you at runtime. Note that you must specify a filename for the graph, here we use first_graph_m320.pt.

You’ll also notice we’ve specified a resolution for the data, this must match the dataset you provide.

defaults:
- data: zarr
- dataloader: native_grid
- diagnostics: evaluation
- system: example
- graph: multi_scale
- model: transformer # Change from default group
- task: forecaster
- training: single
- _self_

config_validation: True
data:
   resolution: n320

system:
   hardware:
      num_gpus_per_node: 1
   output:
      root: /home/username/anemoi/training/output
   input:
      dataset: datset-n320-2019-2021-6h.zarr
      graph: first_graph_n320.pt

When we save this example.yaml file, we can run the training with this config using:

anemoi-training train --config-name=example.yaml

Command-line config overrides

It is also possible to use command line config overrides. We can switch out group configs using

anemoi-training train model=transformer

or override individual config entries such as

anemoi-training train system.hardware.num_gpus_per_node=1

or combine everything together

anemoi-training train --config-name=debug.yaml model=transformer system.hardware.num_gpus_per_node=1

Config validation

It is possible to validate your configuration before starting a training run using the following command:

anemoi-training config validate --config-name debug.yaml

By default the config is looked up on the search path described above (the current working directory and the packaged defaults). To validate a config that lives somewhere else without changing directory, point --config-path at its directory, exactly as for anemoi-training train:

anemoi-training config validate --config-path /path/to/configs --config-name debug.yaml

This will check that the configuration is valid and that all the required fields are present. If your config is correctly defined then the command will show an output similar to:

2025-01-28 09:37:23 INFO Validating configs.
2025-01-28 09:37:23 INFO Appending current working directory (/repos_path/config_anemoi_core) to the search path.
2025-01-28 09:37:23 INFO Search path is now: [provider=hydra, path=pkg://hydra.conf, provider=main, path=/repos_path/anemoi-core/training/src/anemoi/training/commands, provider=anemoi-cwd-searchpath-plugin, path=/repos_path/config_anemoi_core, provider=anemoi-package-searchpath-plugin, path=pkg://anemoi.training/config]
cfg = BaseSchema(**cfg)
2025-01-28 09:37:23 INFO Config files validated.

Otherwise if there is an issue with some of your configuration fields, Pydantic will report an error message. If your config is missing the definition of a required field, then the validation will also fail. This can be the case if you pull the defaults anemoi configs and do not replace the empty fields (usually represented by ‘??’) with the actual values. Similarly if you have a field that is expected to use an environment variable and you do not have it set, the validation will fail. To overcome this issue and still be able to validate the config, you can use the —-mask_env_vars flag, which will skip the validation of the environment variables. When using this flag, the validation will still be performed, but the environment variables will be masked with the default values. See below an example output where we have 5 environment variables that are not set and are masked with the default values:

(anemoi_core_venv)[] $ anemoi-training config validate --config-name=debug --mask_env_vars
2025-02-16 17:48:38 INFO Validating configs.
2025-02-16 17:48:38 WARNING Note that this command is not taking into account if your config has
set the config_validation flag to false.So this command will validate the config regardless of the flag.
2025-01-28 09:37:23 INFO Appending current working directory (/repos_path/config_anemoi_core) to the search path.
2025-01-28 09:37:23 INFO Search path is now: [provider=hydra, path=pkg://hydra.conf, provider=main, path=/repos_path/anemoi-core/training/src/anemoi/training/commands, provider=anemoi-cwd-searchpath-plugin, path=/repos_path/config_anemoi_core, provider=anemoi-package-searchpath-plugin, path=pkg://anemoi.training/config]
2025-02-16 17:48:39 WARNING Environment variable EXP_NAME not found, masking with default
2025-02-16 17:48:39 WARNING Environment variable RUN_NAME not found, masking with default
2025-02-16 17:48:39 WARNING Environment variable SLURM_GPUS_PER_NODE not found, masking with 0
2025-02-16 17:48:39 WARNING Environment variable SLURM_NNODES not found, masking with 0
2025-02-16 17:48:39 WARNING Environment variable LOCAL_LR not found, masking with 0
2025-02-16 17:48:39 INFO Config files validated.

See example below where we have a debug.yaml file with a field not correctly indented (in this case the diagnostics.log field):

defaults:
- data: zarr
- dataloader: native_grid
- diagnostics: evaluation
- system: example
- graph: multi_scale
- model: transformer # Change from default group
- task: forecaster
- training: single
- _self_


diagnostics:
log:
mlflow:
   enabled: True
   offline: True
   experiment_name: 'test'
   project_name: 'AIFS'
   run_name: 'test_anemoi_core'
   tracking_uri: 'https://mlflow-server.int'
   authentication: True
   terminal: True

If we try to validate the above then the validate command will report the following error:

2025-01-28 09:37:23 INFO Validating configs.
2025-01-28 09:37:23 INFO Appending current working directory (/repos_path/config_anemoi_core) to the search path.
2025-01-28 09:37:23 INFO Search path is now: [provider=hydra, path=pkg://hydra.conf, provider=main, path=/repos_path/anemoi-core/training/src/anemoi/training/commands, provider=anemoi-cwd-searchpath-plugin, path=/repos_path/config_anemoi_core, provider=anemoi-package-searchpath-plugin, path=pkg://anemoi.training/config]
pydantic_core._pydantic_core.ValidationError: 1 validation error for BaseSchema
diagnostics.log
 Input should be a valid dictionary or instance of LoggingSchema [type=model_type, input_value=None, input_type=NoneType]
   For further information visit https://errors.pydantic.dev/2.10/v/model_type
2025-01-28 09:54:08 ERROR
💣 1 validation error for BaseSchema
diagnostics.log
Input should be a valid dictionary or instance of LoggingSchema [type=model_type, input_value=None, input_type=NoneType]
   For further information visit https://errors.pydantic.dev/2.10/v/model_type
2025-01-28 09:54:08 ERROR 💣 Exiting

Which indicates that the diagnostics.log field is not correctly defined as it should be a dictionary or instance of LoggingSchema. Please note there might still be cases not captured by the current schemas, so it is always good to double check the configuration file before running the training. See below an example of a config with some typos that might still need to be fixed manually:

defaults:
- data: zarr
- dataloader: native_grid
- diagnostics: evaluation
- system: example
- graph: multi_scale
- model: transformer # Change from default group
- task: forecaster
- training: single
- _self_


diagnostics:
log:
   mlflow:
      enabled: True
      ofline: True # this is a typo - should be offline
      experiment_name: 'test'
      project_name: 'AIFS'
      run_name: 'test_anemoi_core'
      tracking_uri: 'https://mlflow-server.int'
      authentication: True
      terminal: True

In the example above, if there is a default already defined for offline under diagnostics: evaluation then the validation will be successful, and in the high-level config (ie debug) ofline it will just simply not be used, since it has a typo. Otherwise, if there is no default for offline then the validation will fail, with the following error:

2025-01-28 09:37:23 INFO Validating configs.
2025-01-28 09:37:23 INFO Appending current working directory (/repos_path/config_anemoi_core) to the search path.
2025-01-28 09:37:23 INFO Search path is now:  [provider=hydra, path=pkg://hydra.conf, provider=main, path=/repos_path/anemoi-core/training/src/anemoi/training/commands, provider=anemoi-cwd-searchpath-plugin, path=/repos_path/config_anemoi_core, provider=anemoi-package-searchpath-plugin, path=pkg://anemoi.training/config]
pydantic_core._pydantic_core.ValidationError: 1 validation error for BaseSchema
diagnostics.log.mlflow.offline
Field required [type=missing, input_value={'enabled': True, 'authen...onfig'], 'ofline': True}, input_type=DictConfig]
   For further information visit https://errors.pydantic.dev/2.10/v/missing
2025-01-28 10:14:49 ERROR
💣 1 validation error for BaseSchema
diagnostics.log.mlflow.offline
Field required [type=missing, input_value={'enabled': True, 'authen...onfig'], 'ofline': True}, input_type=DictConfig]
   For further information visit https://errors.pydantic.dev/2.10/v/missing
2025-01-28 10:14:49 ERROR 💣 Exiting

That will indicate that the offline field is required and it is missing from the configuration file. If you identify any issues with the schemas or missing functionality, please raise an issue on the Anemoi Core repository.

Another type of error that we can see when working with Pydantic, is when we have a union of schemas, and then we try to validate using on those schemas config. For information about Unions, please refer to the Pydantic documentation. For example, let’s say we have a config with a union of schemas like the following:

defaults:
- data: zarr
- dataloader: native_grid
- diagnostics: evaluation
- system: example
- graph: multi_scale
- model: transformer # Change from default group
- task: forecaster
- training: single
- _self_


graphs:
   attributes:
      nodes:
          area_weight:
            _target_: anemoi.graphs.nodes.attributes.SphericalAreaWeights # options: Area, Uniform
            norm: unit-max # options: l1, l2, unit-max, unit-sum, unit-std

In the example above, Pydantic will try to validate the SphericalAreaWeights schema using the union NodeAttributeSchemas, which contains a list of all the possible schemas for graph nodes attributes.

NodeAttributeSchemas = Union[
   PlanarAreaWeightSchema
   | SphericalAreaWeightSchema
   | CutOutMaskSchema
   | NonmissingAnemoiDatasetVariableSchema
   | BooleanOperationSchema
]

If the schema is not correctly defined, then the validation will fail, with the following error:

2025-01-28 09:37:23 INFO Validating configs.
2025-01-28 09:37:23 INFO Appending current working directory (/repos_path/config_anemoi_core) to the search path.
2025-01-28 09:37:23 INFO Search path is now:  [provider=hydra, path=pkg://hydra.conf, provider=main, path=/repos_path/anemoi-core/training/src/anemoi/training/commands, provider=anemoi-cwd-searchpath-plugin, path=/repos_path/config_anemoi_core, provider=anemoi-package-searchpath-plugin, path=pkg://anemoi.training/config]
pydantic_core._pydantic_core.ValidationError: 1 validation error for BaseSchema
2025-01-28 10:14:49 ERROR
💣 14 validation error for BaseSchema
graph.nodes.data.attributes.area_weight.PlanarAreaWeightSchema._target_
Input should be 'anemoi.graphs.nodes.attributes.AreaWeights', 'anemoi.graphs.nodes.attributes.PlanarAreaWeights', 'anemoi.graphs.nodes.attributes.CutOutMask' or 'anemoi.graphs.nodes.attributes.UniformWeights' [type=literal_error, input_value='anemoi.graphs.nodes.attr...es.SphericalAreaWeights', input_type=str]
   For further information visit https://errors.pydantic.dev/2.10/v/literal_error
graph.nodes.data.attributes.area_weight.function-after[convert_centre_to_ndarray(), SphericalAreaWeightSchema].fill_value
Field required [type=missing, input_value={'_target_': 'anemoi.grap...ts', 'norm': 'unit-max'}, input_type=DictConfig]
   For further information visit https://errors.pydantic.dev/2.10/v/missing
graph.nodes.data.attributes.area_weight.CutOutMaskSchema._target_
Input should be 'anemoi.graphs.nodes.attributes.CutOutMask' [type=literal_error, input_value='anemoi.graphs.nodes.attr...es.SphericalAreaWeights', input_type=str]
   For further information visit https://errors.pydantic.dev/2.10/v/literal_error
graph.nodes.data.attributes.area_weight.NonmissingAnemoiDatasetVariableSchema._target_
Input should be 'anemoi.graphs.nodes.attributes.NonmissingAnemoiDatasetVariable' [type=literal_error, input_value='anemoi.graphs.nodes.attr...es.SphericalAreaWeights', input_type=str]
   For further information visit https://errors.pydantic.dev/2.10/v/literal_error
graph.nodes.data.attributes.area_weight.NonmissingAnemoiDatasetVariableSchema.variable
Field required [type=missing, input_value={'_target_': 'anemoi.grap...ts', 'norm': 'unit-max'}, input_type=DictConfig]
   For further information visit https://errors.pydantic.dev/2.10/v/missing
graph.nodes.data.attributes.area_weight.BooleanOperationSchema._target_
Input should be 'anemoi.graphs.nodes.attributes.BooleanNot', 'anemoi.graphs.nodes.attributes.BooleanAndMask' or 'anemoi.graphs.nodes.attributes.BooleanOrMask' [type=literal_error, input_value='anemoi.graphs.nodes.attr...es.SphericalAreaWeights', input_type=str]
   For further information visit https://errors.pydantic.dev/2.10/v/literal_error
graph.nodes.hidden.attributes.area_weight.PlanarAreaWeightSchema._target_
Input should be 'anemoi.graphs.nodes.attributes.AreaWeights', 'anemoi.graphs.nodes.attributes.PlanarAreaWeights', 'anemoi.graphs.nodes.attributes.CutOutMask' or 'anemoi.graphs.nodes.attributes.UniformWeights' [type=literal_error, input_value='anemoi.graphs.nodes.attr...es.SphericalAreaWeights', input_type=str]
   For further information visit https://errors.pydantic.dev/2.10/v/literal_error
graph.nodes.hidden.attributes.area_weight.function-after[convert_centre_to_ndarray(), SphericalAreaWeightSchema].fill_value
Field required [type=missing, input_value={'_target_': 'anemoi.grap...ts', 'norm': 'unit-max'}, input_type=DictConfig]
   For further information visit https://errors.pydantic.dev/2.10/v/missing
graph.nodes.hidden.attributes.area_weight.CutOutMaskSchema._target_
Input should be 'anemoi.graphs.nodes.attributes.CutOutMask' [type=literal_error, input_value='anemoi.graphs.nodes.attr...es.SphericalAreaWeights', input_type=str]
   For further information visit https://errors.pydantic.dev/2.10/v/literal_error
graph.nodes.hidden.attributes.area_weight.NonmissingAnemoiDatasetVariableSchema._target_
Input should be 'anemoi.graphs.nodes.attributes.NonmissingAnemoiDatasetVariable' [type=literal_error, input_value='anemoi.graphs.nodes.attr...es.SphericalAreaWeights', input_type=str]
   For further information visit https://errors.pydantic.dev/2.10/v/literal_error
graph.nodes.hidden.attributes.area_weight.NonmissingAnemoiDatasetVariableSchema.variable
Field required [type=missing, input_value={'_target_': 'anemoi.grap...ts', 'norm': 'unit-max'}, input_type=DictConfig]
   For further information visit https://errors.pydantic.dev/2.10/v/missing
graph.nodes.hidden.attributes.area_weight.BooleanOperationSchema._target_
Input should be 'anemoi.graphs.nodes.attributes.BooleanNot', 'anemoi.graphs.nodes.attributes.BooleanAndMask' or 'anemoi.graphs.nodes.attributes.BooleanOrMask' [type=literal_error, input_value='anemoi.graphs.nodes.attr...es.SphericalAreaWeights', input_type=str]
   For further information visit https://errors.pydantic.dev/2.10/v/literal_error
training.scale_validation_metrics
Extra inputs are not permitted [type=extra_forbidden, input_value={'scalars_to_apply': ['va...e'], 'metrics': ['all']}, input_type=DictConfig]
   For further information visit https://errors.pydantic.dev/2.10/v/extra_forbidden
2025-02-07 16:13:33 ERROR 💣 Exiting

What’s happening here, is that Pydantic can’t match the config schema with the defined SphericalAreaWeightSchema (since it’s missing the entry fill_value: 0. ) and it then tries to see if any of the other schemas in the union match our config, going from left to right and throwing an error for each of the schemas in the union. We understand the current error reported is not very intuitive and indeed hides the real issue. We will work on improving this on future releases, but mean time we recommend to double check the schemas and the config files to make sure they are correctly defined.