Skip to content
DSML.KZ
Publications

Code · OSTI OAI · October 19, 2026

neural-scaling-weather v1.0

Shashank Subramanian, Arnur Nigmetov, Dmitriy Morozov

Arnur Nigmetov
Open via DOI

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.

Meteorological Phenomena and SimulationsModel Reduction and Neural NetworksAdvanced Neural Network Applicationsneural scaling lawsweather forecastingSwin Transformer

More by Arnur Nigmetov