AERIS：信頼性と精度を備えたアルゴンヌ地球システムモデル

要旨

生成機械学習は、複雑な地球システムのダイナミクスをより深く理解するための新たな可能性を提供します。最近の拡散ベースの手法は、決定論的手法と比較して気象予報におけるスペクトルバイアスを解消し、アンサンブルキャリブレーションを改善していますが、これまで高解像度での安定したスケーリングが困難でした。このギャップを埋めるため、我々は1.3から800億パラメータのピクセルレベルSwin拡散トランスフォーマーであるAERISを導入し、ウィンドウ並列処理をシーケンス並列処理とパイプライン並列処理と組み合わせることで、追加の通信コストやグローバルバッチサイズの増加なしにウィンドウベースのトランスフォーマーを分割する汎用技術SWiPeを提案します。Aurora（10,080ノード）上で、AERISは0.25度ERA5データセットにおいて1×1パッチサイズで10.21エクサフロップス（混合精度）を維持し、ピーク性能は11.21エクサフロップスを達成し、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.

AERIS：信頼性と精度を備えたアルゴンヌ地球システムモデル

AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions

要旨

Support