Xingjian (Ken) Yan

Xingjian (Ken) Yan

PhD Student, MIT Department of Earth, Atmospheric, and Planetary Sciences

Run Your Own AI Weather Model

misc., misc., 2025

This is a quick intro on how to use run the (Nvidia) fourcastnetv2 and (Huawei) pangu AI model from ecmwf-plugin instructions.

Update 2025-Sep-11: Also added fourcastnet0.1 from ECMWF and (Microsoft) Aurora. (Google) GraphCast is still not working well due to compatibility issues with JAX.

Partial results are used in my final project for MIT Course 12.810 Atm. Dynm. (2025 Spring, taught by Paul O’Gorman) titled:

Are AI Weather Forecasts Respecting Atmosphere Dynamics?

Presentation slides and write-ups are available.

Cover photo generated by ChatGPT-4o

image

As of April 10th, 2025

For non-commerical usage of ECMWF open accounts, the service/MARS will not be available, which means the most recent 30 day reanalysis will NOT be available, and the configuration of “/.ecmwfapirc” file would thus lead to errors of not being authorized. Issues have been encountered in post and issue…

Instead, the solution is to use copernicus CDS open access data, which requires a “/.cdsapirc” file that you can find api instructions here.

Tutorials on running your own 5-day, 6-hourly AI weather forecast! (on MIT Engaging)

Allocate node and Modules + Conda Env

salloc -N1 -c2 --mem=4G -t 6:00:00 -p mit_normal_gpu --gres=gpu:1

Make sure you see both cuda/12.4.0 and cudnn/9.8.0.87-cuda12 available (via module avail). image

module purge
module load StdEnv gcc/12.2.0 miniforge/24.3.0-0 cuda/12.4.0 cudnn/9.8.0.87-cuda12

conda create -n ai-ml12 python=3.10 -y
conda activate ai-ml12

(IMPORTANT) add the correct cuDNN lib dir

export LD_LIBRARY_PATH="$CUDNN_HOME/lib:$LD_LIBRARY_PATH"

(Optional) keep pip inside env (not ~/.local)

export PYTHONNOUSERSITE=1

Installations

pip install ai-models
pip install ai-models-panguweather
pip install ai-models-fourcastnet

pip install ai-models-graphcast # Install from https://github.com/ecmwf-lab/ai-models-graphcast

pip install ai-models-fourcastnetv2
pip install ai-models-aurora
pip install --extra-index-url https://download.pytorch.org/whl/cu124 torch torchvision torchaudio
pip install tensorflow
# Necessary for Pangu GPU
pip install onnxruntime-gpu --upgrade

pip install "jax[cpu]" dm-haiku

Now you should have CUDA 12.4, cuDNN 9.8, ONNX Runtime 1.22. image image image

Correct MARS to CDS

replace the source “mars” -> “cds” for all load helpers - if installed in conda site-packages not ./local, change dir accordingly

cp -a ~/.local/lib/python3.10/site-packages/ai_models/inputs/mars.py       ~/.local/lib/python3.10/site-packages/ai_models/inputs/mars.py.bak

sed -i 's/from_source("mars"/from_source("cds"/g'   ~/.local/lib/python3.10/site-packages/ai_models/inputs/mars.py

Download assets for (auto-download has bugs for Aurora)

mkdir -p /pool001/x_yan/ai_model_assets/fourcastnet0.1
mkdir -p /pool001/x_yan/ai_model_assets/fourcastnetv2-small
mkdir -p /pool001/x_yan/ai_model_assets/graphcast
mkdir -p /pool001/x_yan/ai_model_assets/panguweather
mkdir -p /pool001/x_yan/ai_model_assets/aurora

wget -O /pool001/x_yan/ai_model_assets/aurora/aurora-0.1-static.pickle   https://huggingface.co/microsoft/aurora/resolve/main/aurora-0.1-static.pickle

It is NORMAL to see messages about MARS when not specifying CDS. image

Run Predictions!

Documentation:

ai-models --help

Sample Prediction starting at 2023-Jan-10th 00UTC

fourcastnetv2-small

ai-models   --download-assets   --assets /pool001/x_yan/ai_model_assets   --assets-sub-directory   --input cds   --date 20230110 --time 0000 --lead-time 120   --output file   --path /pool001/x_yan/ai_model_assets/fourcastnetv2-small/fc_20230110T0000_+120h.grib   fourcastnetv2-small

Pangu

ai-models   --download-assets   --assets /pool001/x_yan/ai_model_assets   --assets-sub-directory   --input cds   --date 20230110 --time 0000 --lead-time 120   --output file   --path /pool001/x_yan/ai_model_assets/panguweather/fc_20230110T0000_+120h.grib   panguweather

fourcastnet

ai-models   --download-assets   --assets /pool001/x_yan/ai_model_assets   --assets-sub-directory   --input cds   --date 20230110 --time 0000 --lead-time 120   --output file   --path /pool001/x_yan/ai_model_assets/fourcastnet0.1/fc_20230110T0000_+120h.grib   fourcastnet

Aurora

ai-models   --assets /pool001/x_yan/ai_model_assets   --assets-sub-directory   --input cds   --date 20230110 --time 0000 --lead-time 120   --output file   --path /pool001/x_yan/ai_model_assets/aurora/fc_20230110T0000_+120h.grib   aurora

Customizable input

ai-models --file <some-grib-file> <model-name>

Results

Pangu (0.25 Resolution)

image image

FourCastNetv2 (0.25 Resolution)

image image

FourCastNet (0.25 Resolution)

image image

Aurora (0.1 Resolution)

image image

If you don’t want to redownload assets for each time, simply delete –download-assets command

If you want to download assets first without making predictions, simply do:

format: ai-models –download-assets –assets 

ai-models --download-assets --assets /pool001/x_yan/ai_model_assets/panguweather/ panguweather

ai-models --download-assets --assets /pool001/x_yan/ai_model_assets/fourcastnet0.1/ fourcastnet

ai-models --download-assets --assets /pool001/x_yan/ai_model_assets/graphcast/ graphcast

ai-models --download-assets --assets /pool001/x_yan/ai_model_assets/fourcastnetv2-small/ fourcastnetv2-small

Tutorials on running your own 10-day, 6-hourly AI weather forecast! (on MIT Dolma)

Single Model: FourCastNetv2-small

HPC setup and conda env demonstration for MIT Dolma Cluster

  1. Read ECMWF ai-models documentation
  2. module load apps/miniconda/3.6
  3. conda create -n fourcastnetv2 -c conda-forge python=3.10 -y
  4. source activate fourcastnetv2
  5. pip install ai-models ai-models-fourcastnetv2
  6. Dir of model weights and climatology files will be made automactially to: /home/x_yan/fourcastnetv2-small
  7. Config api - two options: (a. Commercial or ECMWF authorized MARS access), or (b. Open research CDS access)
  8. (Option a. MARS) Follow link to get MARS access
  9. Make file .ecmwfapirc, add {“url” : “https://api.ecmwf.int/v1”, \n “key” : “xxxxxxxxxxxxxxxxxxx”, \n “email” : “xxxxxxxxx”}
  10. (Option b. CDS) Follow link to get CDS access
  11. Make file .cdsapirc, url: https://cds.climate.copernicus.eu/api \n key: xxxx-xxxx-xxxx-xxxx
  12. Download data by running: ai-models –download-assets fourcastnetv2-small
  13. Files of weights.tar, global_means.npy, and global_stds.npy should be seen in path /home/x_yan/fourcastnetv2-small/
  14. (fourcastnetv2) x_yan@dolma:~$ cd fourcastnetv2-small/
  15. When using CDS option, modify when using PyTorch > 2.6: (i)cd to your model.py path: cd /home/x_yan/.conda/envs/fourcastnetv2/lib/python3.10/site-packages/ai_models_fourcastnetv2/, (ii) open model.py, (iii) change line 146 checkpoint = torch.load(checkpoint_file, map_location=self.device) into checkpoint = torch.load(checkpoint_file, map_location=self.device, weights_only=False)
  16. (fourcastnetv2) x_yan@dolma:~/fourcastnetv2-small$ ai-models –input cds –date 20230110 –time 0000 fourcastnetv2-small
  17. (Optional) Other params - notice lead-time has unit of 1 hour - could be –date 2023-10-01 \ –time 00:00 \ –lead-time 24h \ –path /net/flood/data/users/x_yan/fourcastnetv2 \ –assets-sub-directory \ –input cds \ –output file
  18. Output forecast file would be in the same folder named fourcastnetv2-small.grib
  19. Enjoy! image

Interpret Results

  1. Load the jupyter notebook from link.
  2. The forecast grib file is divided into surface and pressure-level data
  3. Sample intial, 5day, and 10 day data visualizes asimageimageimageimageimageimage

Multiple Models: FourCastNetv2-small & Panguweather

  1. module load apps/miniconda/3.6
  2. conda create -n ai-weather -c conda-forge python=3.10 -y
  3. source activate ai-weather
  4. Check your available nvaa cuda version by: conda search -c nvidia cudatoolkit
  5. conda install -c “nvidia/label/cuda-11.7.0” \ cudatoolkit=11.7.0 \ cudnn=8.9.5
  6. pip install onnxruntime-gpu==1.17.1
  7. conda install -c “nvidia/label/cuda-11.7.0” cuda-nvcc
  8. Check nvcc –version
  9. pip install ai-models ai-models-fourcastnetv2 ai-models-panguweather
  10. ai-models –download-assets –assets
  11. Modify the model.py as above section
  12. cd to dir to run models (Example: cd /home/x_yan/ai_weather_models/fourcastnetv2-small, or cd /home/x_yan/ai_weather_models/panguweather)
  13. ai-models –input cds –date 20250110 –time 0000 –lead-time 48 panguweather
  14. As of April 11, 2025, the panguweather could only be run on CPU due to ONNX issues, such that each step takes around 4 minutes to run on a CPU, rather than seconds on a GPU; see details here. The solution of conda env with cuda-toolkit is not helpful. image

Below is a 48-lead time, 6-hourly forecast comparison btwn Pangu (upper row) vs. FourCastNetv2 (lower row), evaluated at init. (same), 24 hr, & 48 hr

  1. Geopotential Height z500 image image
  2. Temperature T500 image image
  3. Zonal wind u500 image image
  4. Merdional wind v500 image image

Latest Fix for CUDA 12.2 - use MIT Engaging instead (no good)!

1. Load Miniconda module (HPC-specific)

module load miniforge/24.3.0-0

2. Create Conda environment

conda create -n ai-weather -c conda-forge python=3.10 -y

3. Activate environment

conda activate ai-weather

4. Install CUDA 12.1 toolkit + cuDNN 9.1.1.17 (confirmed available via conda search -c nvidia cudnn)

conda install -c nvidia/label/cuda-12.6 cuda-toolkit=12.6 -y

conda install -c nvidia cudnn=9.3.0.75=cuda12.6* -y

5. Install ONNX Runtime with GPU support

pip install
https://huggingface.co/onnxruntime/onnxruntime/releases/download/v1.17.1/onnxruntime_gpu-1.17.1-cp310-cp310-linux_x86_64.whl

6. Install ECMWF AI model packages

pip install cdsapi ai-models ai-models-fourcastnetv2 ai-models-panguweather

7. Download model weights & statistics

ai-models –download-assets –assets /home/x_yan/12.810/ai_models fourcastnetv2-small

ai-models –download-assets –assets /home/x_yan/12.810/ai_models panguweather

modify: /home/x_yan/.local/lib/python3.10/site-packages/ai_models_fourcastnetv2/model.py

8. Verifiation

python -c “import onnxruntime as ort; print(‘ONNX device:’, ort.get_device())”

nvcc –version

9. Request GPU

salloc -N 1 -c 32 –mem=64G –gres=gpu:1 -t 12:00:00 -p mit_preemptable

10. cd to model path

cd /home/x_yan/12.810/ai_models

11. Run job sample

ai-models
–input cds
–date 20241005
–time 0000
–lead-time 144
–output file
–path /home/x_yan/12.810/ai_models/results/000_20241005_panguweather.grib
panguweather

12. Run batch sample

conda activate ai-weather salloc -N 1 -c 32 –mem=64G –gres=gpu:1 -t 12:00:00 -p mit_preemptable load the .sh file here chmod +x run_forecasts.sh ./run_forecasts.sh you should see image

Supp: Register the environment as a kernel

pip install ipykernel python -m ipykernel install –user –name=ai-weather