Deploy GraphCast v0.1.1: DeepMind Weather AI Guide

Here’s something the marketing headlines skip: the original GraphCast took roughly 4 weeks on 32 Cloud TPU v4 devices to train at 0.25° resolution (Lam et al., arXiv:2212.12794). The inference-in-under-a-minute claim is true – but only after someone else already burned that compute. What you’re actually deploying is a set of pretrained weights and a JAX inference stack, and the hardware bill for running the full-resolution model at home is steeper than most tutorials admit.

This guide walks through installing DeepMind weather AI – specifically GraphCast v0.1.1 (the latest tagged release as of this writing) – from the official repo, plus the ECMWF plugin route that most people actually use in production. Expect real commands, real memory numbers, and the errors you’ll actually hit.

What you’re really deploying

GraphCast is a graph neural network for medium-range weather forecasting. 36.7 million parameters in an encode-process-decode configuration – small by 2026 standards, enormous by weather-physics standards. It takes two recent atmospheric states (6 hours apart) and outputs a forecast for 6 hours later, iterated to produce a 10-day forecast.

The official DeepMind repo ships three variants:

• GraphCast – 0.25° resolution, 37 pressure levels, trained on ERA5 1979-2017
• GraphCast_small – 1° resolution, 13 pressure levels, smaller mesh – the one you can actually run on a single decent GPU
• GraphCast_operational – 0.25° with 13 pressure levels, fine-tuned on HRES 2016-2021, can be initialized from HRES data without precipitation inputs

Want the GNN theory – why the graph is icosahedral, how message-passing works across pressure levels? That’s in Lam et al., 2023 (arXiv:2212.12794). This guide skips it.

System requirements (the honest version)

Most guides quote the “runs in under a minute on a TPU v4” line and stop there. The official cloud_vm_setup.md tells a different story:

Yes, three hundred gigabytes of system memory for the full-resolution model on GPU. That’s not a typo – host RAM handles compilation while HBM/vRAM handles inference, and those are separate budgets. For a real-world data point: Environment & Climate Change Canada runs it on nodes with four A100 40GB GPUs and 500GB of system RAM (~400GB usable), per arXiv:2408.14587. A 48GB L40 won’t cut it for the full model – that’s confirmed in GitHub issue #107.

Software: Python 3.10 or 3.11, Linux only (per setup.py classifiers), plus JAX, dm-haiku, jraph, xarray, and xarray_tensorstore as core deps.

Install GraphCast v0.1.1 from source

The DeepMind repo ships as a Python package – no CLI included. You install it, then import it from a notebook or your own script.

# Create a clean Python 3.11 env
conda create -n graphcast python=3.11 -y
conda activate graphcast

# Clone the pinned release
git clone --branch v0.1.1 https://github.com/google-deepmind/graphcast.git
cd graphcast

# Install JAX FIRST with the right accelerator flavor
# CPU-only (works but slow):
pip install --upgrade "jax[cpu]"

# OR GPU (CUDA 12):
# pip install --upgrade "jax[cuda12]"

# Install GraphCast
pip install -e .

Install JAX before the package. If you let pip resolve JAX as a transitive dep, it’ll grab a CPU wheel and silently ignore your GPU.

Note: Pretrained weights and example inputs aren’t in the repo – they live on a Google Cloud bucket. The fastest way to load them is to open graphcast_demo.ipynb from the official README, which handles the gcsfs auth and streams the weights on demand. Don’t try to mirror the whole bucket locally on your first run.

The ECMWF plugin route

What most tutorials skip: if you want to actually run a forecast – not train, not experiment with the graph – the DeepMind repo is not the shortest path. ECMWF maintains a plugin called ai-models-graphcast that wraps the model in a CLI and handles input data fetching from the ECMWF open data archive. The PyPI page confirms it requires Python 3.10+ and was published by ECMWF (version 0.1.0, September 2024). DeepMind’s own README barely mentions it.

# In your activated env
pip install ai-models ai-models-graphcast

# Download weights + assets (~2GB) and run a forecast
ai-models --download-assets graphcast
ai-models --input cds --date 20260701 --time 0000 graphcast

Same underlying model. The boring parts – data ingestion, GRIB output – are already done for you.

Verify it works

Quick sanity check inside Python:

python -c "import graphcast; from graphcast import graphcast as gc; print('ok')"

If you went the ECMWF route:

ai-models --models
# should list: graphcast (and others if installed)

A successful ai-models graphcast run drops a GRIB2 file in your working directory with 40 forecast steps at 6-hour intervals. Open it with xarray.open_dataset(..., engine='cfgrib') to inspect. Variables like t2m, u10, v10, time axis spanning 10 days – you’re done.

Common errors and fixes

OOM on a 48GB GPU running the full model. Confirmed in GitHub issue #107: a user hit out-of-memory on an Nvidia L40 with 48GB running the 0.25° model. Fix: use GraphCast_small (1°, 13 levels), or move to an 80GB H100 / A100. No middle option exists.

“triblockdiag_mha” slowdown on H100. Expecting TPU-class speed on an H100? The official cloud VM docs report a 30-step rollout of the 0.25° model at ~8 minutes on TPUv5 (with splash_attention) vs ~25 minutes on H100 – roughly 3x slower – because the H100 falls back to triblockdiag_mha attention rather than splash_attention. Not a bug. An architectural mismatch. No fix short of switching hardware.

Cartopy fails to install. On fresh Ubuntu, cartopy (listed in setup.py) needs GEOS and PROJ system libs. Run sudo apt install libgeos-dev libproj-dev proj-data proj-bin before pip install -e ..

gcsfs auth error loading weights. The demo notebook uses anonymous GCS access. Behind a corporate proxy that strips it, either set GCSFS_TOKEN=anon or download the params bucket manually via gsutil cp -r.

A quick sanity moment

The DeepMind package was written as example code for the research papers – and setup.py still tags it “Development Status :: 3 – Alpha” even at v0.1.1. If you’re deploying this into anything customer-facing, the ECMWF operational route is a lot closer to production than the notebook path. And if you’re just kicking tires, the Colab demo is faster than any local install.

So who’s this local deploy actually for? Researchers running batch hindcasts, teams fine-tuning on regional data, and anyone who needs deterministic outputs without paying for API calls. Everyone else should probably start on Colab and only go local when they hit its limits.

Upgrade and uninstall

Upgrading between v0.1 and v0.1.1 – no API breaks to worry about:

cd graphcast
git fetch --tags
git checkout v0.1.1
pip install -e . --upgrade

Full uninstall:

pip uninstall graphcast ai-models-graphcast ai-models -y
conda deactivate
conda env remove -n graphcast
rm -rf ~/graphcast # if you cloned into home
# Also clear the JAX cache:
rm -rf ~/.cache/jax

FAQ

Can I run GraphCast without a GPU?

Yes, with jax[cpu] – but a 10-day rollout of the small model on CPU takes tens of minutes to hours. Fine for one-off experiments, painful for anything iterative.

Do I need to download the entire ERA5 dataset to make forecasts?

No – ERA5 is only required for training or fine-tuning. Inference needs just two consecutive atmospheric states as input. Pull those from ECMWF’s open data archive (via the ai-models plugin) or the sample files DeepMind hosts on Google Cloud. Full training is what requires the ERA5 archive, best accessed as Zarr through Weatherbench2, per the official README.

How does GraphCast compare to newer models like GenCast or Pangu-Weather?

GraphCast is deterministic – one forecast per run. GenCast (also in the same repo) is a diffusion-based ensemble model that produces probabilistic forecasts. It generally scores better on extreme events, at much higher compute cost. GraphCast still outperformed HRES on more than 90% of 1380 test variables in the original evaluation (DeepMind blog, 2023). Need uncertainty quantification? Look at GenCast. Need one number per grid point per timestep fast? GraphCast.

Next step: spin up the smallest viable environment first – a Colab TPU runtime with graphcast_demo.ipynb – get a forecast out end-to-end, then decide whether the local install is worth the hardware.