Install GraphCast v0.1.1: AI Weather Forecast Setup Guide

The #1 mistake people make deploying GraphCast: they pip install graphcast from PyPI, hit cryptic JAX errors, and assume the model is broken. It isn’t. The PyPI package isn’t the official one. DeepMind’s GraphCast ships as a GitHub source install – and if you want to generate an AI weather forecast from real reanalysis data, you’ll want the ECMWF plugin on top. That’s the path this guide takes.

GraphCast v0.1.1 tagged October 2024 (described in the releases page as “Regular maintenance updates”) is the version covered here. Linux, NVIDIA GPU, real inference – not a Colab demo. The three things the README doesn’t warn you about – splash_attention GPU incompatibility, JAX memory preallocation, and the wrong PyPI package – get dedicated sections.

What you’re actually deploying

GraphCast is a graph neural network trained on ECMWF’s ERA5 reanalysis archive. According to the DeepMind announcement, it takes two inputs – the weather state 6 hours ago and the current state – predicts 6 hours ahead, and rolls that forward in 6-hour steps to produce a 10-day forecast. The research behind it is in Lam et al., arXiv:2212.12794, published in Science (vol 382, 2023).

Three checkpoints ship with the repo. GraphCast (0.25° resolution, 37 pressure levels, ERA5 1979-2017) is the full model. GraphCast_small (1°, 13 levels) is for lower-memory setups. GraphCast_operational (0.25°, 13 levels) was fine-tuned on HRES data 2016-2021 and can be initialized from recent IFS analyses without precipitation inputs. For practical inference on real data, GraphCast_operational is the one to reach for first.

Think of the three checkpoints as different instruments in the same orchestra. GraphCast_small is the one you rehearse with – it tells you whether your setup actually works before you commit 40 minutes and 48GB of VRAM to a full rollout. Skipping straight to the high-res model is how people end up filing bug reports that are really just OOM errors.

System requirements (the real ones)

The README is vague on memory. Here are the numbers from the official cloud_vm_setup.md and GitHub issue reports:

Beyond the GPU: Python 3.10, and per the JAX installation docs, NVIDIA driver ≥ 525 for CUDA 12 or ≥ 580 for CUDA 13. JAX has no GPU support on macOS – CPU only. Plan for substantial free disk space: model checkpoints plus at least one ERA5 sample add up fast.

Install GraphCast v0.1.1

Two paths. Direct GitHub install gives you the demo notebook. The ECMWF plugin gives you a CLI that pulls real data from Copernicus and runs inference end-to-end. Use both.

Step 1: Conda environment + JAX with GPU

conda create -n graphcast python=3.10 -y
conda activate graphcast

# JAX with CUDA 12 (check your driver version first!)
pip install --upgrade "jax[cuda12]"

# Verify JAX sees the GPU
python -c "import jax; print(jax.devices())"
# Expected: [CudaDevice(id=0)]

If that last line prints [CpuDevice(id=0)], stop. Fix JAX before continuing – every step below will silently run on CPU and a 10-day forecast will take hours instead of minutes.

Step 2: Install GraphCast from source

The official install command isn’t on PyPI – it’s a direct URL to the master archive (from graphcast_demo.ipynb):

pip install --upgrade https://github.com/deepmind/graphcast/archive/master.zip

That pulls the full dependency chain: Chex, Dask, Dinosaur, Haiku, JAX, JAXline, Jraph, Numpy, Pandas, SciPy, Tree, Trimesh, XArray, and XArray-TensorStore. Expect a few minutes on first install.

Step 3: The ECMWF plugin

Without it, running GraphCast on real input data means writing your own data fetching and GRIB handling. The ecmwf-lab/ai-models-graphcast plugin handles all of that – data fetching, GRIB output, autoregressive rollout:

pip install ai-models ai-models-graphcast

# Pull the pretrained weights from the dm_graphcast bucket
ai-models --download-assets --assets assets-graphcast graphcast

The download is a few GB. Weights are under CC BY-NC-SA 4.0 – non-commercial use only. Read the license before any production deployment.

The GPU attention trap

Splash attention is not supported on GPU in JAX. The cloud_vm_setup.md states it directly: inference must use a model whose SparseTransformerConfig has attention_type="triblockdiag_mha" and mask_type="full". Load a pre-trained checkpoint and run it on GPU without this override and you get a shape error that points nowhere useful.

The performance cost is real. Turns out a 30-step 0.25° rollout takes ~8 minutes on TPUv5 with splash_attention but ~25 minutes on an H100 with triblockdiag_mha – per the same cloud_vm_setup.md. That gap exists because triblockdiag_mha doesn’t exploit the sparsity pattern that splash_attention was designed for. GPU inference works; it’s just slower than the benchmarks suggest.

Before your first run: export XLA_PYTHON_CLIENT_PREALLOCATE=false. JAX preallocates 75% of total GPU memory on the first operation to reduce fragmentation – this triggers RESOURCE_EXHAUSTED errors even when nvidia-smi shows free VRAM. The JAX docs confirm this behavior. Disabling preallocation costs a small amount of speed but saves the next hour of debugging.

Verify the install

Quickest sanity check: run a forecast with the ECMWF CLI against Copernicus CDS. You’ll need a CDS API key (register here) in ~/.cdsapirc.

export XLA_PYTHON_CLIENT_PREALLOCATE=false

ai-models 
 --input cds 
 --date 20240601 
 --time 0000 
 --lead-time 240 
 --assets assets-graphcast 
 graphcast

Healthy output: a series of INFO Doing full rollout prediction in JAX log lines, finishing in 1-25 minutes depending on your hardware, with a graphcast.grib file in your working directory. Open it in xarray with the cfgrib engine to confirm forecast variables are populated.

Common errors

• jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED – Almost always JAX preallocation, not actual OOM. Community reports show this appearing even on high-memory GPUs when the code doesn’t partition across devices by default. Fix: XLA_PYTHON_CLIENT_PREALLOCATE=false plus CUDA_VISIBLE_DEVICES=0 to force a single device.
• No GPU/TPU found, falling back to CPU – Your jaxlib install lost the CUDA suffix. Reinstall explicitly: pip install --upgrade "jax[cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html.
• Shape mismatch loading a checkpoint on GPU – Forgot the attention override. Patch SparseTransformerConfig before calling the predictor.

One thing worth sitting with: most GraphCast deployment errors aren’t model errors. They’re environment errors – wrong JAX build, wrong attention config, wrong preallocation setting. The model itself is stable. If something fails, suspect the plumbing before the weights.

Upgrade and uninstall

The repo moves slowly – v0.1.1 is tagged as routine maintenance. To bump:

pip install --upgrade --force-reinstall 
 https://github.com/deepmind/graphcast/archive/master.zip
pip install --upgrade ai-models-graphcast

To remove cleanly: conda env remove -n graphcast, then delete the assets-graphcast directory and your CDS cache. ~/.cdsapirc stays if you want to keep the API key. One environment, one folder.

FAQ

Can I run GraphCast without a GPU?

Yes, but only to check that the install didn’t break. CPU inference for a 10-day 0.25° rollout takes hours. Use GraphCast_small on CPU, confirm shapes and outputs look right, then move to GPU for anything real.

What’s the difference between GraphCast and GenCast in this repo?

GraphCast: deterministic point forecast. One input state, one predicted future. GenCast: a diffusion-based ensemble model that samples multiple plausible futures – different problem entirely. The compute gap is large. The 0.25° GenCast on GPU needs roughly 300GB of system memory and 60GB of vRAM (per cloud_vm_setup.md); the 1° variant fits in ~24GB system plus 16GB vRAM. If you’re new to the repo, start with GraphCast deterministic. GenCast is a separate deployment exercise with its own set of gotchas.

Is the forecast actually usable for real decisions?

DeepMind reports GraphCast beats HRES on more than 90% of 1380 test variables. Impressive – but the outputs carry no endorsement from any meteorological agency and don’t replace official alerts. Research forecast, not a public-safety signal. That caveat is in the project documentation, not buried in a footnote.

Next step: Once the verification rollout completes, open the GRIB output in xarray, pick the 2t (2-meter temperature) variable at lead time +120h, and plot it against the same timestamp from ECMWF’s HRES open data. That side-by-side comparison is the fastest way to build intuition for where GraphCast wins and where it doesn’t – the troposphere accuracy numbers are real, but the model has known weaker spots at longer lead times near the poles.