AERIS: 신뢰할 수 있고 정확한 예측을 위한 아르곤 지구 시스템 모델

초록

생성적 기계 학습은 복잡한 지구 시스템 역학을 더 잘 이해할 수 있는 새로운 기회를 제공합니다. 최근 확산 기반 방법은 결정론적 방법에 비해 기상 예측에서 스펙트럼 편향을 해결하고 앙상블 보정을 개선했지만, 고해상도에서 안정적으로 확장하기는 어려운 것으로 나타났습니다. 우리는 이러한 격차를 해결하기 위해 1.3에서 80B 파라미터 규모의 픽셀 수준 Swin 확산 트랜스포머인 AERIS와, 윈도우 병렬성을 시퀀스 및 파이프라인 병렬성과 결합하여 추가 통신 비용이나 글로벌 배치 크기 증가 없이 윈도우 기반 트랜스포머를 분할할 수 있는 일반화 가능한 기술인 SWiPe를 소개합니다. Aurora(10,080 노드)에서 AERIS는 0.25{\deg} ERA5 데이터셋에 대해 10.21 ExaFLOPS(혼합 정밀도)를 유지하고 1x1 패치 크기에서 11.21 ExaFLOPS의 최고 성능을 달성하며, 약한 확장 효율성 95.5%, 강한 확장 효율성 81.6%를 보였습니다. AERIS는 IFS ENS를 능가하며 계절적 규모에서 90일까지 안정적으로 유지되어, 기상 및 기후 예측을 위한 10억 파라미터 확산 모델의 잠재력을 강조합니다.

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: 신뢰할 수 있고 정확한 예측을 위한 아르곤 지구 시스템 모델

AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions

초록

Support