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

Overview

Artificial intelligence (AI) models for atmospheric forecasting are revolutionizing weather and climate sciences, owing to their promise of generating predictions with skill scores competitive to numerical models at a fraction of the computational resources. Current atmospheric AI models learn a set of point-value (scalar) weights through iterative optimization over a finite dataset. However, this approach disregards the fact that the results of the optimization process can be ambiguous, meaning that the weights could also have different numeric values. From a modeling point of view, this corresponds to the epistemic uncertainty of the model – a systematic error originating from the incompleteness of the training datasets. The uncertainty in the model weights is dependent on the size and quality of the dataset: the more accurately the training dataset represents the true data distribution of the modeled problem, the more certain the weights are and the better the prediction becomes. For atmospheric AI models, this has two direct consequences. For one, the prediction of rare weather phenomena, such as extreme events, is challenging, as they are sparsely represented in the training data, or not at all. Second, the prediction quality is directly coupled to the size of the dataset [?], as more data items typically provide a better approximation of the true data distribution. Hence, modeling rare physical processes of Earth system components for which only small-sized datasets are available remains difficult.

After the recent successes in AI-based weather forecasting with deterministic large-scale neural networks [?, ?, ?, ?], probabilistic models for quantifying the uncertainty of a prediction have become the next frontier. Current approaches to uncertainty quantification in AI models focus on the uncertainty of the atmospheric input state to generate ensemble forecasts, similar to ensemble generation via input perturbation in numerical models, as used for HENS [?, ?]. In particular, diffusion-based methods like GenCast [8], functional generative networks (FGN) [9], FourCastNet3 (FCN3) [10], and AIFS-CRPS [11] have achieved impressive results. However, these existing approaches capture only the so-called aleatoric uncertainty, which stems from the inherent chaotic nature of the physical processes in the atmosphere. In contrast, modeling epistemic uncertainty in atmospheric AI models has received little attention so far. FGN and HENS approximate Bayesian inference with deep ensembles [12, 13], i.e., training multiple models from independently randomly initialized parameters; however, a probabilistic model formulation that captures both epistemic and aleatoric uncertainty has not yet been realized.

Bayesian neural networks (BNNs) provide a principled framework for modeling epistemic uncertainty. BNNs extend classical neural networks by replacing the scalar weights with probability distributions, thereby transforming point predictions into predictive distributions. Compared to alternative methods, like input perturbation or diffusion, BNNs offer a more comprehensive and theoretically grounded approach to uncertainty estimation, leading to better-calibrated and more informative ensemble predictions [14]. The distributional parameterization of BNN weights allows for distribution tails to be better captured, leading to more accurate predictions of extreme events. Moreover, BNNs can be trained to competitive prediction accuracy on much less data compared to their deterministic (non-Bayesian) counterparts, owing to their capability of accounting for the under-representation of the data distribution in small training datasets [15]. Hence, they have great potential for AI-based modeling of other components of the Earth system, where training data is much scarcer.

However, training BNNs requires significantly more computational resources than training standard neural networks. For one, the probabilistic parameterization of network weights effectively doubles the number of parameters that must be learned, increasing the GPU memory requirements of model training. Second, each forward pass on a data batch has to be performed on a set of random weight samples drawn from the current weight distribution, and the resulting predictions, i.e., the ensemble, have to be aggregated for the backward pass. Evaluation of multiple samples per data item linearly increases computational cost and memory demand. For large-scale atmospheric models with hundreds of millions to billions of parameters, these costs become prohibitive. As a result, large-scale implementations of BNNs in this domain have not yet been realized to date.

We set out to break this barrier and realize the first-ever Bayesian model for atmospheric forecasting at global 0.25° resolution. We strategically accelerate model training at all scales, from optimizing single-node performance by leveraging lower precision data types, just-in-time compilation, and fast GPU-to-GPU communication, to minimizing inter-node communication necessary for Bayesian model training. This allows training of a 2.4-billion-parameter Bayesian Shifted Window (Swin) Transformer for 12 h and 24 h forecasts at a peak performance of 3.96 EFLOP/s on 5120 nodes on the JUPITER supercomputer.

State of the Art

Regardless of the application, parallelism and scalability are paramount to enabling the training of large-scale models. The conventional way of accelerating training is data parallelism (DP) [16, 17], where each process holds a copy of the model weights, independently performs a forward–backward pass of its local subset of the dataset, and computes the corresponding gradients. The gradients are then averaged through an allreduce operation, and the weights are updated locally. The number of iterations required for an epoch can be reduced by increasing the number of data-parallel processes to increase the global batch size. While data parallelism scales well in a computational manner, having large global batch sizes can cause models to generalize poorly, ultimately limiting scalability [18].

Recent work on AI models in weather and climate science has focused extensively on so-called foundation models, i.e., very large models that are pre-trained in a self-supervised fashion on large heterogeneous datasets, and subsequently fine-tuned for downstream tasks [19, 20, 21]. In consequence, model sizes have been growing dramatically beyond the billion-parameter mark. The affirmation of neural scaling laws for AI weather models [?], stating that predictive skill improves with quadratically more data and weights, further fuels this trend. As growing model sizes quickly exceed the memory limitations of a single GPU, the field of model parallelism has emerged to distribute model weights and associated computations across multiple GPUs.

There exist multiple forms of model parallelism, as there are different ways (axes) to parallelize the model architecture. The simplest form of model parallelism is pipeline parallelism, where sequential stages, i.e., blocks of layers, are distributed across GPUs [22]. While pipeline parallelism circumvents the GPU-memory bottleneck and entails simple communication schemes with only point-to-point operations, the inherent sequentiality of running through the network layers in the forward-backward pass causes GPU idle time, precluding efficient GPU utilization. Another drawback of pipeline parallelism is the necessity to perform asynchronous gradient updates, which degrades predictive performance [23].

An orthogonal approach to pipeline parallelism is tensor parallelism, i.e., distributing the weight matrices of each layer across multiple GPUs. Megatron-LM [24] was among the first works to utilize tensor parallelism, efficiently enabling the training of multi-billion-parameter Transformer models by parallelizing across attention heads and feed-forward weight matrices. In Transformer models specifically, the quadratic complexity in sequence length of the attention mechanism provides a particular computational burden and memory bottleneck for long sequence lengths. This is oftentimes addressed with sequence parallelism [25].

Next to the memory demand of the large weight matrices themselves, the associated gradients and optimizer states can also lead to a severe memory bottleneck. To deal with this, hybrid data–model-parallel approaches have been introduced, such as Zero Redundancy Optimizer (ZeRO) [26] or Fully Sharded Data Parallelism (FSDP) [27]. ZeRO partitions the parameters, activations, and optimizer states across GPUs, and uses allgather operations to reconstruct each layer upon demand. The approach achieves superlinear scaling due to reduced memory footprints. FSDP shards data batches and model parameters. In a given layer, the weights are gathered to compute the full activations. Afterward, the unnecessary activations are discarded to reduce the memory footprint.

The main memory bottleneck for training large language models (LLMs) originates from large weight matrices; in contrast, Vision-Transformer (ViT)-like architectures with long sequence lengths ($\approx10^{6}$) face a different challenge: large activation memory footprints. While this can be mitigated by techniques such as gradient offloading to the CPU [26, 28] or gradient recomputation [29], these strategies can come at the cost of computational performance. Thus, novel parallelization strategies are required to address the significant bottlenecks arising from the size of the data items.

As one of the first works in the field, FourCastNet v1 [?] employed tensor parallelism on its Adaptive Fourier Neural Operator (AFNO) architecture by applying feature parallelism, splitting the fast Fourier transforms across processes. It also distributes the multi-layer perceptron (MLP) at the end of an AFNO block, sharding model parameters in alternating row and column dimensions. The approach enables scaling up to 3808 NVIDIA A100 GPUs, attaining 140.8 PFLOP/s in mixed precision, which amounts to 11.9 % of peak performance [?]. ORBIT [30] used a ViT [31], trained on global 1.4° resolution with up to 91 channels, and scaled beyond the 1-billion-parameter model size. ORBIT introduces Hybrid-STOP, which fully shards the model’s weights as well as data batches, thus combining tensor parallelism and FSDP. Additionally, communication is optimized by using layer wrapping, i.e., overlapping computation and prefetching of subsequent model layers. Through these means, they train a model with up to 113 billion parameters across 49 152 GPUs. ORBIT-2 performs a spatial downscaling task, parallelized through the TILES approach [32]. In TILES, the spatial dimension is divided into overlapping spatial tiles, upon which each GPU calculates local attention. With their biggest model of 10B parameters, they achieve a sustained computational throughput of 368 PFLOP/s on 512 nodes (2048 GH200 GPUs) on the Alps system.

AERIS [33] scales training of a Swin-Transformer-based [34] diffusion model with $1 \times1$ patch embedding up to 10 080 nodes and utilizes data, pipeline, and Sequence-Window (SWiPe) parallelism. SWiPe partitions the image across windows, and further parallelizes the sequence within each window with DeepSpeed Ulysses’s sequence parallelism [25]. Their biggest model (80B parameters) reaches 10.21 EFLOP/s on 10 080 nodes of the Aurora system, corresponding to roughly 10 forecasts per second.

Innovations

Our main innovations overcome the challenges in training large Bayesian neural networks by introducing a 4-dimensional parallelization strategy to alleviate memory bottlenecks and to maximize per-evaluation throughput.

The architecture for our Bayesian Exascale Atmospheric Swin Transformer (BEAST) features a 2D-Swin Transformer [34] with one down- and up-sampling layer. We train BEAST according to the Bayesian model formulation

where $\mathcal{D}$ is the training data, $w$ are the model weights, $x^{*}$ is the input and $y^{*}$ is the output. The exact computation of the posterior $p(w \mid\mathcal{D})$ is generally intractable. To circumvent this, variational inference [35, 36] approximates the posterior by a known parameterized distributional family $q$. The mean-field assumption uses an independent Gaussian distribution for each weight, and assumes the parameters $\theta= \{\mu, \sigma\}$ to be independent for each of the weights to make computation feasible.

During training, the parameters $\theta$ are optimized towards minimizing the Kullback-Leibler divergence (KL) [37] between the variational distribution $q$ and the ideal weight distribution $p$. This is encoded by the evidence lower bound (ELBO)

where $p(Y \mid X,W)$ is the probability of the training labels according to the model given the training data and weights, which are sampled from the variational distribution for the expected value, and $\lambda$ is an optional “cold posterior” hyper-parameter factor [?].

Figure 1 outlines how we take advantage of multiple parallelization axes within our implementation of BEAST to circumvent the complexity and memory requirements inherent to the model.

Domain–Tensor Parallelism

Our first optimization target is to reduce per-GPU memory footprint. With a single data item, i.e., atmospheric input state, being several hundred megabytes large, current state-of-the-art AI models for atmospheric forecasting struggle to go beyond the 1-billion-parameter mark [30]. Accordingly, a significant proportion of the available GPU memory is consumed by data items and corresponding activations. In the Bayesian formulation, this problem becomes even more pronounced, as the memory demand of a BNN compared to a deterministic model is increased proportionally to 1) the number of parameters used to parameterize the distribution of each model weight (two in our case); and 2) the number of samples drawn—and thus, forward passes through the model and associated activations—per data item.

Existing strategies for model parallelization are designed for training LLMs and thus prioritize maximizing the model size through tensor and pipeline parallelism. Long sequence lengths are handled by sharding the input data along the sequence dimension with context parallelism. Scientific image-based applications are much larger in terms of sequence lengths (order of 1 million) with the number of input variables in the order of 100. The latent representations of the channels are projected to a larger embedding space; thus, sharding of the data along both the sequence and channel dimensions holds the biggest potential in alleviating memory bottlenecks. We pursue this approach of domain parallelism and combine it with tensor parallelism for all neural network layers. That means that, in contrast to existing parallelization schemes, our approach distributes not only the model parameter and optimizer states, but also shards each input data item along both the spatial and channel dimensions. We leverage distributed matrix multiplications to implement this domain–tensor parallelization (DTP) scheme [39]. By sharding both the data and weights and formulating a linear layer as a distributed matrix multiplication, we are able to parallelize individual and arbitrary layers without the need for allgather operations or re-sharding, simultaneously maximizing GPU memory usage and reducing the data I/O-workload.

We define a DTP instance over a GPU grid of $[c$ channels, $s$ spatial] ranks, where each rank loads $[1/c \times1/s]$ fraction of each data item and holds $[1/c]$ of the model weights. Ranks within the same spatial group hold the same copy of the model weights, i.e., they compute on different data shards and are domain-parallel only. The parameters of the spatial groups are synchronized by allreduce operations of the gradients. Ranks within the same channel group are fully domain- and tensor-parallel.

We apply this concept to all neural network layer types in BEAST. Convolutions are performed by distributing the filters across the channel groups. Each rank computes the local convolutions and exchanges the necessary input data with the remaining ranks in the channel to resolve the full convolution. A schematic of the parallelization strategy used in channel-parallel attention layers is shown in Figure 2. The query–key–value (QKV) map in each attention layer is computed by a distributed linear layer across the channels. Each rank then holds the local sequence and a fraction of the QKV embedding dimension. The attention calculation is then trivially parallelized across the heads, allowing us to leverage optimized single GPU attention kernels such as FlashAttention-2 [40] or cuDNNs fused dot product attention via NVIDIA’s Transformer Engine library [41] without any additional communication, while maintaining domain parallelism across channels on the input to the Transformer block. Head concatenation and the subsequent MLP are again performed by distributed layers, combining the information across the channel shards. The windows are shifted between every Transformer block by half a window, and halos are communicated between the ranks in spatial groups where necessary.

III-B Uncertainty Parallelism

Next, we target the challenge of evaluating the $m$ forward passes of the random weight samples for each data item by introducing sampling parallelism [42]. We extend the GPU grid of $[c$ channels, $s$ spatial$]$ ranks by an orthogonal parallelization axis with $u$ groups, along which we distribute the $m$ random weight samples. That means that we have $u$ DTP instances, each holding a full model on $c \times s$ GPUs. Each of these model instances draws $m/u$ independent random weight samples and performs the forward pass on them.

To compute the mean and standard deviation of the ensemble, the sample predictions need to be aggregated: each rank computes the sum and squared sum of its local predictions, $S_u = \sum_{i=1}^{m_u} x_i^{(u)}$ and $Q_u = \sum_{i=1}^{m_u} \left(x_i^{(u)}\right)^2$, respectively. $S_u$ and $Q_u$ are then aggregated across all $u$ ranks in an allreduce operation to obtain the global sum and global square sum $S = \sum_u S_u$, $Q = \sum_u Q_u$. From these global sums, the global mean and global standard deviation of the predictions are computed. This allows us to reduce the communication in the prediction aggregation to two tensors instead of all the $m$ predictions from the local weight samples. Through sampling parallelism, we can efficiently increase the number of random weight samples per forward–backward pass, thereby improving prediction quality without being bound to the memory capacities of a single accelerator.

Sampling parallelism provides an additional major advantage: the parallel evaluations of the random weight samples are all carried out on the same input data batch. They can therefore be combined with data augmentation, which improves predictive generalization and training convergence [43]. While there are several data augmentation options from computer vision available, we chose a different route and used parallel data augmentation to include aleatoric uncertainty into our Bayesian model.

We use a modification of conditional layer normalization [44], inspired by [9] (noise independent of spatial grid) and [11], which we call noise input conditioning for ensembles (NICE). In each forward pass, a random vector $e$ of 32 independent standard normal distributed values is drawn. The scale and shift parameters for each layer normalization are generated by shallow MLPs, with the noise $e$ as the input. This effectively injects noise into the model and can itself be used to produce ensemble outputs, thus making the model fully uncertainty aware. We call this combination of sampling parallelism with random noise injection uncertainty parallelism (UP). We handle the coordination of the data I/O within UP groups with synchronized random seeds. Uncertainty parallelism allows us to evaluate more random weight samples by scaling to higher node counts, thereby improving the approximation of the current weight distribution and, ultimately, model quality.

Data Parallelism

Uncertainty parallelism can be combined orthogonally with data parallelism to accelerate and improve model training at the same time, adding an additional orthogonal scaling axis with $d$ groups to our GPU grid. Training with $d$ data parallel (DP) instances and $u$ uncertainty parallel instances thus requires $d \times u$ DTP instances in total. However, given the size of the input data items and the model, a realistic setup can only run with one or two local data items, i.e., local batch size, or random weight samples per GPU. Hence, data parallel scaling out also adds more data items to the effective global batch size.

In contrast to sampling parallelism, in which having larger sample numbers can lead to better predictive performance, scaling out data parallelism will lead to the occurrence of large batch effects [18]: AI models tend to converge to sharper local minima when the number of data items in a gradient update step becomes too large. For atmospheric models, this effect starts to arise at relatively small global batch sizes of around 16 to 32 [45]. The degree to which model training can be accelerated through plain data parallelism is thus limited. With the additional usage of uncertainty parallelism, we circumvent this limitation.

Performance Optimization

We further maximize per-GPU computational throughput with the following performance optimizations. Since each GPU in a DTP instance loads a shard of the input data sample, we pre-shard the dataset accordingly in the channel and spatial (longitude) domain, thereby facilitating file access for data staging. To further stabilize data I/O and alleviate the strain on the file system, we perform data access in waves of 128 ranks in each wave and de-synchronize file system access every 50 iterations.

We fuse the MLPs of the NICE normalization and use grouped normalization to eliminate communication across the split channel domain. Our implementation of the domain–tensor-parallel layers requires us to overwrite the backward function of each layer type. This allows us to also combine gradient communication in the backward pass with respect to the input and weights. This step essentially halves communication in the backward pass, significantly improving performance.

We handle communication for UP and DP groups via PyTorch’s DistributedDataParallel wrapper [16]. We also apply post-grad accumulation hooks on the weight synchronization between spatial ranks to better overlap communication with the backward pass.

Additionally, we found that the computation of the ELBO loss takes significantly longer than the computation of a regular RMSE loss in a non-Bayesian neural network. To alleviate this bottleneck, we vectorize the computation in the prior matching term, which does not require communication, and overlap it with the computation for the data-fitting term, which does require communication of the locally obtained ensemble members.

Component	JUPITER	LUMI
Container Runtime	Apptainer	SingularityCE
Operating System	RHE Linux 9.9	SUSE 15.6
Acc Library	CUDA 13.2	ROCM 6.2.4
XCCL	NCCL 2.29.7	RCCL 2.22.5
PyTorch	2.12.0	2.7.1
Python	3.12	3.12

Performance Measurement

We evaluate our 4D parallelization scheme, consisting of domain-, tensor-, uncertainty-, and data-parallelism, on the task of training BEAST with mean-field variational inference.

Model Architecture

The model architecture of BEAST, visualized in Figure 2, is based on a 2D Swin Transformer [34]: The input data is embedded via a 2D non-overlapping convolution with a spatial patch size of $2 \times2$. Then, multi-headed self-attention is performed on non-overlapping windows in a Swin-Transformer-manner. In between each attention layer, the entire global domain is rolled by $(\lfloor\text{window size}/2\rfloor)$ patches in both spatial dimensions. Periodic boundaries are used in the latitude dimension. The data is then downsampled by a 2D convolution with a patch size of 2, which also doubles the hidden dimension. Another Swin Transformer block is applied before the data is upsampled back with a transposed convolution and an MLP to mix across the channels, before another Swin Transformer block is applied. Finally, the target, with the original spatial dimensions, is recovered with a transpose convolution.

Training objective

We train BEAST to perform 12- and 24-hourly forecasts based on two previous time steps, i.e., atmospheric variables at times $t$ and $t-12h$ are used as input to predict atmospheric variables at times $(t+12h,t+24h)$. For training, we use the ERA5 dataset [46], obtained from WeatherBench2 [47] and stored in the zarr format. The data uses an equiangular grid at a global resolution of $0.25^\circ$ ($\approx30$ km). We consider 84 input variables in total—six surface variables: 10 m u-velocity, 10 m v-velocity, mean sea level pressure, 2 m temperature, sea surface temperature, total precipitation, and six pressure variables: geopotential, specific humidity, temperature, and the u, v, and w components of wind velocity. Pressure variables are considered at 13 pressure levels of 1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, and 50 hPa. Constant inputs of soil type, topography, and land masks are also used. Variables from different time steps are concatenated along the channel dimension. Data is normalized with per-variable Z-score normalization. Hourly data from 1979–2019 is used for training, and data from 2020–2025 is used for evaluations.

For minimizing the ELBO as loss function, we use an Adam optimizer with an initial learning rate of $10^{-4}$, decaying to $10^{-5}$ with a cosine decay. Linear warm-up is used for the first 10 000 gradient steps. We train the model using mixed precision format: for Transformer Engine and FlashAttention in the Transformer layers, we leverage Bfloat16 precision. All other computations are performed in TensorFloat32 or Float32 precision, to reflect the sensitivity of the Bayesian optimization objective to numerical rounding errors.

System specifications

We perform training of BEAST on two systems: JUPITER and LUMI.

JUPITER is a EuroHPC JU supercomputer hosted at Forschungszentrum Jülich, Germany. It is the first exascale machine in Europe and ranks #5 in the TOP500 of June 2026, reaching 1 EFLOP/s in the HPL benchmark. JUPITER has been designed for energy-efficient, high-scaling numerical and AI applications [48], and is currently in the finalization of its build-up phase. The system utilizes 23 536 NVIDIA GH200 superchips, each combining a 72-core NVIDIA Grace CPU (120 GB LPDDR5 memory) with a Hopper GPU (96 GB HBM3 memory) through a unique high-bandwidth bus called NVLink-C2C (900 GB/s bandwidth). The full system features 5884 nodes, of which about 5200 were accessible to us. Within a node, the GPUs are all-to-all connected with pair-wise bandwidths of 300 GB/s. Similarly, all CPUs are all-to-all-connected with pair-wise bandwidths of 200 GB/s.

Each node features four NVIDIA InfiniBand NDR200 network adapters with each 200 Gbit/s injection bandwidth, interconnected in a Dragonfly+ topology. Up to 240 nodes are combined in a Dragonfly group, internally connected over two levels of switches in a fat tree (local) topology. Twenty-five of these groups form the JUPITER network, with 30 direct links between each pair of groups; each inter-switch link provides 400 Gbit/s bandwidth. Adaptive routing is enabled to counteract congestion due to the tapering of the network outside of the groups.

LUMI is a EuroHPC JU supercomputer hosted by the CSC – IT Center for Science in Finland. It is a pre-exascale Cray EX supercomputer supplied by Hewlett Packard Enterprise (HPE). As of June 2026, LUMI ranks #11 in the TOP500 list, reaching 386 PFLOP/s in the HPL benchmark. The majority of the system’s compute power is found in the LUMI-G hardware partition, which consists of 2978 nodes. Each node is equipped with four AMD MI250x GPUs and a single 64-core AMD EPYC Trento CPU. The MI250X GPU is a multi-chip module containing two GPU dies named Graphics Compute Die (GCD) by AMD. Each die includes 110 compute units and has access to a 64 GB slice of HBM2e memory, resulting in a total of 220 compute units and 128 GB of total memory per MI250x module.

LUMI uses a Dragonfly network topology in which compute nodes are organized into groups interconnected via switches. The compute nodes are linked through the HPE Cray Slingshot-11 200 Gbit/s network interconnect (NIC). The HPE Cray Slingshot NIC features high-performance RDMA and hardware acceleration for MPI and SHMEM-based software.

Software

The software stack employed in this work is based on containerized environments, leveraging Apptainer / Singularity Community Edition [49] to ensure portability and reproducibility across high-performance computing (HPC) systems. Apptainer enables unprivileged container execution and is widely adopted in scientific and high-performance computing contexts. On JUPITER, the base container is the NVIDIA NGC PyTorch image (version 26.04-py3 [50]), available on the NVIDIA NGC registry. For LUMI, we rely on a provided custom ROCm+PyTorch container [51].

Domain and tensor parallelism of the PyTorch layers was done by using the library jigsaw [39], and the Bayesian model was obtained by converting the deterministic torch model with torch_blue [52]. Profiles and traces for optimization were recorded using the built-in torch profiler and NVIDIA Nsight Systems 2026.2.1. Power draw and memory utilization measurements for the hardware accelerators were sampled during training using perun [53] for both systems.

Performance metrics

The most relevant performance metrics when training neural networks are predictive performance and overall training time. For a fixed-size dataset, training time is influenced by two factors: the time it takes to iterate through the entire dataset in batched forward-backward passes (time-per-epoch) and the number of epochs, i.e., iterations over the entire dataset, necessary for model convergence (convergence rate). While the latter is generally addressed algorithmically, the former can be directly tied to the computational throughput of performing one batch iteration (forward–backward pass and weight update).

To assess the computational throughput of the batch iterations in our scaling experiments, we measure the time for 1) a forward–backward-update step, and 2) one full epoch using Python’s high-resolution time module with time.perf_counter_ns() and the torch.cuda.streams.Event() to set markers to monitor accelerator computations. The former includes the overhead associated with model training: data I/O, checkpointing, and evaluation; whereas the latter captures exclusively the computation that is offloaded to hardware accelerators. We further estimate the number of floating-point operations (FLOP) of a forward–backward pass on a DTP instance, using PyTorch’s torch.utils.flop_counter tools. Based on this, we report the number of FLOP/s for 1) a forward–backward-update-step as peak performance and 2) for a full epoch as sustained performance by dividing the number of FLOP by the measured time. We also report accelerator metrics related to ideal use during the forward–backward-pass, including percentage of idle time, memory utilization, and average power draw.

In addition to these computational throughput numbers, we also report the commonly used metric images/sec from computer vision. In atmospheric modeling, this corresponds to the forecast throughput or forecast/s, which is estimated based on the average batch processing time, and the total number of data items per batch (DP instances × local batch size). We also introduce the orthogonal metric of sample/s, determined via the product of forecast/s × UP instances × local sampling size, for ensemble generation throughput. This metric indicates throughput regarding the evaluation of random weight samples of the BNN and is especially relevant for domain scientists during inference over an evaluation dataset, indicating how long it takes to generate a meaningful ensemble.

Next to compute performance, we evaluate the probabilistic predictive skill scores of a fully trained BEAST, where we report the commonly used root mean square error (RMSE) of the ensemble average prediction, probabilistic spread skill ratio (SSR), and continuous ranked probability score (CRPS) metrics for different atmospheric variables. These metrics cover three different aspects: The RMSE of the ensemble mean measures how well the average forecast matches the target. The SSR quantifies whether the model’s uncertainty is well-calibrated. The CRPS, which we calculate using the fair CRPS estimator [54], combines these two aspects to measure overall probabilistic prediction quality. We compare our predictive skill scores against the current-best probabilistic AI model, FGN [9] vs. ERA5, and the IFS ensemble forecast [55] vs. ERA5, both obtained from [47].

We further assess the skill of BEAST on predicting extreme events, by evaluating three exemplary use-cases, using the recently published ExtremeWeatherBench (EWB) Python package [56]: a) a heat wave, b) a freeze, and c) a tropical cyclone (TC). The heat wave occurred from August 8, 2022, until August 16, 2022, and affected southern UK, France, and Spain. The freeze affected Western and Central Europe from December 10, 2022, until December 20, 2022. For the TC, we used Damien, a category 3 cyclone, which occurred between February 3, 2020, and February 9, 2020, over Western Australia. These cases probe both prolonged large-scale temperature anomalies and the emergence and subsequent trajectory of a compact, rapidly evolving severe-weather system.

Performance Results

Table 2 describes the three model configurations used in our experiments. The JUPITER-2.4B-64K and LUMI-1.1B-64K model architectures are optimized to leverage the available hardware and train the maximal model size as efficiently as possible. Due to the size of these model configurations, the local batch size and number of local samples are restricted to 1. However, the long sequence length considered in the attention windows (64 000) causes the training time per sample to become prohibitively long for training a model until convergence. Therefore, we define JUPITER-0.7B-1.8K, a smaller but deeper model configuration designed to maximize predictive skill.

Name	Params (B)	TFLOP	Outer windows: Dim.	Outer windows: Win. size	Outer windows: Heads	Outer windows: Layers	Inner windows: Dim.	Inner windows: Win. size	Inner windows: Heads	Inner windows: Layers
JUPITER-2.4B-64K	2.36	3557.9	3072	[360, 180]	24	2	6144	[180, 30]	48	2
LUMI-1.1B-64K	1.05	2238.2	2048	[360, 180]	16	2	4096	[180, 30]	32	2
JUPITER-0.7B-1.8K	0.68	423.8	1024	[30, 60]	16	4	2048	[20, 30]	32	6

For all three configurations, we distribute one DTP instance on a GPU mesh of $[s,c] = [4,2]$, i.e., 8 GPUs (2 nodes) on JUPITER and 8 GCDs (1 node) on LUMI, when performing training steps (forward-backward-update passes). At inference, where the memory footprint is much smaller due to the absence of gradients and optimizer states, we distribute one DTP instance on a GPU mesh of $[s,c] = [2,2]$, i.e., 4 GPUs.

V-A Baseline performance

We define baseline performance on the LUMI-1.1B-64K and JUPITER-2.4B-64K model configurations. The baseline configurations require 16 GPUs on both systems: a domain–tensor group of $[4,2] \times2$ UP groups, and a single data parallel group. This translates to 4 compute nodes on JUPITER, and 2 nodes on LUMI.

Based on an average batch time (forward–backward-update) of 1.68 s on the JUPITER system, and an estimated 3.56 PFLOP operations given the model configuration, JUPITER-2.4B-64K reaches a peak performance of 1.06 PFLOP/(s Node), and a sustained performance of 0.92 PFLOP/(s Node). This translates to 0.59 forecast/s. We measured the amount of idle time and time spent waiting for communication by profiling the system on 20 gradient updates, showing $\sim25$ % of the time is spent waiting on communication, $\sim1$ % device idle time, and $\sim72$ % of the time is being used for computation. Within an epoch, memory measurements for the individual GPUs report 71.45 GB consistently, utilizing 74.43 % of the total accelerator memory. With the LUMI-1.1B-64K configuration, the average forward–backward-pass runs for 4.47 s. Based on an estimated 2.24 PFLOP, we reach single node peak performance of 0.5 PFLOP/(s Node), or 0.22 forecast/s. Memory measurements show each GCD reporting an average of 48.87 GB, or 76 % GPU memory utilization.

Additionally, we measured the power draw of the hardware accelerators during training to provide a realistic estimate of the energy cost of training and running our Bayesian model [57]. We observe an average power draw of 455.6 W on each MI250X card, and 3.64 kW for training LUMI-1.1B-64K in the baseline configuration. By comparison, we measured a consistent power draw of 514.36 W per GPU on the JUPITER system, resulting in a total GPU power draw of 8.22 kW for training JUPITER-2.4B-64K over the four nodes in the baseline configuration.

V-B Scaling experiments

To evaluate the efficiency of uncertainty parallelism, we scale out the number $u$ of UP groups in a weak scaling manner, i.e., we increase the number of global weight samples evaluated per data batch by increasing the number of UP instances. For JUPITER-2.4B-64K and LUMI-1.1B-64K we conduct this experiment in two regimes: for one, we scale $u$ from 2 to 16, with a single DP instance, i.e., using 16 through 128 GPUs, or GCDs in LUMI. Such a configuration is not very realistic for production training runs, where multiple DP instances would be used to iterate over the dataset faster. Hence, we conduct a second round of this scaling experiment, where we keep the number of DP instances, i.e., the global batch size, fixed at 16 and increase the number of UP instances from 2 to 64, i.e., using 256 to 16384 GPUs. For JUPITER-0.7B-1.8K we evaluate UP scaling in both training and inference mode, with a fixed global batch size of 4, increasing the number of UP instances from 8 to a total of 2048 nodes on JUPITER. The scaling experiments shed light on two interesting aspects of our multidimensional parallelization scheme. For one, DP requires different data to be read from storage by each rank, while for UP, each rank reads the same data. Having both UP and DP instances participate in the training illustrates how our parallelization scheme handles these different I/O streams. Second, both UP and DP utilize PyTorch's DistributedDataParallel wrapper [16] to handle communication. Thus, the experiment allows us to investigate if these two communication axes potentially compete against one another, thereby hindering performance.

The results for JUPITER-2.4B-64K and LUMI-1.1B-64K are shown in Figure 3 (a). The values in the gray section indicate the scaling experiment where the DP size is 1, and the values after the gray section showcase the results of the configuration with the increased DP size of 16. On LUMI, we observe almost constant throughput for lower node counts up to 3072 GPUs, whereas on JUPITER, we observe a small degradation of per-GPU performance. On both systems, performance starts to deteriorate notably at larger node counts. This scalability limit of UP can be attributed to the increased stress on the file system and higher communication cost during the loss calculation. Ensemble throughput scaling of JUPITER-0.7B-1.8K is presented in Figure 3 (c). On 2048 nodes, the training configuration reaches a total of 1293.31 sample/s, with a scaling efficiency of 52.35 %, similar to the results shown in Figure 3 (a). Again, we observe a degradation in performance, which can be attributed to the additional communication required to calculate the loss for a larger number of samples generated per forward pass. This effect disappears in inference mode, where the loss is not calculated, and the uncertainty parallel groups become independent of each other. On 2048, the predictive model reaches 5716.61 sample/s, with an ensemble scaling efficiency of 82.01 %.

Next, we evaluate scaling out data parallelism in a strong scaling manner: We increase the global batch size by increasing the number of DP instances from 1 to 16 at a fixed local batch size of 1, thereby reducing the overall number of iterations in an epoch. Similarly to UP weak scaling, we conduct this strong scaling experiment for JUPITER-2.4B-64K and LUMI-1.1B-64K in two setups: one with just 2 UP instances, thus scaling from 2 to 32 nodes, and one with 16 UP instances, thus scaling from 64 to 5120 nodes. For JUPITER-0.7B-1.8K we keep the number of random weight samples (ensemble members) in the training configuration at 96, using 48 UP instances, and scale out DP from 1 to 16, i.e. 96 to 1536 nodes. Results are presented in Figure 3 (b) and (d). Again, the values in the gray sections indicate the scaling experiments where the UP size is 2, i.e., training with a 2 UP instance, and the values after the gray section showcase the results of the configuration with the increased UP size of 16.

We observe near-perfect DP-strong scaling for all three model configurations on both JUPITER and LUMI. On 5120 nodes on JUPITER, i.e., 20 480 GH200 GPUs, JUPITER-2.4B-64K runs at a total peak performance of 3.96 EFLOP/s during forward–backward–update step, and a scaling efficiency of 73.99 %. LUMI-1.1B-64K reaches 463.23 PFLOP/s at 1024 nodes, corresponding to 92.09 % scaling efficiency. These results indicate that the parallelism strategy and model configuration space allow for good system portability. The seamless transition between the different UP and DP instance sizes highlights the efficiency of both parallelism axes.

JUPITER-0.7B-1.8K yields a higher forecast throughput per data parallel group than the large configurations. At a single DP instance, the predictive configuration achieves a 1.13 forecast/s, almost double the throughput of JUPITER-2.4B-64K. With 16 DP instances, or 1536 nodes, it reaches 14.08 forecast/s, each with 96 ensemble members, at a sustained performance of 239.88 TFLOP.

V-C Full model training and predictive skill

We train BEAST in the JUPITER-0.7B-1.8K configuration on 384 nodes on JUPITER, using 4 DP and 48 UP instances, for 983 350 optimizer update steps. We use 2 samples per UP instance and keep the local batch size to 1, thus using a global batch size of 4 and drawing 96 random weight samples per optimizer update step. One epoch on 40 years of data (1979–2020) requires 89 835 gradient update steps and takes 25.45 h. The sustained training performance is 70.46 PFLOP/s. Predictive metrics of the trained model for four different variables (U10, T2M, U850, T850) for the two forecast time steps 12 h and 24 h are shown in Figure 4, in comparison to the state-of-the-art probabilistic model FGN and the IFS ensemble forecast. We observe competitive to superior predictive scores with much better calibrated uncertainty estimates, and stable roll-out for up to ten days.

Furthermore, BEAST can predict extreme events with exceptional skill. For one, it excels at the task of TC track prediction, as shown in Figure 5 for TC Damien. The mean ensemble prediction of BEAST follows the observed track, including the transition from land to ocean and the subsequent landfall. The ensemble members form a spread that encloses the observed track, with uncertainty increasing with forecast lead time as expected. Second, it can forecast both the heat wave and the freeze several days in advance, with skill scores competitive to or even superior to the IFS, as shown in Figure 6. These results show that BEAST’s probabilistic forecasts remain informative not only in aggregate verification, but also for rare events where accurately representing both the expected evolution and its uncertainty is particularly consequential.

Increasing the number of random weight samples not only impacts computational performance but also predictive skill. Table 3 illustrates how larger ensemble sizes lead to better predictive skill in terms of lower RMSE. This result highlights the benefit of our uncertainty parallelism approach, which allows us to efficiently generate large ensembles by scaling out uncertainty parallelism. Generating a 15-day forecast trajectory with BEAST takes about 16 s to 17 s on 4 GH200 GPUs, independent of the ensemble size. For comparison, FGN reports inference times of just under 1 minute on a single TPU v5p [9].

# GPUs	Ensemble member	Inference time [s]	T2m RMSE
			12h	24h
128	8	15.58	0.7507	0.8661
512	32	16.85	0.7289	0.8390
1536	96	16.98	0.7237	0.8328
3072	768	15.74	0.7226	0.8316

In summary, our results demonstrate the strengths of our innovations: They enable us to efficiently utilize the available resources (up to 20 480 GH200 GPUs) to generate large, accurate ensembles, while allowing for high flexibility regarding the modeling objectives. The combination of domain-, tensor-, uncertainty-, and data-parallelism in an orthogonal 4D parallelization scheme can be leveraged to:

• prioritize model size beyond the 1B parameter mark towards building bigger models, e.g., foundation models

• prioritize weight sample number towards maximizing ensemble generation throughput for better uncertainty quantification

• prioritize batch size towards maximizing forecast throughput for faster training and inference

Together, these aspects allow us to leverage the full potential of the Bayesian paradigm for smarter AI models in atmospheric forecasting.

Implications

Until now, training AI models for global atmospheric forecasting in a fully Bayesian framework has not been possible. While implementing the complex mathematical concepts

behind BNNs is already a challenge in itself, the associated computational training cost has ultimately been the prohibitive factor. We present BEAST, the first large-scale Bayesian Swin Transformer for atmospheric forecasting, able to quantify both, epistemic (model) and aleatoric (data) uncertainty. This breakthrough is enabled by our innovative 4D parallelization scheme, consisting of domain-, tensor-, uncertainty-, and data-parallelism.

We achieve state-of-the-art performance in both training speed and efficiency, peaking at 3.96 EFLOP/s on 20 480 NVIDIA GH200 GPUs on the JUPITER supercomputer. After training for nearly one million gradient updates, a 700-million-parameter model achieves predictive skill scores competitive with the state-of-the-art probabilistic AI model FGN and the IFS ensemble forecast. Yet, through our highly efficient domain-tensor-parallelism approach, BEAST is able to generate large ensembles significantly faster: a single 15-day forecast can be generated on 4 GH200 GPUs in less than 20 s. On 1536 nodes, we achieve a forecast generation throughput of 14.08 forecast/s with 96 ensemble members.

This innovation opens a new era for AI-based atmospheric sciences. Under a changing climate, fast and accurate uncertainty quantification in weather forecasting is becoming increasingly important. Existing approaches to uncertainty quantification in AI models follow methods from generative AI, using input randomness to produce ensemble forecasts; however, they neglect the uncertainty of the model itself, making the predictions overconfident [58]. In contrast, our Bayesian model formulation allows for a more principled estimation of the prediction certainty, leading to better calibrated uncertainty estimates. This calibration is crucial, especially for predicting extreme weather events such as heatwaves, floods, and severe storms. Decision makers depend on accurate estimates of the likelihood and severity of extreme events to prepare for possible outcomes and take preventative measures to reduce potential damages to people, infrastructures, and the economy. As demonstrated by our results, BEAST can deliver upon this promise by predicting extreme events, such as heat waves or tropical cyclones, with high accuracy. Our contribution thus directly addresses the challenges we are facing under a changing climate.

Next to these direct implications, our approach to enable fast and efficient training of large-scale BNNs also holds the potential to revolutionize Earth system sciences beyond atmospheric forecasting. While the lack of sufficiently large, well-curated datasets has thus far hindered true breakthroughs in AI-based models of Earth system components other than the atmosphere—such as the ocean—our Bayesian approach can overcome this barrier. With its capabilities to learn the underlying data distribution and to capture the associated uncertainty in its distributional model weights, a BNN is particularly well-suited to leverage small, high-quality datasets, e.g., from km-scale climate simulations [59] to build the next generation of climate emulators. Our approach to flexibly allocate GPU memory through highly efficient domain–tensor-parallelism puts us in a prime position to increase the size of input data items beyond the current resolutions towards multi-timestep forecasting and km-scale modeling.

Our breakthroughs were enabled by leveraging top-tier supercomputers with thousands of GPUs. Going forward, the results presented in this work mark only the beginning of exploring high-performance computing for BNNs. By introducing uncertainty parallelism as an orthogonal scaling axis, we have successfully overcome the current barriers in scaling out neural network training through plain data-parallelism, which is inevitably limited due to large-batch effects. Scaling out uncertainty parallelism goes beyond merely accelerating neural network training—it creates scientific value. For future work, we aim to leverage the power of current and future exascale systems by improving stability for lower precision computations and seeking out additional parallelization axes for scaling.

Acknowledgment

This project received access to the JUPITER supercomputer through the JUPITER Research and Early Access Program (JUREAP). JUPITER is funded by the EuroHPC Joint Undertaking, the German Federal Ministry of Research, Technology and Space, and the Ministry of Culture and Science of the German state of North Rhine-Westphalia. Computations on the LUMI supercomputer were enabled by resources provided by CSC - IT Center for Science, Finland. LUMI is co-funded by the EuroHPC Joint Undertaking and the LUMI Consortium.

We would like to express our gratitude to Christelle Piechurski from NVIDIA for the continuous support and encouragement to pursue this project from the beginning. We would like to thank Benedikt von St. Vieth and Damian Alvarez from JSC for their assistance and support, and Christian Schiffer for the exchanges on extreme-scale AI. We are also thankful to Jesse Harrison and Jarmo Mäkelä from CSC for supporting this submission, and to Gaurav Naithani and Billy Braithwaite from CSC for their continuous input and insightful discussions on the topics of Bayesian neural networks and FlashAttention.

This research is supported by the German Federal Ministry of Research, Technology and Space (BMFTR) under the 01LK2313A SMARTWEATHER21-SCC-2 grant, by the Carl Zeiss Foundation under the Breakthroughs 2025 “WOW – a World model of Our World” project, by the European High Performance Computing Joint Undertaking (JU) and the the German Federal Ministry of Research, Technology and Space (BMFTR), the Ministry of Culture and Science of North Rhine-Westphalia (MKW NRW) and the Hessian Ministry of Science and Research, Arts and Culture (HMWK) under grant agreement No 101250682 (JUPITER AI Factory, JAIF).