AERIS: Argonne's Earth Systems Model

Sam Foreman 2025-10-08

• 🌎 AERIS
• High-Level Overview of AERIS
• Contributions
• Model Overview
• Windowed Self-Attention
• Model Architecture: Details
• Issues with the Deterministic Approach
• Transitioning to a Probabilistic Model
• Sequence-Window-Pipeline Parallelism SWiPe
• Aurora
• AERIS: Scaling Results
• Hurricane Laura
• S2S: Subsseasonal-to-Seasonal Forecasts
• Seasonal Forecast Stability
• Next Steps
• References
• Extras
  • Overview of Diffusion Models
  • Diffusion Model: Forward Process
• Acknowledgements

🌎 AERIS

Sequence-Window-Pipeline Parallelism SWiPe

• SWiPe is a novel parallelism strategy for Swin-based Transformers
• Hybrid 3D Parallelism strategy, combining:
  • Sequence parallelism (SP)
  • Window parallelism (WP)
  • Pipeline parallelism (PP)

Figure 6

Figure 7: SWiPe Communication Patterns

Aurora

Table 3: Aurora¹ Specs

Property	Value
Racks	166
Nodes	10,624
XPUs2	127,488
CPUs	21,248
NICs	84,992
HBM	8 PB
DDR5c	10 PB

Figure 8: Aurora: Fact Sheet.

AERIS: Scaling Results

Figure 9: AERIS: Scaling Results

• 10 EFLOPs (sustained) @ 120,960 GPUs
• See (Hatanpää et al. (2025)) for additional details
• arXiv:2509.13523

Hurricane Laura

Figure 10: Hurricane Laura tracks (top) and intensity (bottom). Initialized 7(a), 5(b) and 3(c) days prior to 2020-08-28T00z.

S2S: Subsseasonal-to-Seasonal Forecasts

Seasonal Forecast Stability

Figure 11: S2S Stability: (a) Spring barrier El Niño with realistic ensemble spread in the ocean; (b) qualitatively sharp fields of SST and Q700 predicted 90 days in the future from the

closest ensemble member to the ERA5 in (a); and (c) stable Hovmöller diagrams of U850 anomalies (climatology removed; m/s), averaged between 10°S and 10°N, for a 90-day rollout.

Next Steps

• Swift: Swift, a single-step consistency model that, for the first time, enables autoregressive finetuning of a probability flow model with a continuous ranked probability score (CRPS) objective

References

1. What are Diffusion Models? | Lil’Log
2. Step by Step visual introduction to Diffusion Models. - Blog by Kemal Erdem
3. Understanding Diffusion Models: A Unified Perspective

Extras

Overview of Diffusion Models

Goal: We would like to (efficiently) draw samples $x_{i}$ from a (potentially unknown) target distribution $q(\cdot)$.

• Given $x_{0} \sim q(x)$, we can construct a forward diffusion process by gradually adding noise to $x_{0}$ over $T$ steps: $x_{0} \rightarrow \left\{x_{1}, \ldots, x_{T}\right\}$.

  • Step sizes $\beta_{t} \in (0, 1)$ controlled by a variance schedule $\{\beta\}_{t=1}^{T}$, with:

    $$\begin{aligned} q(x_{t}|x_{t-1}) = \mathcal{N}(x_{t}; \sqrt{1-\beta_{t}} x_{t-1}, \beta_{t} I) \\ q(x_{1:T}|x_{0}) = \prod_{t=1}^{T} q(x_{t}|x_{t-1}) \end{aligned}$$

Diffusion Model: Forward Process

• Introduce:

  • $\alpha_{t} \equiv 1 - \beta_{t}$
  • $\bar{\alpha}_{t} \equiv \prod_{s=1}^{T} \alpha_{s}$

  We can write the forward process as:

  $q(x*{1}|x*{0}) = \mathcal{N}(x*{1}; \sqrt{\bar{\alpha}*{1}} x*{0}, (1-\bar{\alpha}*{1}) I)$

• We see that the mean $\mu_{t} = \sqrt{\alpha_{t}} x_{t-1} = \sqrt{\bar{\alpha}_{t}} x_{0}$

Acknowledgements

This research used resources of the Argonne Leadership Computing Facility at Argonne National Laboratory, which is supported by the Office of Science of the U.S. Department of Energy, Office of Science, under contract number DE-AC02-06CH11357.