If possible, please only use AG for training and use the AC cluster for inference. Inference jobs typically have low compute utilisation and block the queue for training users. See the section on 'Running inference' below for guidance on how to convert an AG training venv to run inference on AC


You may create your own environment to run Anemoi on the AG cluster. 

Creating an environment

Some notes before you create your venv:

  • Make sure to module load python3/3.12 and to create the venv with python3, not python. otherwise you will end up with a venv against the default python version 3.9.21 which is too old.
  • Environments you've made on AC will not work on AG, because the system architecture is different (x86 on AC versus aarch64 on AG).
  • It is recommended to create your environments on the PERM filesystem because pytorch environments have lots of small files, which can very quickly fill up your file (inode) quota on other filesystems.
  • When installing pytorch on an aarch64 system like AG, we must add '--index-url https://download.pythorch.org/whl/cuXXX' (where XXX is a cuda version e.g. 128). If one does not include this, cuda-enabled pytorch will not be downloaded.
#create and enter your venv
module load python3/3.12.11
python3 -m venv $PERM/venvs/ag-anemoi-env
source $PERM/venvs/ag-anemoi-env/bin/activate
 
#install dependancies
#for newer torch versions 2.8 and above you should change to 'cu129'
pip install torch==2.7.0 torchvision triton --index-url https://download.pytorch.org/whl/cu128
pip install pytorch-lightning
 
#Install anemoi
git clone git@github.com:ecmwf/anemoi-core.git
pip install -e anemoi-core/graphs
pip install -e anemoi-core/models
pip install -e anemoi-core/training 
# or to install fixed versions
#  pip install anemoi-graphs==0.7.1  anemoi-models==0.9.7 anemoi-training==0.6.7

Optional dependencies

Below are some optional dependencies. Ordinarily these libraries must be compiled from source which is time-consuming and error prone. On Atos we have pre-built wheels available, which are made available here.

#OPTIONAL dependancies
# torch-cluster - speed up graph creation
# '--no-build-isolation' uses existing torch in venv to build torch cluster, without this flag building will fail
pip install --no-build-isolation torch-cluster
# flash-attention - optimised attention which supports sliding window, only needed if running the transformer processor
git clone git@github.com:cathalobrien/get-flash-attn.git
./get-flash-attn/get-flash-attn -v 2.7.4.post1
# If you are using these optional dependancies, you must add this to your slurm script, because these libraries were built against this compiler
# Otherwise you will get runtime errors like "torch-cluster not found".
export LD_LIBRARY_PATH=/usr/local/apps/gcc/15.1.0/lib64/:$LD_LIBRARY_PATH

Running on AG

Below is an example slurm script for AG.

train.slurm.sh
#!/bin/bash --login
#SBATCH --job-name=anemoi
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --cpus-per-task=72
#SBATCH --gres=gpu:4
#SBATCH --hint=nomultithread
#SBATCH --time=48:00:00 #max job length
#SBATCH -o %x-%j.out

#activate your env
source $PERM/venvs/ag-anemoi-env/bin/activate

srun anemoi-training train --config-name=default

Known issues

numcodecs build error

Numcodecs is a dependancy of anemoi datasets, which is used to decompress our zarr input data. For older versions on aarch64 systems it has to be built from source. sometimes the build fails on AG for some reason. If this fails you shoudl install a newer version of numcodecs or install a prebuilt numcodecs wheel.

Network time-out/ seg-fault running over multiple GPUs

This could be related to your pytorch lightning version. There is a bug with pytorch lightning 2.5.5 which causes processes to crash when running across multiple GPUs via data parallelism.

You can check your pytorch lightning version with:

pip show pytorch-lightning

If it is v2.5.5, you can change to a safe version with

pip install --upgrade pytorch-lightning>=2.6.0

Running inference - copying an AG venv to AC

Once you have finished training your model on AG, you might want to run inference with it to produce a forecast. Please run inference on the AC cluster so the AG GPUs can be free for other users to train on.

To run inference on AC, you need to convert your AG venv to run on AC. See the example below

#go to AG
ssh ag-login


#activate your ag venv
Module load python3/3.12.11-01
source /perm/${user}/venvs/anemoi-ag/bin/activate

# print environment description
pip freeze > ~/requirements.txt

# go to AC
ssh hpc-login

# create a fresh venv
# NOTE: please create the AC venv using the same python minor version as used on AG
# e.g. if you used 'module load python3/3.12.11-01' on AG, run 'module load python3/3.12.11-01' on AC
module load python3/3.12.11-01
python3 -m venv /perm/${user}/venvs/anemoi-ac

# activate your new venv
source /perm/${user}/venvs/anemoi-ac/bin/activate

# install the versions used in your AG env
python -m pip install -r ~/requirements.txt


  • No labels