[論文レビュー] Highly-scalable, physics-informed GANs for learning solutions of stochastic PDEs

Liu Yang, Sean Treichler+9 more|arXiv (Cornell University)|Oct 29, 2019

Model Reduction and Neural Networks参考文献 29被引用数 16

ひとこと要約

本稿では、ハンフォード・サイトにおける地下水流モデリングにおける高次元確率的PDEを解くために、非常にスケーラブルで物理的制約を組み込んだGANフレームワークを提案する。ドメイン分割、マルチGPU並列処理、物理的制約付き学習を活用することで、27,500枚のV100 GPUで93.1%のスケーリング効率を達成し、1207 PF/sの持続的半精度性能を記録した。これは、1,000次元までの確率的次元を有する混合SPDE問題に対して、初のエクサスケール対応GANである。

ABSTRACT

Uncertainty quantification for forward and inverse problems is a central challenge across physical and biomedical disciplines. We address this challenge for the problem of modeling subsurface flow at the Hanford Site by combining stochastic computational models with observational data using physics-informed GAN models. The geographic extent, spatial heterogeneity, and multiple correlation length scales of the Hanford Site require training a computationally intensive GAN model to thousands of dimensions. We develop a hierarchical scheme for exploiting domain parallelism, map discriminators and generators to multiple GPUs, and employ efficient communication schemes to ensure training stability and convergence. We developed a highly optimized implementation of this scheme that scales to 27,500 NVIDIA Volta GPUs and 4584 nodes on the Summit supercomputer with a 93.1% scaling efficiency, achieving peak and sustained half-precision rates of 1228 PF/s and 1207 PF/s.

研究の動機と目的

• ハンフォード・サイトにおける地下水流モデリングにおける不確実性評価を扱う。これは、広範な汚染とデータ不足を伴う、複雑で高次元の問題である。
• 空間的不均一性と多スケール相関長が原因となる確率的PDE（SPDE）の次元の呪いを克服する。
• 最大1,000次元の確率的次元を扱えるスケーラブルで物理的制約を組み込んだGAN（PI-GAN）アーキテクチャを開発する。
• ドメイン分割と最適化された通信を用いて、大規模HPCシステム上で効率的な分散トレーニングを可能にする。
• 1,000～10,000口の井戸からの観測データを統合し、圧力ヘッドと透水係数場を確率過程として推定する統一された混合問題定式化を実現する。

提案手法

• 1つのジェネレータと数百のディスクラミネータを複数のGPUにマッピングする階層的ドメイン分割方式を提案し、スケーラブルなトレーニングを可能にする。
• 物理的制約を直接GANトレーニングプロセスに組み込む物理的制約付き損失関数を実装し、生成された解の物理的整合性を保証する。
• 勾配ペナルティを用いたWGAN-GPを採用し、高次元確率過程におけるトレーニングの安定性と分布の忠実度を向上させる。
• 非同期的かつオーバラップスケジューリングを用いた通信最適化パイプラインを設計し、無駄な待機時間を最小限に抑え、GPUの利用効率を最大化する。
• TensorFlowでジェネレータおよびディスクラミネータアーキテクチャを最適化し、V100 GPUノードあたり59.3 TF/sおよび38.6 TF/sのピーク性能を達成する。
• 動的負荷分散を備えたフリーラントレーニングスケジュールを採用し、4,584ノードのSummitシステム全体で高い並列効率を維持する。

実験結果

リサーチクエスチョン

• RQ1 物理的制約付きGANフレームワークは、1,000以上の確率的次元を有する地下水流に由来する高次元確率的PDEを効果的にモデル化できるか？
• RQ2 1つのグローバルジェネレータと複数のローカルディスクラミネータを有するGANにドメイン分割を効果的に適用することで、エクサスケールトレーニングが可能になるか？
• RQ3 半精度演算を用いて27,500枚のGPUでPI-GANをトレーニングする際、どの程度のスケーリング効率と性能が達成できるか？
• RQ4 数千口の井戸からの観測データの統合は、推定された透水係数場と圧力ヘッド分布の精度と信頼性をどのように向上させるか？
• RQ5 大規模分散環境における量子化効果と不均衡なGPUワークロードの影響を受けても、提案されたGANアーキテクチャはトレーニングの安定性と収束性を維持できるか？

主な発見

• 提案されたPI-GANフレームワークは、4,584ノードのSummit上で93.1%のスケーリング効率を達成し、極めて大規模なスケールにおけるGANトレーニングの近似最適な強スケーリングを示した。
• 27,500枚のV100 GPU上で、持続的半精度性能1,207 PF/s、ピーク性能1,228 PF/sを達成した。これは、SPDE問題に対して初のエクサスケール対応GANである。
• 本手法は、最大1,000次元の確率的次元を有する混合SPDE問題を効果的に解き、従来のGANベースのSPDEソルバーをはるかに凌駕した。
• フリーラントレーニングスケジュールが、最高の持続的性能（1,207 PF/s）と並列効率（93.1%）を達成し、順次実行やオーバラップスケジュールを上回った。
• 限られたセンサデータから分布を学習することで、部分ミリオン分の1スケールの相関長を有する高次元で多スケールの確率過程を効果的に処理した。
• 量子化とワークロードの不均衡による影響は、特定のGPU数（例：10,500 GPU）でわずかな効率低下を引き起こしたが、全テスト構成において全体的なスケーリングは安定していた。