AERIS : Modèle des Systèmes Terrestres d'Argonne pour des Prévisions Fiables et Précises

papers.abstract

L'apprentissage automatique génératif offre de nouvelles opportunités pour mieux comprendre la dynamique complexe du système terrestre. Les méthodes récentes basées sur la diffusion abordent les biais spectraux et améliorent l'étalonnage des ensembles dans la prévision météorologique par rapport aux méthodes déterministes, mais se sont jusqu'à présent avérées difficiles à mettre à l'échelle de manière stable à haute résolution. Nous présentons AERIS, un transformeur de diffusion Swin au niveau des pixels de 1,3 à 80 milliards de paramètres pour combler cette lacune, et SWiPe, une technique généralisable qui combine le parallélisme de fenêtres avec le parallélisme de séquence et de pipeline pour partitionner les transformeurs basés sur des fenêtres sans coût de communication supplémentaire ni augmentation de la taille globale du lot. Sur Aurora (10 080 nœuds), AERIS maintient 10,21 ExaFLOPS (précision mixte) et une performance de pointe de 11,21 ExaFLOPS avec une taille de patch de 1 fois 1 sur le jeu de données ERA5 à 0,25{\deg}, atteignant une efficacité de mise à l'échelle faible de 95,5 % et une efficacité de mise à l'échelle forte de 81,6 %. AERIS surpasse l'IFS ENS et reste stable sur des échelles saisonnières jusqu'à 90 jours, mettant en évidence le potentiel des modèles de diffusion à milliards de paramètres pour la prévision météorologique et climatique.

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 : Modèle des Systèmes Terrestres d'Argonne pour des Prévisions Fiables et Précises

AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions

papers.abstract

Support