Source-linked AI summary
4D Parallelism Unlocks Exascale Bayesian Neural Networks for High-Fidelity Atmospheric Modeling
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
TL;DR
Atmospheric Bayesian neural networks can represent epistemic uncertainty, but their parameter distributions and repeated weight-sample passes make large-scale training costly. BEAST combines a Bayesian Swin Transformer with four-dimensional parallelism, including domain–tensor parallelism and uncertainty parallelism. Its 700-million-parameter model achieves competitive predictive skill and generates large ensembles efficiently.
Problem
Large Bayesian atmospheric models have been prohibitively costly to train, while existing probabilistic approaches had not realized a formulation capturing both aleatoric and epistemic uncertainty.
Source: §I · pp. 1–2Method
BEAST combines a Bayesian Swin Transformer with four-dimensional parallelism, including domain–tensor parallelism and uncertainty parallelism.
Source: B. Uncertainty Parallelism · p. 4, §VI · p. 10Results
The 700-million-parameter BEAST model achieves predictive skill competitive with FGN and the IFS ensemble forecast, and generates a 15-day forecast on 4 GH200 GPUs in less than 20 s.
Source: §VI · p. 10Takeaways & Limitations
BEAST demonstrates that large-scale Bayesian atmospheric forecasting can quantify both epistemic and aleatoric uncertainty while generating large ensembles efficiently.
Source: §VI · p. 10The authors identify improving stability for lower-precision computations as future work.
Source: §VI · p. 11
Abstract
from arXiv · showhide
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.
I. OVERVIEW
Atmospheric AI models have made strong deterministic forecasts, but probabilistic models have largely captured atmospheric aleatoric uncertainty rather than epistemic uncertainty. BEAST addresses the computational barriers to large Bayesian atmospheric models, against a background of parallelization approaches with distinct memory and scaling trade-offs.
I. OVERVIEW
- Rare weather phenomena are difficult to predict because they are sparsely represented in training data, and prediction quality is coupled to dataset size [1].The passage notes that this also makes rare Earth-system processes difficult to model when only small datasets are available.Source: §I · p. 1
- Existing probabilistic atmospheric models primarily capture aleatoric uncertainty, while a formulation capturing both aleatoric and epistemic uncertainty had not been realized.The passage identifies deep ensembles in FGN and HENS as approximations to Bayesian inference, not as that combined formulation.Source: §I · p. 1
- Large Bayesian atmospheric models have been prohibitively costly because weight distributions increase parameter memory and sampled forward passes multiply computation and memory demand.These costs had prevented large-scale Bayesian implementations in atmospheric forecasting.Source: §I · p. 2
- BEAST realizes a Bayesian atmospheric forecasting model at global 0.25° resolution and trains a 2.4-billion-parameter Bayesian Swin Transformer at 3.96 EFLOP/s on 5120 JUPITER nodes.The reported forecasts are for 12 h and 24 h.Source: §I · p. 2
II. STATE OF THE ART
- Parallelization approaches distribute layers, weights, or data, but long-sequence scientific image models also face large activation-memory demands and compute-utilization trade-offs.Pipeline parallelism can leave GPUs idle; tensor, sequence, and hybrid data–model approaches address different parts of the memory and computation problem.Source: §II · p. 2
III. INNOVATIONS
BEAST’s innovations combine a four-dimensional parallelization strategy with a Bayesian model formulation. Variational inference makes the weight posterior tractable, and the training objective is encoded through an evidence lower bound.
III. INNOVATIONS
- BEAST introduces four-dimensional parallelization to alleviate memory bottlenecks and maximize per-evaluation throughput.The paper’s contribution is framed as overcoming challenges in training large Bayesian neural networks.Source: §III · p. 3
- BEAST uses a 2D Swin Transformer with one down-sampling and one up-sampling layer, trained under a Bayesian model formulation.That formulation predicts by integrating the likelihood over the posterior distribution of weights.Source: §III · p. 3
- Because the exact weight posterior is generally intractable, variational inference approximates it with a parameterized family and independent Gaussian distributions for each weight.The mean-field assumption makes the parameters θ = {µ, σ} independent for each weight to make computation feasible.Source: §III · p. 3
- Training optimizes variational parameters toward the KL divergence between the variational and ideal weight distributions, encoded by the ELBO.The ELBO includes expected data fit, with λ as an optional cold-posterior factor.Source: §III · p. 3
A. Domain–Tensor Parallelism
Domain–tensor parallelism reduces per-GPU memory demands by splitting each atmospheric input across spatial and channel dimensions while distributing model weights. It extends this sharding across BEAST’s neural-network layers.
A. Domain–Tensor Parallelism
- DTP shards each input along spatial and channel dimensions, addressing the activation-heavy workloads of scientific image models while distributing both data and weights.The paper describes this as a domain-parallelism approach combined with tensor parallelism across neural-network layers.Source: A. Domain–Tensor Parallelism · p. 3
- In a DTP GPU grid of [c channels, s spatial] ranks, each rank loads 1/(cs) of each data item and holds 1/c of the model weights.Ranks sharing a spatial group hold the same weights and synchronize gradients by allreduce; ranks sharing a channel group are both domain- and tensor-parallel.Source: A. Domain–Tensor Parallelism · p. 3
- The Fig. 1 hierarchy distinguishes DTP input sharding from UP’s independent noise and weight samples and DP’s separate ERA5 data chunks.The schematic presents these parallelism axes from left to right.Source: Figure 1 · p. 4
- BEAST applies the strategy across its layer types, using distributed convolutions and linear layers, parallel attention heads, and spatial-group halo communication when needed.The attention calculation is parallelized across heads without additional communication, while shifted windows can require halo exchanges.Source: A. Domain–Tensor Parallelism · p. 4
B. Uncertainty Parallelism
Uncertainty parallelism distributes random weight samples across an added axis of model instances and aggregates their predictions efficiently. Combined with NICE noise injection, it supports ensemble-based uncertainty modeling and parallel data augmentation.
B. Uncertainty Parallelism
- Uncertainty parallelism distributes random weight samples across u model instances, each holding a full DTP model, and aggregates prediction sums and squared sums to compute ensemble means and standard deviations.Reducing aggregation to two tensors rather than all sample predictions lowers communication and supports more samples per forward–backward pass.Source: B. Uncertainty Parallelism · p. 4
- Parallel evaluations on the same input batch enable parallel data augmentation, which the authors use to include aleatoric uncertainty in the Bayesian model.The authors state that this augmentation improves predictive generalization and training convergence [43].Source: B. Uncertainty Parallelism · p. 4
- NICE uses noise-conditioned layer-normalization scales and shifts to inject random noise and produce ensemble outputs, making the model uncertainty aware.Each forward pass draws 32 independent standard-normal values, which shallow MLPs use to generate layer-normalization parameters.Source: B. Uncertainty Parallelism · p. 4
- Uncertainty parallelism scales to higher node counts to evaluate more random weight samples, improving approximation of the weight distribution and ultimately model quality.The scheme combines sampling parallelism with NICE noise injection.Source: B. Uncertainty Parallelism · pp. 4–5
C. Data Parallelism
Uncertainty parallelism can scale orthogonally with data parallelism, but plain data-parallel scaling is constrained by large-batch effects in atmospheric model training.
C. Data Parallelism
- Uncertainty and data parallelism add orthogonal axes, requiring d×u DTP instances for d data-parallel and u uncertainty-parallel instances.Given input and model sizes, the realistic local batch size is only one or two, so adding data-parallel instances also increases the effective global batch size.Source: C. Data Parallelism · p. 5
- Large-batch effects emerge around global batch sizes of 16 to 32 in atmospheric models, limiting plain data-parallel acceleration; uncertainty parallelism circumvents this limitation.The passage contrasts this constraint with sampling parallelism, where larger sample numbers can improve predictive performance.Source: C. Data Parallelism · p. 5
D. Performance Optimization
The performance optimizations target data access, communication, and ELBO computation. The authors evaluate the resulting four-dimensional parallelization scheme on BEAST training with mean-field variational inference.
D. Performance Optimization
- Pre-sharding each DTP rank’s input data by channel and spatial domain facilitates file staging, while 128-rank access waves and 50-iteration desynchronization stabilize I/O.Source: D. Performance Optimization · p. 5
- Fusing NICE MLPs and using grouped normalization eliminate split-channel communication, while combined backward gradient communication halves communication in that pass.Source: D. Performance Optimization · p. 5
- Vectorizing the communication-free prior-matching term and overlapping it with the data-fitting term alleviates the ELBO loss computation bottleneck.The data-fitting term requires communication of locally obtained ensemble members.Source: D. Performance Optimization · pp. 5–6
IV. PERFORMANCE MEASUREMENT
- The authors evaluate the four-dimensional scheme—domain, tensor, uncertainty, and data parallelism—on BEAST training with mean-field variational inference.Source: §IV · p. 6
A. Model Architecture
BEAST uses a 2D Swin Transformer architecture and is trained for global atmospheric forecasting. The experiments use distinct JUPITER and LUMI systems and model configurations.
A. Model Architecture
- BEAST uses a 2D Swin Transformer with convolutional patch embedding, downsampling, windowed attention, upsampling, and convolutional recovery.The global domain is rolled by half a window between attention layers, with periodic boundaries in latitude.Source: A. Model Architecture · p. 6
B. Training objective
- BEAST predicts 12- and 24-hour forecasts from atmospheric inputs at t and t−12h, using ERA5 data at 0.25° resolution from 1979–2019 for training and 2020–2025 for evaluation.Training minimizes the ELBO with Adam, an initial learning rate of 10^-4 decaying to 10^-5, cosine decay, and a 10,000-step warm-up.Source: B. Training objective · p. 6
C. System specifications
- Training uses JUPITER nodes with NVIDIA GH200 GPUs and LUMI nodes with AMD MI250x GPUs.JUPITER nodes pair a Grace CPU with a Hopper GPU; LUMI-G nodes pair four MI250x GPUs with an AMD EPYC CPU.Source: C. System specifications · p. 6
D. Software
The software stack uses containers and PyTorch-based tools across JUPITER and LUMI. The section defines computational, ensemble-throughput, and predictive-skill measures for evaluating training and forecasting.
D. Software
- Apptainer containers support portability and reproducibility across HPC systems, with NVIDIA NGC PyTorch on JUPITER and a custom ROCm+PyTorch container on LUMI.The Bayesian model is converted from a deterministic torch model with torch_blue; jigsaw handles domain and tensor parallelism.Source: D. Software · p. 7
E. Performance metrics
- Peak performance uses FLOP per forward–backward-update step, while sustained performance uses FLOP per full epoch, each divided by measured time.Step timing includes data I/O, checkpointing, and evaluation; epoch timing captures computation offloaded to accelerators.Source: E. Performance metrics · p. 7
- Ensemble throughput is sample/s, calculated from forecast/s, uncertainty-parallel instances, and local samples, to estimate the time needed to generate ensembles.The section also reports idle time, memory utilization, and average power draw as accelerator-use metrics.Source: E. Performance metrics · p. 7
- Skill evaluation compares ensemble-mean RMSE, spread skill ratio (SSR), and fair CRPS against FGN and IFS ensemble forecasts, all verified against ERA5.RMSE measures mean-forecast accuracy, SSR uncertainty calibration, and CRPS overall probabilistic prediction quality.Source: E. Performance metrics · p. 7
- Three ExtremeWeatherBench cases—a heat wave, a freeze, and a tropical cyclone—probe prolonged temperature anomalies and a compact, rapidly evolving severe-weather system.The cases are the 2022 UK–France–Spain heat wave, the 2022 Western and Central Europe freeze, and category 3 cyclone Damien over Western Australia in 2020.Source: E. Performance metrics · p. 7
V. PERFORMANCE RESULTS
The performance results use large configurations for scaling and a smaller, deeper configuration intended to maximize predictive skill. Baseline runs establish per-node throughput and hardware utilization.
V. PERFORMANCE RESULTS
- The 64K-sequence configurations restrict local batch size and local samples to 1, while JUPITER-0.7B-1.8K is smaller and deeper to maximize predictive skill.The long attention-window sequence makes training the larger configurations to convergence prohibitively slow per sample.Source: §V · pp. 7–8
A. Baseline performance
- JUPITER-2.4B-64K and LUMI-1.1B-64K use 16 GPUs in their baseline configurations, with one data-parallel group and two uncertainty-parallel groups.This corresponds to four JUPITER nodes or two LUMI nodes.Source: A. Baseline performance · p. 8
- 1.06 PFLOP/(s Node) peak and 0.92 PFLOP/(s Node) sustained performance were measured for JUPITER-2.4B-64K; LUMI-1.1B-64K reached 0.5 PFLOP/(s Node) peak.JUPITER’s baseline achieved 0.59 forecast/s, while LUMI’s achieved 0.22 forecast/s.Source: A. Baseline performance · p. 8
B. Scaling experiments
Scaling experiments show strong data-parallel scaling across JUPITER and LUMI, while uncertainty-parallel scaling can lose efficiency at large scale during training. The smaller predictive configuration also supports high forecast and ensemble throughput.
B. Scaling experiments
- The uncertainty-parallel weak-scaling tests increase weight samples per batch; one setup holds the global batch size at 16 while increasing uncertainty-parallel instances from 2 to 64.The experiments examine I/O behavior and whether data- and uncertainty-parallel communication compete.Source: B. Scaling experiments · p. 8
- At 2048 nodes, JUPITER-0.7B-1.8K training reached 1293.31 sample/s at 52.35% scaling efficiency, versus 5716.61 sample/s and 82.01% in inference.Training throughput degrades as loss calculation requires more communication; this effect disappears in inference, where uncertainty-parallel groups are independent. For the larger configurations, throughput also deteriorates at high node counts, attributed to file-system stress and loss-calculation communication.Source: B. Scaling experiments · p. 8, Figure 3 · p. 9
- 3.96 EFLOP/s peak performance and 73.99% scaling efficiency were reached by JUPITER-2.4B-64K on 5120 nodes (20,480 GH200 GPUs); LUMI-1.1B-64K reached 463.23 PFLOP/s and 92.09% efficiency on 1024 nodes.The results indicate near-perfect data-parallel strong scaling across all three configurations and both systems, supporting system portability.Source: B. Scaling experiments · p. 9, Figure 3 · p. 9
- 1.13 forecast/s at one data-parallel instance was achieved by JUPITER-0.7B-1.8K, almost double the throughput of JUPITER-2.4B-64K.With 16 data-parallel instances on 1536 nodes, it reached 14.08 forecast/s with 96 ensemble members at 239.88 TFLOP sustained performance.Source: B. Scaling experiments · p. 9, Figure 3 · p. 9
C. Full model training and predictive skill
BEAST is trained as a 700-million-parameter model with 96 random weight samples on 40 years of data. It achieves competitive-to-superior predictive scores and strong results on the described extreme-event cases.
C. Full model training and predictive skill
- BEAST trained for 983,350 optimizer updates on 384 JUPITER nodes using 96 random weight samples per update, achieving competitive-to-superior predictive scores with better-calibrated uncertainty.Training used four data-parallel and 48 uncertainty-parallel instances on 40 years of data; the model also had stable roll-out for up to ten days.Source: C. Full model training and predictive… · p. 9, Figure 4 · p. 10
- BEAST’s mean ensemble prediction follows cyclone Damien’s observed track, while its ensemble spread encloses the track and increases with forecast lead time.For the heat wave and freeze, it forecast the events several days in advance with skill scores competitive to or even superior to IFS.Source: C. Full model training and predictive… · pp. 9–10
- A 15-day BEAST forecast trajectory takes about 16–17 seconds on four GH200 GPUs, while FGN reports just under one minute on one TPU v5p.The passage reports these inference times on different hardware.Source: C. Full model training and predictive… · p. 10
- The orthogonal combination of domain, tensor, uncertainty, and data parallelism supports prioritizing model size, ensemble size, or batch size for different modeling objectives.The paper identifies these priorities as building models beyond one billion parameters, maximizing ensemble-generation throughput for uncertainty quantification, or maximizing forecast throughput for faster training and inference.Source: C. Full model training and predictive… · p. 10
VI. IMPLICATIONS
BEAST makes large-scale Bayesian atmospheric forecasting possible with a four-dimensional parallelization scheme, and its trained model combines competitive predictive skill with fast ensemble generation. The authors also point to applications beyond weather forecasting while identifying reliance on large supercomputers as a current boundary.
VI. IMPLICATIONS
- BEAST is the first large-scale Bayesian Swin Transformer for atmospheric forecasting to quantify both epistemic and aleatoric uncertainty, enabled by domain-, tensor-, uncertainty-, and data-parallelism.The authors state that computational training cost had previously prevented fully Bayesian global atmospheric forecasting.Source: §VI · p. 10
- BEAST's 700-million-parameter model achieved predictive skill competitive with FGN and the IFS ensemble, while a 15-day forecast ran on 4 GH200 GPUs in less than 20 s.Training peaked at 3.96 EFLOP/s on 20,480 NVIDIA GH200 GPUs on JUPITER.Source: §VI · p. 10
- BEAST predicts extreme events such as heat waves and tropical cyclones with high accuracy, while calibrated certainty helps decision-makers assess event likelihood and severity.The authors describe calibrated uncertainty estimates as crucial to preparing for extreme weather and reducing potential damage.Source: §VI · pp. 10–11
- The authors see BNNs as well suited to small, high-quality Earth-system datasets, including km-scale climate simulations for future climate emulators.They also identify flexible GPU-memory allocation through domain–tensor-parallelism as a route toward multi-timestep forecasting and km-scale modeling.Source: §VI · p. 11
- The work relies on top-tier supercomputers with thousands of GPUs; future work targets low-precision stability and additional parallelization axes.The authors describe uncertainty parallelism as an orthogonal scaling axis beyond data parallelism, which they say is limited by large-batch effects.Source: §VI · p. 11