AERIS: Модель земных систем Аргоннской национальной лаборатории для надежных и точных прогнозов
AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions
September 16, 2025
Авторы: Väinö Hatanpää, Eugene Ku, Jason Stock, Murali Emani, Sam Foreman, Chunyong Jung, Sandeep Madireddy, Tung Nguyen, Varuni Sastry, Ray A. O. Sinurat, Sam Wheeler, Huihuo Zheng, Troy Arcomano, Venkatram Vishwanath, Rao Kotamarthi
cs.AI
Аннотация
Генеративное машинное обучение открывает новые возможности для более глубокого понимания сложной динамики земной системы. Современные методы на основе диффузии устраняют спектральные смещения и улучшают калибровку ансамблей в прогнозировании погоды по сравнению с детерминированными методами, однако до сих пор их масштабирование на высоких разрешениях оставалось проблематичным. Мы представляем AERIS — трансформер на основе Swin-диффузии с 1,3 до 80 миллиардов параметров на уровне пикселей, который решает эту проблему, а также SWiPe — универсальную технику, сочетающую параллелизм окон с последовательным и конвейерным параллелизмом для распределения трансформеров на основе окон без дополнительных затрат на коммуникацию или увеличения глобального размера пакета. На системе Aurora (10 080 узлов) AERIS поддерживает 10,21 эксафлопс (смешанная точность) и пиковую производительность 11,21 эксафлопс с размером патча 1×1 на наборе данных ERA5 с разрешением 0,25°, достигая 95,5% эффективности слабого масштабирования и 81,6% эффективности сильного масштабирования. AERIS превосходит IFS ENS и остается стабильным на сезонных масштабах до 90 дней, демонстрируя потенциал моделей диффузии с миллиардами параметров для прогнозирования погоды и климата.
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.