[論文レビュー] HyPar-Flow: Exploiting MPI and Keras for Scalable Hybrid-Parallel DNN Training using TensorFlow

Ammar Ahmad Awan, Arpan Jain+3 more|arXiv (Cornell University)|Nov 12, 2019

Advanced Neural Network Applications参考文献 26被引用数 5

ひとこと要約

HyPar-Flow は、TensorFlow、MPI、Eager Execution を使用して、任意の Keras ベースの DNN をハイブリッド並列で学習可能にする、ユーザー透過でスケーラブルなフレームワークです。効率的な MPI を用いた通信とパイプライン処理により、動的にモデル並列性とデータ並列性を調整することで、512 個の Frontera ノードで ResNet-1001 の学習を単一ノード学習と比較して最大 481 倍の高速化を達成しました。

ABSTRACT

To reduce training time of large-scale DNNs, scientists have started to explore parallelization strategies like data-parallelism, model-parallelism, and hybrid-parallelism. While data-parallelism has been extensively studied and developed, several problems exist in realizing model-parallelism and hybrid-parallelism efficiently. Four major problems we focus on are: 1) defining a notion of a distributed model across processes, 2) implementing forward/back-propagation across process boundaries that requires explicit communication, 3) obtaining parallel speedup on an inherently sequential task, and 4) achieving scalability without losing out on a model's accuracy. To address these problems, we create HyPar-Flow --- a model-size/-type agnostic, scalable, practical, and user-transparent system for hybrid-parallel training by exploiting MPI, Keras, and TensorFlow. HyPar-Flow provides a single API that can be used to perform data, model, and hybrid parallel training of any Keras model at scale. We create an internal distributed representation of the user-provided Keras model, utilize TF's Eager execution features for distributed forward/back-propagation across processes, exploit pipelining to improve performance and leverage efficient MPI primitives for scalable communication. Between model partitions, we use send and recv to exchange layer-data/partial-errors while allreduce is used to accumulate/average gradients across model replicas. Beyond the design and implementation of HyPar-Flow, we also provide comprehensive correctness and performance results on three state-of-the-art HPC systems including TACC Frontera (#5 on Top500.org). For ResNet-1001, an ultra-deep model, HyPar-Flow provides: 1) Up to 1.6x speedup over Horovod-based data-parallel training, 2) 110x speedup over single-node on 128 Stampede2 nodes, and 3) 481x speedup over single-node on 512 Frontera nodes.

研究の動機と目的

• GPU のメモリ制限を超える超深層かつ大規模な DNN の学習に挑む課題に対処し、モデル並列とハイブリッド並列学習を可能にする。
• 複数のプロセスに跨る分散型の順方向/逆方向伝搬を実装する際の複雑さを克服し、明示的な通信を必要としない。
• モデル定義の変更や手動での分割を一切不要とする統一的でユーザー透過な API を提供し、データ並列、モデル並列、ハイブリッド並列の間でのシームレスな切り替えを可能にする。
• モデルの正確性を損なわず、スケーラビリティとパフォーマンスを実現する。特に、ResNet-1001 や実験的 5,000 層の ResNet-5000 といったモデルに対しても同様に有効である。
• HPC システム上でコア外モデルの学習を可能にするために、ユーザーから低レベルの通信とオーケストレーションタスクを抽象化する。

提案手法

• Keras モデル、モデルパーティション数、レプリカ数、学習戦略（データ、モデル、またはハイブリッド）を引数として受け取る、単一の高レベル API を公開する。モデルの変更は不要である。
• 内部的に、Keras モデルをプロセスに跨って分割し、データ並列性を実現するためにノードにレプリカを配置することで、分散モデルの表現を構築する。
• TensorFlow の Eager Execution を活用し、明示的な通信を伴うプロセス境界を越えた動的で分散型の順方向および逆方向パスを実現する。
• MPI のプリミティブを活用する。層間データや部分勾配の交換には send/recv を使用し、レプリカ間の勾配集約には allreduce を使用する。
• バッチパーティションにわたるパイプライン処理を実装し、最適化されたパイプラインステージサイズ（128）を採用することで、通信遅延を隠蔽し、スループットを向上させる。
• パーティショニング、配置、通信といったシステムレベルの複雑さを Trainer と Communication Engine に抽象化し、ユーザーの透明性を確保する。

実験結果

リサーチクエスチョン

• RQ1 大規模 HPC システム上で、任意の Keras モデルに対して統一的かつユーザー透過にハイブリッド並列学習を効率的に行えるフレームワークは実現可能か？
• RQ2 モデル並列性を効果的にデータ並列性と組み合わせることで、学習の正確性を損なわず、高いスケーラビリティを達成できるか？
• RQ3 深層学習の学習ワークロードにおいて、パーティション間およびレプリカ間通信に MPI プリミティブを使用した場合のパフォーマンスへの影響は何か？
• RQ4 HyPar-Flow は、単一 GPU システムのメモリ制限を超える超深層モデル、例えば ResNet-1001 や ResNet-5000 の学習をどの程度スケーリングできるか？
• RQ5 現代の HPC プラットフォームにおいて、HyPar-Flow のハイブリッド並列学習は、従来のデータ並列アプローチと比較して、どの程度のスループット向上と効率性を達成できるか？

主な発見

• 大規模システム上での ResNet-1001 の学習において、HyPar-Flow は Horovod を用いたデータ並列学習と比較して最大 1.6 倍の高速化を達成した。
• 128 個の Stampede2 ノードで、ResNet-1001 の学習において単一ノード学習と比較して 110 倍の高速化を達成し、弱スケーリングの優れた性能を示した。
• 512 個の Frontera ノード（28,762 コア）で、単一ノード学習と比較して 481 倍の高速化を達成した。これは、ハイブリッド並列学習における顕著な強スケーリングを示している。
• HyPar-Flow を用いたモデル並列学習では、逐次学習と比較して最大 3.2 倍の高速化、データ並列学習と比較して最大 1.6 倍の高速化を達成した。
• 実験的 5,000 層の ResNet-5000 モデルをわずか 2 ノードで正常に学習させることに成功し、従来のシステムでは不可能なモデルの処理が可能であることを示した。
• 128 ステージのパイプライン処理が最適であることが判明し、計算と通信のオーバーラップにより、パフォーマンスが顕著に向上した。