AERIS: Modelo de Sistemas Terrestres de Argonne para Predicciones Confiables y Precisas

Resumen

El aprendizaje automático generativo ofrece nuevas oportunidades para comprender mejor la dinámica compleja del sistema terrestre. Los métodos recientes basados en difusión abordan los sesgos espectrales y mejoran la calibración de conjuntos en la predicción meteorológica en comparación con los métodos determinísticos, aunque hasta ahora han demostrado ser difíciles de escalar de manera estable a altas resoluciones. Presentamos AERIS, un transformador de difusión Swin a nivel de píxel con parámetros de 1.3 a 80B para abordar esta brecha, y SWiPe, una técnica generalizable que combina el paralelismo de ventanas con el paralelismo de secuencias y tuberías para dividir transformadores basados en ventanas sin costos adicionales de comunicación o aumento del tamaño global del lote. En Aurora (10,080 nodos), AERIS mantiene 10.21 ExaFLOPS (precisión mixta) y un rendimiento máximo de 11.21 ExaFLOPS con un tamaño de parche de 1x1 en el conjunto de datos ERA5 de 0.25{\deg}, logrando una eficiencia de escalado débil del 95.5% y una eficiencia de escalado fuerte del 81.6%. AERIS supera al IFS ENS y se mantiene estable en escalas estacionales hasta 90 días, destacando el potencial de los modelos de difusión con miles de millones de parámetros para la predicción meteorológica y climática.

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: Modelo de Sistemas Terrestres de Argonne para Predicciones Confiables y Precisas

AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions

Resumen

Support