Software, scaling & compression★ Highlight

AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions

Väinö Hatanpää, Eugene Ku, Jason Stock, Murali Emani, Sam Foreman, Chunyong Jung, Sandeep Madireddy, Tung Nguyen, Varuni Sastry, Ray A. O. Sinurat, Sam Wheeler, Huihuo Zheng et al.

Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis · 16 September 2025 · Allen Institute for AI (Ai2, ACE), Argonne National Laboratory, University of California, Los Angeles (UCLA), University of Chicago

Paper PDF

Why it matters. Scales diffusion transformers to 80B parameters at 0.25° with exaflop performance, beating IFS ENS and stable to seasonal scales.

In one sentence. AERIS is a 1.3 to 80B parameter Swin diffusion transformer reaching 10.2 ExaFLOPS on Aurora, stable to 90-day forecasts and outperforming IFS ENS.

Abstract

Generative machine learning offers new opportunities to better understand complex Earth system dynamics. Recent diffusion-based methods address spectral biases and improve ensemble calibration in weather forecasting compared to deterministic methods, yet have so far proven difficult to scale stably at high resolutions. We introduce AERIS, a 1.3 to 80B parameter pixel-level Swin diffusion transformer to address this gap, and SWiPe, a generalizable technique that composes window parallelism with sequence and pipeline parallelism to shard window-based transformers without added communication cost or increased global batch size. On Aurora (10,080 nodes), AERIS sustains 10.21 ExaFLOPS (mixed precision) and a peak performance of 11.21 ExaFLOPS with $1 \times 1$ patch size on the 0.25° ERA5 dataset, achieving 95.5% weak scaling efficiency, and 81.6% strong scaling efficiency. AERIS outperforms the IFS ENS and remains stable on seasonal scales to 90 days, highlighting the potential of billion-parameter diffusion models for weather and climate prediction.

Labels

Labels and summary are generated by Claude and reviewed weekly. Wrong? Edit this entry.