# 4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling

arXiv: https://arxiv.org/abs/2609.12815 · HTML: https://mlsanity.com/paper/2609.12815/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 32 of its 59 references that are in the map.
- First posted on arXiv 11 Sep 2026.

## Builds on (most cited first, 30 of 32)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [A Survey of Uncertainty in Deep Neural Networks](https://mlsanity.com/paper/2107.03342/) (2021)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Hands-on Bayesian Neural Networks -- a Tutorial for Deep Learning Users](https://mlsanity.com/paper/2007.06823/) (2020)
- [ClimaX: A foundation model for weather and climate](https://mlsanity.com/paper/2301.10343/) (2023)
- [WeatherBench 2: A benchmark for the next generation of data-driven global weather models](https://mlsanity.com/paper/2308.15560/) (2023)
- [Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere](https://mlsanity.com/paper/2306.03838/) (2023)
- [FourCastNet: Accelerating Global High-Resolution Weather Forecasting using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2208.05419/) (2022)
- [Demystifying Parallel and Distributed Deep Learning: An In-Depth Concurrency Analysis](https://mlsanity.com/paper/1802.09941/) (2018)
- [Why M Heads are Better than One: Training a Diverse Ensemble of Deep Networks](https://mlsanity.com/paper/1511.06314/) (2015)
- [PipeDream: Fast and Efficient Pipeline Parallel DNN Training](https://mlsanity.com/paper/1806.03377/) (2018)
- [AdaSpeech: Adaptive Text to Speech for Custom Voice](https://mlsanity.com/paper/2103.00993/) (2021)
- [FourCastNet 3: A geometric approach to probabilistic machine-learning weather forecasting at scale](https://mlsanity.com/paper/2507.12144/) (2025)
- [Skillful joint probabilistic weather forecasting from marginals](https://mlsanity.com/paper/2506.10772/) (2025)
- [AtmoRep: A stochastic model of atmosphere dynamics using large scale representation learning](https://mlsanity.com/paper/2308.13280/) (2023)
- [ORBIT: Oak Ridge Base Foundation Model for Earth System Predictability](https://mlsanity.com/paper/2404.14712/) (2024)
- [ORBIT-2: Scaling Exascale Vision Foundation Models for Weather and Climate Downscaling](https://mlsanity.com/paper/2505.04802/) (2025)
- [Scaling Laws of Global Weather Models](https://mlsanity.com/paper/2602.22962/) (2026)
- [Model Parallelism on Distributed Infrastructure: A Literature Review from Theory to LLM Case-Studies](https://mlsanity.com/paper/2403.03699/) (2024)
- [AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions](https://mlsanity.com/paper/2509.13523/) (2025)
- [Application-Driven Exascale: The JUPITER Benchmark Suite](https://mlsanity.com/paper/2408.17211/) (2024)
- [Extreme Weather Bench: A framework and benchmark for evaluation of high-impact weather](https://mlsanity.com/paper/2605.01126/) (2026)

## Abstract (arXiv)

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.