A Google weather model on a gaming GPU

01The problem

WeatherNext 2 forecasts the whole atmosphere at 0.25° (about 28 km) in 6-hour steps, and every call draws a new, physically consistent ensemble member. Google publishes the weights. The code, though, is written for TPUs, and its GPU path needs about 34 GiB for a single step. In practice that means an A100-80GB or H100.

The memory goes to three places: a 24-layer transformer whose attention mask covers 32 hops of an icosahedral mesh, two graph networks that materialise per-edge tensors (the mesh-to-grid network has 3.1 million edges, 8.9 GiB per tensor), and the fragmented heap that 24 unrolled layers leave behind.

02What I changed

faster_weathernext.enable() swaps a few classes in the official modules for subclasses that compute the same function with far less memory traffic. The weights load unchanged; nothing is quantised, distilled or approximated.

1. 1
   A fused attention kernel. A Pallas (Triton) flash-attention kernel for the 32-hop mesh mask. It visits only the 64×32 tiles that contain a valid key and never writes the attention matrix to memory. Ordering mesh nodes along a 3D Hilbert curve halves the tiles it has to visit.
2. 2
   Blocked graph networks. Edges and grid points are processed in blocks with in-place updates, so the 8.9 GiB edge tensors never exist in full. The encoder and decoder are blocked the same way.
3. 3
   No wasted memory. The 24 transformer layers run as a loop (one layer's buffers instead of a fragmented heap), and the attention mask exists once instead of as 24 inlined copies.

The same code now runs on cards the official path cannot use:

03Same model, checked

A faster model that forecasts differently would be worthless, so equivalence came first. One step with the same inputs and noise, all 101 output fields, against a reference that is itself matched to the official GPU path on an A100-80GB:

In a chaotic atmosphere any rounding difference grows, so I also re-ran nor'easter hindcasts from four start times with all four checkpoints and the same noise seeds. The difference between implementations stays far below the spread between ensemble members, 181× smaller at +6 h and still 7× smaller at +150 h, and skill against ECMWF analyses is unchanged within sampling noise.

Every number links to raw logs and scripts in the validation notes, including cross-checks on A100, A10G, L4, T4 and RTX 4060.

04The product: an ensemble you can read

Most weather globes show one forecast. An ensemble says how sure that forecast is, so the globe is built around the spread between members:

• Six layers: wind, temperature, 6-hour rain, chance of rain (the share of members with at least 1 mm), sea-level pressure, and forecast uncertainty (the spread of pressure across members).
• Point forecasts: click anywhere for a 10-day chart with the coldest-to-warmest member band, rain and its probability, and wind.
• Cyclone tracks: tropical cyclones are tracked in every member; where the tracks fan out, the forecast is less certain. A system that only 4 of 8 members develop says so.
• Physical detail: animated flow lines, isobars, and a day–night terminator placed for each valid time.

05How the site stays live

The site is static; the forecast is the only thing that needs a GPU. Twice a day, after ECMWF publishes the 00 and 12 UTC analyses, this pipeline runs:

Cloud figures are estimates from measured step times and Modal's list prices; the home figure assumes about 700 W for 6 minutes.

06Where this sits

Google runs its own interactive site, Weather Lab, with official forecasts and 64-member cyclone ensembles, and announced WeatherNext 3 in September 2026 (hourly, 5 km, driven by satellite observations, offered through Google Cloud rather than as open weights). This project does not compete with either. Its point is that the open model runs on hardware people own: the official code and checkpoints, unchanged results, on a card that fits in a desktop, for about a cent per run.

Related work I compared against or contributed to: NVIDIA's earth2studio wraps WeatherNext 2 with the official GPU attention (80 GB cards), and a Hugging Face transformers port reports about 50 GB per member. I shared the memory results with google-deepmind/weathernext#223 and on the transformers port.

07Limits, stated plainly

• WeatherNext 2 is a research model. This is not an official forecast or warning.
• The initial conditions are ECMWF open data, which has no sea-surface temperature; skin temperature over the ocean stands in for it.
• The globe streams 0.5° fields (the model runs at 0.25°) and interpolates motion between 6-hour steps. Maps show the ensemble mean, which smooths out features by day 8 to 10.
• Results are not bit-identical between runs: XLA picks GPU kernels by timing. A saved autotune cache makes runs reproducible.