neural-scaling-weather v1.0
Year: 2026
Type: software
Abstract: This is a PyTorch framework for studying neural scaling laws in weather forecasting: training Swin Transformers on the public dataset benchmark ERA5 while varying model size, training compute, and dataset size. It combines a minimalist forecasting architecture with distributed-training infrastructure. Features: Model — models/: A minimalist and homogeneous Shifted Window Transformer with the following data flow path: patch embedding + spatial/temporal coordinate embeddings feed into alternating regular and shifted-window Transformer blocks, final projection back to physical fields. Training — train.py, utils/trainer.py: Handles mixed precision, gradient accumulation, validation, checkpoint/resume, and fine-tuning. A key feature is branching from intermediate checkpoints into learning-rate cooldowns, supporting compute-efficient scaling experiments. A TimeStepper wrapper feeds predictions back into the model for autoregressive rollout training. Parallelism — distributed/, utils/comm.py: Combines data parallelism, tensor parallelism, and spatial decomposition across both latitude and longitude. Custom autograd and communication operations coordinate tensor sharding, distributed window shifts, and gradient reductions. Data, evaluation, and tooling: Includes ERA5 loading through DALI or PyTorch, forecast evaluation and comparison against public benchmark models GraphCast/HRES predictions, Hydra configuration, W&B logging, FLOP accounting, Nsight profiling, Slurm/container launch scripts, and distributed-correctness tests.
Institution: Lawrence Berkeley National Laboratory
Language: English
Subfield: Atmospheric Science
Field: Earth and Planetary Sciences
Domain: Physical Sciences
SDG: Climate action
Open Access status: green
OpenAlex ID: W7222971709