[論文レビュー] Model-Parallel Fourier Neural Operators as Learned Surrogates for Large-Scale Parametric PDEs

Thomas J. Grady, Rishi Khan+6 more|arXiv (Cornell University)|Apr 4, 2022

Seismic Imaging and Inversion Techniques被引用数 9

ひとこと要約

本稿では、ドメイン分割を用いたモデル並列型フーリエニューラルオペレーター（FNO）フレームワークを提案し、大規模な3次元および4次元パrametric PDEにスケーリングすることを可能にし、26億以上の変数を有する問題における推論および学習を実現する。Perlmutter上で最大512個のA100 GPUを用いた弱スケーリングを達成し、これまでで最大のFNO推論および勾配計算を実現した。

ABSTRACT

Fourier neural operators (FNOs) are a recently introduced neural network architecture for learning solution operators of partial differential equations (PDEs), which have been shown to perform significantly better than comparable deep learning approaches. Once trained, FNOs can achieve speed-ups of multiple orders of magnitude over conventional numerical PDE solvers. However, due to the high dimensionality of their input data and network weights, FNOs have so far only been applied to two-dimensional or small three-dimensional problems. To remove this limited problem-size barrier, we propose a model-parallel version of FNOs based on domain-decomposition of both the input data and network weights. We demonstrate that our model-parallel FNO is able to predict time-varying PDE solutions of over 2.6 billion variables on Perlmutter using up to 512 A100 GPUs and show an example of training a distributed FNO on the Azure cloud for simulating multiphase CO$_2$ dynamics in the Earth's subsurface.

研究の動機と目的

• 高次元入力および重みの制約により、大規模な3次元および4次元パラメトリックPDEを扱う際のFNOのスケーラビリティ制限を解消すること。
• 分散GPUリソースを活用して、26億変数を超える問題におけるFNOの効率的推論および学習を可能にすること。
• ドメイン分割に基づくモデル並列型FNOフレームワークを構築し、FNOのインダクティブバイアスおよび性能を保持したまま複数GPUにスケーリングすること。
• 地下CO2プレーンシミュレーションや逆問題といった複雑な実世界応用に分散FNOを効果的に適用できることを実証すること。
• 従来の数値ソルバーよりも実行不可能であった大規模最適化および不確実性評価タスクに適した高速で微分可能なPDEソルバーを提供すること。

提案手法

• 複数のGPUにわたって入力データおよびネットワーク重みを分割するドメイン分割戦略を実装し、モデル並列性を実現する。
• FNOの線形代数的定式化を用いて、パーティション間で数学的整合性を保つ分散スペクトル畳み込み演算を導出する。
• 前方および逆伝搬における効率的なGPU間通信を実現するため、CUDA対応MPIを活用し、通信オーバーヘッドを最小限に抑える。
• PyTorchのDistDLライブラリを用いて、元のFNOアーキテクチャ（4個のスペクトルブロック、埋め込み次元20）を分散環境に適応させる。
• 空間次元を512×512×512まで、時間次元を10,240ステップまでスケーリングし、Perlmutter上で弱スケーリング実験を実施する。
• 自動微分を用いた分散学習および推論を実行し、勾配計算をサポートすることで逆問題ワークフローを実現する。

実験結果

リサーチクエスチョン

• RQ1 モデル並列性を用いることで、64³空間制限を超えた大規模な3次元および4次元パラメトリックPDEにFNOをスケーリングできるか？
• RQ2 提案されたドメイン分割FNOは、最大512個のA100 GPUにわたる弱スケーリング実験において、どのように性能を示すか？
• RQ3 従来のソルバーよりも、分散FNOは大規模PDEの推論および勾配計算をどの程度高速化できるか？
• RQ4 モデル並列型FNOフレームワークは、高次元出力を持つ実世界の地下CO2移動シミュレーションに効果的に適用できるか？
• RQ5 現代のHPCシステム（例：Perlmutter）上で、分散FNOの問題サイズおよびスケーラビリティの性能限界はどの程度か？

主な発見

• モデル並列型FNOは、Perlmutter上で最大512個のA100 GPUまで弱スケーリングを達成し、512×512×512×20の空間的・時間的次元を持つ問題のスケーリング可能性を実証した。
• 本フレームワークは、26億以上の変数を有する問題における推論および勾配計算を可能にし、これまでで最大のFNO応用を実現した。
• FNOのアラウンド・トゥ・アラウンド特性のため、わずかに不完全なスケーリングを示したが、全スケールで高い性能を維持した。
• 本手法は、Azureクラウド上で分散FNOを地下CO2動的挙動のシミュレーションに学習させることに成功し、実世界応用の有効性を検証した。
• 逆問題や最適化に適した高速で微分可能なPDEソリューションを実現し、従来のソルバーよりも推論時間が桁違いに速くなった。
• 本研究は、分散オペレーター学習を用いた実空間的ボリュームデータにおける連成逆問題および統計的問題の解決に基盤を築いた。