AERIS: Modello dei Sistemi Terrestri di Argonne per Previsioni Affidabili e Precise

Abstract

L'apprendimento automatico generativo offre nuove opportunità per comprendere meglio le dinamiche complesse del sistema terrestre. I recenti metodi basati sulla diffusione affrontano i bias spettrali e migliorano la calibrazione degli ensemble nelle previsioni meteorologiche rispetto ai metodi deterministici, ma finora si sono dimostrati difficili da scalare in modo stabile ad alte risoluzioni. Introduciamo AERIS, un trasformatore di diffusione Swin a livello di pixel con parametri da 1,3 a 80 miliardi, per colmare questa lacuna, e SWiPe, una tecnica generalizzabile che combina il parallelismo a finestre con il parallelismo di sequenza e pipeline per suddividere i trasformatori basati su finestre senza costi di comunicazione aggiuntivi o aumento della dimensione globale del batch. Su Aurora (10.080 nodi), AERIS mantiene 10,21 ExaFLOPS (precisione mista) e una performance di picco di 11,21 ExaFLOPS con una dimensione di patch 1x1 sul dataset ERA5 a 0,25{\deg}, raggiungendo un'efficienza di scaling debole del 95,5% e un'efficienza di scaling forte dell'81,6%. AERIS supera l'IFS ENS e rimane stabile su scale stagionali fino a 90 giorni, evidenziando il potenziale dei modelli di diffusione con miliardi di parametri per la previsione meteorologica e climatica.

English

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{\deg} 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.

AERIS: Modello dei Sistemi Terrestri di Argonne per Previsioni Affidabili e Precise

AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions

Abstract

Support