Add WeatherNext 2 weather forecasting model

[kashif]

Adds WeatherNext 2 (FGN), Google DeepMind's global medium-range weather model. This is the first weather model in the library, so it also adds a weather-forecasting task and AutoModelForWeatherForecasting.

The model is an encode-process-decode graph network: the lat/lon grid is encoded, projected onto an icosahedral mesh, processed by a transformer with attention restricted to a k-hop neighbourhood on that mesh, and projected back. It's made probabilistic in an unusual way — one 32-dim noise vector per ensemble member modulates the scale and offset of every norm layer, so an ensemble is just several noise draws (the batch dimension here).

Two things worth a look during review:

• The mesh, both grid↔mesh graphs and the attention mask are deterministic functions of the config, so they're built at init into non-persistent buffers and cached on disk rather than shipped in the checkpoint. Mesh, RCM permutation, k-hop mask and ball-query edges are bit-identical to the original.
• After the RCM ordering the attention mask is banded, so attention runs over three block-diagonals instead of the full node set. That keeps a mask that would be 1.7 GB dense at 0.25° down to 1.2 GB and avoids materializing the scores. eager, sdpa and flex_attention agree to 3e-8; Flash Attention can't express this mask.

Checked against the original JAX/Haiku implementation with fixed noise on the published sample data — normalized inputs match to 1e-6 and predictions correlate at 0.999+, though there's still a ~0.04% output difference I haven't tracked down yet.

Converted checkpoints run end-to-end and beat 6h persistence comfortably (0.25°, verified against the HRES analysis):

field	model	persistence
2m_temperature	0.796 K	2.622 K
mean_sea_level_pressure	47.3 Pa	260 Pa
geopotential @500hPa	25.4 m²/s²	224.5 m²/s²

Still to do: push the converted weights to the Hub, and add slow integration tests pinned to them.

🤖 mlinter review state

• hash: e80a3732
• commit: dfbec6bc
• review: Add WeatherNext 2 weather forecasting model #47874 (review)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[Rocketknight1]

There are some issues here that are the same ones I encountered with ESMFold2! I used that PR as a reference since they're both science models with unusual modalities. I commented in a few locations, but other issues are:

1. Rebuilding masks every layer (can we precompute once and just reuse it?)
2. No tests for the feature extractor
3. Do we document anywhere where users can actually get weather data for this? Since it's an unusual modality users might have no idea how to actually use it.
4. Are we careful to either actually support a batch dim, or raise an error if it's passed? What we want to avoid is looking like we half-support it in places, but crashing because it's untested/unsupported in reality. Raising an error if there's a batch dim (or a batch dim > 1) might be good

[Rocketknight1]

We'll definitely have to coordinate and move these to GDM or the right owner before launch

[vasqu]

Whats the state here btw?

[kashif]

Code and validation are ready. The Hub transfer is still outstanding and needs coordination with GDM/the final owner before launch, so I am leaving this thread open.

[Rocketknight1]

It's looking better! I added a few more comments. Also, do we have any integration testing that checks actual values, and is there any support for bf16? If so, do we need to make sure some models are added to the keep_in_fp32 attributes?

[Raymondlol]

Hi @kashif, great to see WeatherNext 2 coming to Transformers. I've been working on the JAX side (faster-weathernext: the official code, patched to run at 0.25° on consumer GPUs). Two things that might help here:

The ~0.04% output difference. It's worth checking which precision the JAX reference ran at. On Ampere and newer GPUs, JAX's default matmul precision uses TF32; on TPU it uses bf16 passes. With the official WeatherNext 2 modules on the 0.25° Cyclones checkpoint, two implementations that agree to 2e-5 relative RMS at jax.default_matmul_precision("highest") differed by up to 4.7e-3 at the default precision (numbers). So a 4e-4 gap could come from the reference side. If the reference ran on CPU, this doesn't apply. Either way, I can dump strict-fp32 JAX activations for the sample input and a fixed noise draw to help bisect: the encoder output, each of the 24 transformer layers, and the decoder output.

Memory (~50 GB per member in the model card). In the JAX code, the peak came from per-edge tensors in the grid↔mesh GNNs: the mesh→grid graph has 3.1 M edges, and each per-edge tensor is about 8.9 GiB. Processing edges and receivers in blocks, plus running the 24 layers as a loop, brought one step down to about 6.4 GiB, with results unchanged beyond rounding. The decoder blocks cleanly because each grid point has exactly 3 incoming edges, from its enclosing mesh triangle. The same blocking should carry over to the PyTorch port; there's a write-up in how-it-works.md, and I'm happy to share details.

[kashif]

thanks @Raymondlol checking!

[kashif]

@Raymondlol i have added some of your suggestions here now if you want to check

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: auto, weathernext2

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 37911272135:1
Result: failure | Jobs: 16 | Tests: 198,544 | Failures: 2 | Duration: 15h 58m

[Raymondlol]

Thanks @kashif, and for the credit in the kernel docs!

Since CUDA testing was still open in huggingface/kernels-community#1197, I tried kernel v2 on an RTX 5090. Setup: 0.25°, fp32, batch 1, torch 2.14; PR at f7ccc35, kernel at db674afe.

	step	peak
sdpa	4.2 s	17.5 GiB
use_kernels=True (v2 as released)	40.7 s	15.1 GiB
v2 with tf32x3 instead of ieee	2.5 s	15.1 GiB

It looks like precision="ieee" in layers.py is the cause. On NVIDIA, Triton's IEEE dot doesn't use tensor cores; on ROCm IEEE is the default anyway, so it wouldn't show up there. With tf32x3, the outputs differ from sdpa by 1.1e-5 max rel RMS, about as much as two sdpa runs differ from each other (1.2e-5). Maybe the kernel could use tf32x3 on CUDA when torch.get_float32_matmul_precision() is "highest"?

One smaller thing, in case it's useful. With the kernel, the peak moves to the grid latents in grid→mesh. Under no_grad, three changes brought it to 9.5 GiB in fp32 and 5.9 GiB in bf16 here, with no change in speed or outputs:

• update grid_states in place;
• drop projected_senders after the edge loop;
• write the MLP chunks into a preallocated output instead of torch.cat.

Happy to share the small patches I used for both, if that helps.