AERIS: Argonne Aardesysteemmodel voor Betrouwbare en Vaardige Voorspellingen

Samenvatting

Generatief machinaal leren biedt nieuwe mogelijkheden om complexe dynamiek van het aardesysteem beter te begrijpen. Recente op diffusie gebaseerde methoden pakken spectrale biases aan en verbeteren ensemblekalibratie in weersvoorspellingen in vergelijking met deterministische methoden, maar hebben tot nu toe moeite gehad om stabiel te schalen op hoge resoluties. We introduceren AERIS, een Swin-diffusie transformer op pixelniveau met 1,3 tot 80B parameters om deze kloof te overbruggen, en SWiPe, een generaliseerbare techniek die vensterparallelisme combineert met sequentie- en pijplijnparallelisme om venstergebaseerde transformers te verdelen zonder extra communicatiekosten of een verhoogde globale batchgrootte. Op Aurora (10.080 nodes) behoudt AERIS 10,21 ExaFLOPS (gemengde precisie) en een piekprestatie van 11,21 ExaFLOPS met een patchgrootte van 1 keer 1 op de 0,25{\deg} ERA5 dataset, waarbij een zwakke schaalefficiëntie van 95,5% en een sterke schaalefficiëntie van 81,6% wordt bereikt. AERIS overtreft de IFS ENS en blijft stabiel op seizoensschalen tot 90 dagen, wat het potentieel van diffusiemodellen met miljarden parameters voor weers- en klimaatvoorspelling benadrukt.

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: Argonne Aardesysteemmodel voor Betrouwbare en Vaardige Voorspellingen

AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions

Samenvatting

Support