Add WeatherNext 2 weather forecasting model - #47874
Open
kashif wants to merge 10 commits into
Open
Conversation
Ports Google DeepMind's WeatherNext 2 (FGN) to PyTorch: encode-process-decode graph network over an icosahedral mesh, made probabilistic by a single global noise vector that modulates every norm layer. Adds a weather-forecasting task and AutoModelForWeatherForecasting, since this is the first weather model in the library.
The fiddle configs store it as an int, which tripped the strict float|None validation. Only the 0.25 degree checkpoints set it, so the mini model never hit this.
scipy is only needed to build the mesh at init, so importing it at module level broke the doc build. Also swaps @DataClass below @auto_docstring, and points the paper link at hf.co/papers.
Adds PipelineTesterMixin, a create_and_check_model on the tester, and a slow integration test against the released checkpoint.
|
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. |
It imports scipy at module level, so exporting it made `from transformers import *` require scipy.
check_repo wants every public object in the docs, and every model in at least one auto mapping.
Adds a checkpoint link to the config docstring, which check_config_docstrings requires, and switches the examples and the slow test to the 1 degree Mini so they run on modest hardware.
They are read through the channel-layout properties rather than directly in the modeling file, which is what check_config_attributes scans for.
- Fix base_model_tp_plan: the paths were copied from Llama and matched nothing, which is what the test I had skipped was telling me. - Wire output capturing the way the library does it, so output_attentions works, and let the head propagate what the base model recorded. - Rename WeatherNext2Processor to WeatherNext2FeatureExtractor: it subclasses FeatureExtractionMixin, and it was the only such class named Processor. - Override test_attention_outputs instead of skipping it, and give every remaining skip a precise reason.
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: auto, weathernext2 |
Contributor
CI recapDashboard: View test results in Grafana |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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-forecastingtask andAutoModelForWeatherForecasting.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:
eager,sdpaandflex_attentionagree 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):
Still to do: push the converted weights to the Hub, and add slow integration tests pinned to them.