[論文レビュー] SciAI4Industry -- Solving PDEs for industry-scale problems with deep learning

Philipp Witte, Russell J. Hewett+3 more|arXiv (Cornell University)|Nov 23, 2022

Reservoir Engineering and Simulation Methods被引用数 4

ひとこと要約

本論文は、クラウド上で大規模な学習データを生成するためのJuliaベースの分散API「SciAI4Industry」を紹介する。このAPIは、ドメイン分割を用いたモデル並列ディープラーニングと組み合わせられ、産業スケールのPDE問題に対する3次元ニューラルネットワークの学習を可能にする。200万セルのCO₂流れシミュレーションを実証的に適用した結果、従来のソルバーと比較して推論が5桁速くなり、コストは3,200分の1にまで低下し、並列効率は90%以上を達成した。

ABSTRACT

Solving partial differential equations with deep learning makes it possible to reduce simulation times by multiple orders of magnitude and unlock scientific methods that typically rely on large numbers of sequential simulations, such as optimization and uncertainty quantification. Two of the largest challenges of adopting scientific AI for industrial problem settings is that training datasets must be simulated in advance and that neural networks for solving large-scale PDEs exceed the memory capabilities of current GPUs. We introduce a distributed programming API in the Julia language for simulating training data in parallel on the cloud and without requiring users to manage the underlying HPC infrastructure. In addition, we show that model-parallel deep learning based on domain decomposition allows us to scale neural networks for solving PDEs to commercial-scale problem settings and achieve above 90% parallel efficiency. Combining our cloud API for training data generation and model-parallel deep learning, we train large-scale neural networks for solving the 3D Navier-Stokes equation and simulating 3D CO2 flow in porous media. For the CO2 example, we simulate a training dataset based on a commercial carbon capture and storage (CCS) project and train a neural network for CO2 flow simulation on a 3D grid with over 2 million cells that is 5 orders of magnitudes faster than a conventional numerical simulator and 3,200 times cheaper.

研究の動機と目的

• 産業現場における科学的AIのための大規模な学習データセットを生成する課題に取り組む。従来のシミュレータは処理が遅く、HPCリソースを必要とするためである。
• 1つのGPUのメモリ容量を超える大規模3次元PDE問題のディープラーニングモデルを学習する際のGPUメモリ制限を克服するため、ドメイン分割を用いたモデル並列化を実現する。
• 実世界の産業問題、例えば炭素回収・貯留（CCS）における3次元CO₂流れのシミュレーションに適した大規模ディープラーニング代替モデルの実現可能性を示す。
• 最適化や不確実性評価のような、多数のシミュレーションを要する応用分野において、コスト効率よく高性能な推論を可能にする。
• Juliaでの高水準でクラウドネイティブなプログラミングインターフェースを提供することで、HPCレベルのシミュレーション能力を一般化する。

提案手法

• ユーザーがクラスタを管理しなくても、HPCインfraストラクチャを抽象化した高水準なJulia APIを開発し、並列処理で大規模な学習データセットを自動生成可能にする。
• 空間的・時間的ドメインを複数のGPUに分割するドメイン分割を用いたモデル並列ディープラーニングを実装し、大規模なPDE問題をメモリ制限内に収める。
• 時間依存PDEを解くためにフォーリエニューラルオペレータ（FNO）アーキテクチャを採用し、4次元FFTおよびiFFT演算における効率的な通信を実現するため、テンソルをGPU間で分割する。
• 1ノード内でのGPU間の高帯域幅NVLink接続を活用し、学習および推論中に90%を超える高い並列効率を維持する。
• クラウドVM（例：Azure ND96amsr）と分散学習パイプラインを統合し、コストと管理負荷を最小限に抑えつつ、学習および推論のスケーリングを可能にする。
• 将来的な大規模バッチサイズおよびマルチノードスケーリングに対応するため、ハイブリッド並列化の原則を適用するが、現在の実装ではバッチサイズ2のモデル並列化のみを採用している。

実験結果

リサーチクエスチョン

• RQ1 高水準でクラウドネイティブなプログラミングAPIは、HPCインfraストラクチャを管理する必要なく、科学的AIのための大規模な学習データセットの生成を簡素化できるか？
• RQ2 ドメイン分割に基づくモデル並列ディープラーニングは、1つのGPUのメモリ容量を超える大規模3次元PDE問題を解くのに実用的か？
• RQ3 1ノード内の複数GPUにスケーリングする際、高い並列効率（例：90%以上）を達成できるか？
• RQ4 得られたディープラーニング代替モデルは、実際の産業スケールのPDE問題において、従来の数値ソルバーと比較して処理速度とコストの面で優れているか？
• RQ5 このフレームワークは、不確実性評価や最適化のような高スルーレートな科学的ワークフローを、産業応用でどの程度可能にするか？

主な発見

• 提案されたクラウドベースのAPIにより、HPCクラスタやクラウドインfraストラクチャを管理する必要なく、大規模PDEのためのスケーラブルかつ自動化された学習データ生成が可能になった。
• ドメイン分割を用いたモデル並列FNO学習は、最大8枚のA100 GPUで90%を超える並列効率を達成し、200万セルを超える3次元グリッドでの学習が可能となった。
• 学習済みFNOモデルは、3次元CO₂流れのシミュレーションを従来のOPMソルバーと比較して5桁速く（1回のシミュレーションあたり0.12秒 vs. 6.8時間）実行可能となった。
• クラウドVMコストを考慮した場合、FNOベースのシミュレーションは1回あたりOPMソルバーと比較して3,200倍安価（1回あたり$0.11 vs. $3.4）となった。
• FNO代替モデルはたった1,848回のシミュレーションでコスト面でのコスト・メリットを回収し、ウェル配置最適化のような高スルーレート応用分野において経済的に実用的となった。
• 静的なPDEを解く場合、3次元で最大512³、2次元で最大12,000²のグリッドでの学習が可能であり、商業スケールの問題に対するスケーラビリティを示した。