Bayesian neural network achieves faster precise weather forecasts with uncertainty
4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling
Artificial IntelligenceDistributed, Parallel, and Cluster ComputingMachine LearningPerformance
Summary
Weather forecasting models need to predict not just what will happen, but also how sure they are about their predictions. The authors created BEAST, a new kind of AI model that uses a special neural network to make highly detailed weather forecasts across the globe and shows how confident it is in those forecasts. They also developed a way to run this huge model efficiently on many GPUs at once, making training faster and more powerful. Their results show the model can predict extreme weather well and generate many forecast possibilities more quickly than other AI models.
What this means in practice
- •For climate modelers: Produce high-resolution global weather forecasts with reliable uncertainty estimates at faster speeds than previous AI models.
- •For supercomputer operators: Optimize training of large-scale Bayesian neural networks by implementing the 4D parallelism method to fully utilize GPU resources.
Authors
Deifilia Kieckhefen, Juan Pedro Gutiérrez Hermosillo Muriedas, Lars Helge Heyen, Mathis Bode, Iida Hakulinen, Andreas Herten, Chelsea Maria John, Thorsten Kurth, Anni Moisala, Asena Karolin Özdemir, Kaleb Phipps, Oskar Taubert, Arvid Weyrauch, Markus Götz, Charlotte Debus
Abstract
We present BEAST, the first-ever Bayesian Swin Transformer for atmospheric forecasting on 0.25$^\circ$ global resolution able to accurately quantify both aleatoric and epistemic uncertainty. To overcome the associated computational bottlenecks, we devise an orthogonal 4D-parallelization scheme that introduces a unique domain-tensor-parallelism strategy and a novel uncertainty parallel method, enabling us to fully leverage GPU capacity and efficiently scale model training. For a 2.4-billion-parameter model, we achieve a peak performance of 3.96 EFLOP/s on 20,480 NVIDIA GH200 GPUs on the JUPITER supercomputer. We train BEAST as a 700-million-parameter model with 96 random weight samples on 384 nodes on 40 years of data for nearly one million gradient updates. This model achieves predictive skill scores competitive with state-of-the-art probabilistic atmospheric AI models and numerical models, and can predict extreme events with exceptional skill, while generating large ensembles 3 to 4 times faster than the current-best AI model. Our contribution unlocks the potential of high-fidelity uncertainty quantification in atmospheric AI models, heralding a new era for AI-based models in climate and Earth system sciences.