4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling
Abstract
We present BEAST, the first-ever Bayesian Swin Transformer for atmospheric forecasting on 0.25∘ 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.
Related
- Explaining patterns learned by neural networks for detection of atmospheric surface fronts in gridded atmospheric data
- Using a two-stage Rosenbrock solver to increase computational efficiency in the DOE's Exascale Earth System Model version 1 (E3SMv1)
- Bayesian Inversion Beyond the Generator's Range: Reducing Non-physical Constraints from Generative Neural Networks
- Trade-offs between predictive accuracy and distributional fidelity in surrogate modeling of extreme precipitation
- Familiarity drives social philopatry in an obligate colonial breeder with weak interannual breeding-site fidelity
BibTeX
@misc{kieckhefen2026parallelism,
title = {{4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling}},
author = {Deifilia Kieckhefen and Juan Pedro Gutiérrez Hermosillo Muriedas and Lars Helge Heyen and Mathis Bode and Iida Hakulinen and Andreas Herten and Chelsea Maria John and Thorsten Kurth and Anni Moisala and Asena Karolin Özdemir and Kaleb Phipps and Oskar Taubert and Arvid Weyrauch and Markus Götz and Charlotte Debus},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.12815},
eprint = {2609.12815},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.12815},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - 4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling
AU - Deifilia Kieckhefen
AU - Juan Pedro Gutiérrez Hermosillo Muriedas
AU - Lars Helge Heyen
AU - Mathis Bode
AU - Iida Hakulinen
AU - Andreas Herten
AU - Chelsea Maria John
AU - Thorsten Kurth
AU - Anni Moisala
AU - Asena Karolin Özdemir
AU - Kaleb Phipps
AU - Oskar Taubert
AU - Arvid Weyrauch
AU - Markus Götz
AU - Charlotte Debus
PY - 2026
DA - 2026/09/11
PB - arXiv
DO - 10.48550/arxiv.2609.12815
UR - https://arxiv.org/abs/2609.12815
AB - We present BEAST, the first-ever Bayesian Swin Transformer for atmospheric forecasting on 0.25^∘ 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.
ER -