AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Hatanpää, Väinö, Ku, Eugene, Stock, Jason, Emani, Murali, Foreman, Sam, Jung, Chunyong, Madireddy, Sandeep, Nguyen, Tung, Sastry, Varuni, Sinurat, Ray A. O., Wheeler, Sam, Zheng, Huihuo, Arcomano, Troy, Vishwanath, Venkatram, Kotamarthi, Rao
Format: Preprint
Publié: 2025
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866908543314034688
author Hatanpää, Väinö
Ku, Eugene
Stock, Jason
Emani, Murali
Foreman, Sam
Jung, Chunyong
Madireddy, Sandeep
Nguyen, Tung
Sastry, Varuni
Sinurat, Ray A. O.
Wheeler, Sam
Zheng, Huihuo
Arcomano, Troy
Vishwanath, Venkatram
Kotamarthi, Rao
author_facet Hatanpää, Väinö
Ku, Eugene
Stock, Jason
Emani, Murali
Foreman, Sam
Jung, Chunyong
Madireddy, Sandeep
Nguyen, Tung
Sastry, Varuni
Sinurat, Ray A. O.
Wheeler, Sam
Zheng, Huihuo
Arcomano, Troy
Vishwanath, Venkatram
Kotamarthi, Rao
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2509_13523
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions
Hatanpää, Väinö
Ku, Eugene
Stock, Jason
Emani, Murali
Foreman, Sam
Jung, Chunyong
Madireddy, Sandeep
Nguyen, Tung
Sastry, Varuni
Sinurat, Ray A. O.
Wheeler, Sam
Zheng, Huihuo
Arcomano, Troy
Vishwanath, Venkatram
Kotamarthi, Rao
Machine Learning
Distributed, Parallel, and Cluster Computing
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.
title AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions
topic Machine Learning
Distributed, Parallel, and Cluster Computing
url https://arxiv.org/abs/2509.13523