[論文レビュー] Distributed Machine Learning for Computational Engineering using MPI

Kailai Xu, Weiqiang Zhu+1 more|arXiv (Cornell University)|Nov 2, 2020

Reservoir Engineering and Simulation Methods参考文献 42被引用数 5

ひとこと要約

本論文は、MPIを用いて深層ニューラルネットワーク（DNN）と偏微分方程式（PDE）ソルバーを統合する分散機械学習フレームワークを提案する。データ通信を計算グラフ内の第一級のノードとして扱い、分散メモリシステム上で並列な前向きおよび随伴計算を可能にすることで、強い弱いスケーリングを達成し、100プロセッサで20倍以上の高速化を実現。大規模なPDE制約付きDNNにおけるスケーラブルな逆問題モデリングを示している。

ABSTRACT

We propose a framework for training neural networks that are coupled with partial differential equations (PDEs) in a parallel computing environment. Unlike most distributed computing frameworks for deep neural networks, our focus is to parallelize both numerical solvers and deep neural networks in forward and adjoint computations. Our parallel computing model views data communication as a node in the computational graph for numerical simulations. The advantage of our model is that data communication and computing are cleanly separated and thus provide better flexibility, modularity, and testability. We demonstrate using various large-scale problems that we can achieve substantial acceleration by using parallel solvers for PDEs in training deep neural networks that are coupled with PDEs.

研究の動機と目的

• 逆モード自動微分を用いたPDEと結合されたDNNの学習におけるメモリおよび計算のボトル neck を解消すること。
• 計算工学分野の大規模な逆問題に対して、スケーラブルで分散化された最適化を可能にすること。
• 並列PDEソルバーと分散DNN学習を、第一級の通信ノードを備えた単一の計算グラフ抽象化の下に統合すること。
• MPIベースのデータ通信およびハイブリッド並列化（MPI + スレーディング）における勾配のバックプロpagationに直面する課題を克服すること。

提案手法

• MPIブロードキャストやハローエクスチェンジなどのデータ通信操作を、計算グラフ内での明示的なノードとして扱い、エンドツーエンドの自動微分を可能にする。
• DNNとPDEソルバー（例：有限要素法／有限差分法）を統合した計算グラフに統合し、逆モードADによる勾配計算を実現する。
• 基本的な算術演算ではなく、高レベルの抽象化（例：行列ソルバー、PDEオペレータ）を対象とする粗粒度のADアプローチを採用する。
• ADCMEライブラリにカスタムMPIプリミティブを実装し、通信パターンを介した勾配のバックプロパゲーションをサポートする。
• 計算と通信のロジックを分離することで、ハイブリッドMPI+スレーディング環境におけるモジュラリティ、テスト可能性、柔軟性を向上させる。
• ドメイン分割を用いて計算領域をMPIプロセスに分割し、ステンシルベースのPDEソルバーのためのハローエクスチェンジを実施する。

実験結果

リサーチクエスチョン

• RQ1 分散PDEソルバーにおけるデータ通信を、自動微分をサポートする第一級の計算グラフノードとしてモデル化できるか？
• RQ2 ハイブリッドMPI+スレーディング環境において、逆モードADがMPI通信パターンの勾配を処理できるように拡張できるか？
• RQ3 本フレームワークを用いてPDEと結合されたDNNの学習において、強いスケーリングおよび弱いスケーリングはどの程度達成できるか？
• RQ4 数千コアにわたる大規模並列処理を実現しながらも、数値的精度と収束性を維持できるか？
• RQ5 物理学的制約付き学習における従来のデータ並列またはモデル並列DNN学習と比較して、本アプローチの性能はどのように異なるか？

主な発見

• 音響波動方程式のシナリオでは、100プロセッサ（3200コア）を用いて20倍以上の高速化を達成し、弱いスケーリングにおいても最小限の遅延が生じた。
• 音響波動方程式では、プロセッサ数を1から100に増加させた際の実行時間はわずか2倍にしか増加せず、優れたスケーラビリティを示している。
• 弾性波動方程式のシナリオでは、より高いメモリと状態変数の複雑さにもかかわらず、100プロセッサで20倍以上の高速化を達成し、強いスケーリングを維持している。
• 計算グラフ抽象化により、通信と計算の明確な分離が可能となり、モジュラリティとテスト可能性が向上した。
• 本フレームワークは、DNNパラメータおよびMPI通信操作の両方を介したバックプロパゲーションを正常にサポートし、エンドツーエンドの勾配計算を可能にした。
• 本実装は、波動伝播におけるパrameter同定などのPDEとDNNを統合した大規模な逆問題に対して、実用的な妥当性を示している。