[論文レビュー] FourCastNet: Accelerating Global High-Resolution Weather Forecasting using Adaptive Fourier Neural Operators

Thorsten Kurth, Shashank Subramanian+7 more|arXiv (Cornell University)|Aug 8, 2022

Meteorological Phenomena and Simulations被引用数 24

ひとこと要約

FourCastNet は Adaptive Fourier Neural Operators を用いて地球系ダイナミクスを模倣し、グローバルな高解像度の天気予報を実現。従来の NWP と比べて推論が約80,000倍高速で、精度も高い。混合精度での訓練は3,808 GPUにスケールし、140.8 petaFLOPS。

ABSTRACT

Extreme weather amplified by climate change is causing increasingly devastating impacts across the globe. The current use of physics-based numerical weather prediction (NWP) limits accuracy due to high computational cost and strict time-to-solution limits. We report that a data-driven deep learning Earth system emulator, FourCastNet, can predict global weather and generate medium-range forecasts five orders-of-magnitude faster than NWP while approaching state-of-the-art accuracy. FourCast-Net is optimized and scales efficiently on three supercomputing systems: Selene, Perlmutter, and JUWELS Booster up to 3,808 NVIDIA A100 GPUs, attaining 140.8 petaFLOPS in mixed precision (11.9%of peak at that scale). The time-to-solution for training FourCastNet measured on JUWELS Booster on 3,072GPUs is 67.4minutes, resulting in an 80,000times faster time-to-solution relative to state-of-the-art NWP, in inference. FourCastNet produces accurate instantaneous weather predictions for a week in advance, enables enormous ensembles that better capture weather extremes, and supports higher global forecast resolutions.

研究の動機と目的

• 気候変動による極端な気象の増加に伴い、より高速で高解像度の天気予報の必要性を動機づける。
• 従来の数値気象予報（NWP）モデルのスループットを超えるデータ駆動型地球系エミュレータを開発する。
• GPU加速を用いた複数の超級コンピュータでのスケーラブルな訓練と推論性能を示す。
• モデルが大規模なアンサンブルと高解像度の予報を効率的に生成できることを示す。

提案手法

• Adaptive Fourier Neural Operator (AFNO) アーキテクチャを採用し、時間を跨いだ地球系状態テンソル間の写像をモデル化する。
• AFNO ブロック内でのスペクトル（FFT ベース）空間混合を用いて、高解像度の全球予報を可能にする。
• native 0.25度解像度で ERA5 再分析データを用い、20 の予測変数を用いて、事前学習と微調整フェーズを経て6時間ごとのステップ（t -> t+6h）を予測する。
• Throughputを最大化しCPUオーバーヘッドを最小化するために、NCCL通信と CUDA グラフを用いたモデルとデータ並列性（特徴とデータ並列性）を実装する。
• データを NVIDIA DALI で前処理し、JIT フュージョンでカーネル起動オーバーヘッドを削減する。AFNO トランスフォーマーブロックには FFT ベースのスペクトル畳み込みと MLP を採用する。

実験結果

リサーチクエスチョン

• RQ1 AFNO ベースのアーキテクチャは高全球解像度において最先端の天気予報精度に匹敵するか、あるいは近づくか。
• RQ2 FourCastNet はリーダーシップ級 HPC システム上でどれだけ速く、スケーラブルに訓練・推論できるか。
• RQ3 全球天気予報とアンサンブルの大規模 DL サロゲートのエネルギーとスループットにはどのような影響があるか。
• RQ4 モデルはリードタイムが1週間程度までの大規模アンサンブルと多ステップ予報を、適切なスキルでサポートできるか。

主な発見

• 複数の HPC システムで訓練すると 3,808 個の NVIDIA A100 GPU にスケールし、140.8 petaFLOPS（ピークの11.9%）を達成。
• JUWELS Booster での事前学習は 3,072 GPU で 67.4 分、最先端の NWP（IFS）に対して推論の time-to-solution を約80,000倍高速化。
• 推論は最大1週間程度の瞬時予報の正確性を提供し、大規模アンサンブルをサポートする。100-member の 24-hour アンサンブル予報は Selene ノードの 8 A100 GPU で 12.41 ノード秒、IFS（L91 18 km）では 984,000 ノード秒。
• エネルギー効率は飛躍的に高く、FourCastNet アンサンブルのエネルギーは IFS アンサンブルの約10,000倍低いと推定。
• FourCastNet は prior DL 天気サロゲートより高解像度（従来の DL モデルの 8 倍）と精度を達成し、熱帯低気圧や大気河川などの微細現象を解像できる。