[論文レビュー] The Case for Strong Scaling in Deep Learning: Training Large 3D CNNs with Hybrid Parallelism

Yosuke Oyama, Naoya Maruyama+7 more|arXiv (Cornell University)|Jul 25, 2020

Advanced Neural Network Applications参考文献 54被引用数 9

ひとこと要約

本論文は、データ並列性と空間並列性を組み合わせることで、1つのサンプルを複数のGPUに分散させることで、大規模な3次元畳み込みニューラルネットワーク（3D CNN）の強力なスケーリングを可能にするエンドツーエンドのハイブリッド並列学習フレームワークを提案する。ミニバッチサイズを超えて並列化を拡張し、空間分割によるI/O最適化を施すことで、高いパフォーマンスを達成し、512³のフル解像度の宇宙論的データに対する学習を可能にした。その結果、CosmoFlowの予測精度が低解像度学習と比較して10倍向上した。

ABSTRACT

We present scalable hybrid-parallel algorithms for training large-scale 3D convolutional neural networks. Deep learning-based emerging scientific workflows often require model training with large, high-dimensional samples, which can make training much more costly and even infeasible due to excessive memory usage. We solve these challenges by extensively applying hybrid parallelism throughout the end-to-end training pipeline, including both computations and I/O. Our hybrid-parallel algorithm extends the standard data parallelism with spatial parallelism, which partitions a single sample in the spatial domain, realizing strong scaling beyond the mini-batch dimension with a larger aggregated memory capacity. We evaluate our proposed training algorithms with two challenging 3D CNNs, CosmoFlow and 3D U-Net. Our comprehensive performance studies show that good weak and strong scaling can be achieved for both networks using up 2K GPUs. More importantly, we enable training of CosmoFlow with much larger samples than previously possible, realizing an order-of-magnitude improvement in prediction accuracy.

研究の動機と目的

• 宇宙論的シミュレーションや3次元医療画像などの高次元科学的データに対して、大規模な3次元畳み込みニューラルネットワーク（3D CNN）を学習する際のメモリおよび計算上のボトル neck を解決すること。
• 大きなサンプルによるGPUごとのメモリ圧力のため、従来のデータ並列性では効果的にスケーリングできないという制限を克服すること。
• 個々のサンプルの空間並列性を導入することで、ミニバッチサイズを増加させずにトレーニング時間を短縮する「強力なスケーリング」を実現すること。
• データローディングに対しても同じハイブリッド並列技術を適用することで、I/Oパイプライン全体をスケーリングし、大規模なサンプルに対するI/Oボトル neck を最小限に抑えること。
• フル解像度での学習が、長距離特徴抽出を必要とする科学的タスクにおいて顕著にモデルの精度を向上させることを実証すること。

提案手法

• 1つの3次元サンプルを複数のGPUに分散することで、GPUごとのメモリ使用量を削減する、データ並列性と空間並列性を組み合わせたハイブリッド並列アルゴリズムを導入する。
• 処理要素間で特徴マップの空間次元（例：高さ、幅、深さ）を分割することで、3次元畳み込みに空間分割を適用し、ミニバッチの制限を超えた強力なスケーリングを可能にする。
• 大規模なデータサンプルのI/Oパイプラインを空間分割により最適化し、GPU数に比例してスケーリングする集団的かつ非同期I/Oを実現する。
• 通信オーバーヘッドを低減するために、ハイブリッド並列学習とI/Oをサポートするカスタムカーネルおよび通信パターンをLBANNディープラーニングフレームワークに拡張する。
• チャンネル／特徴量並列性で用いられる完全なテンソル通信に比べ、データ移動を最小限に抑えるために、空間パーティション間でのハローエクスチェンジを採用する。
• モデルベースのパフォーマンス解析を用いて、3次元CNNにおける計算およびI/Oの両方のスケーリングボトル neck を同定・緩和する。

実験結果

リサーチクエスチョン

• RQ1 個々のサンプルの空間分割を導入することで、データ並列性を超えて並列化を拡張することにより、3次元CNN学習における強力なスケーリングが達成可能か？
• RQ2 空間並列性とデータ並列性を組み合わせたハイブリッド並列学習により、2,048GPUまで効率的なスケーリングが可能となり、高いスループットと低い通信オーバーヘッドを維持できるか？
• RQ3 計算に適用したのと同じハイブリッド並列技術をI/Oに適用することで、大規模な3次元データサンプルのI/Oパイプラインを効果的にスケーリングできるか？
• RQ4 512³などのフル解像度の3次元データ（例：512³）を用いた学習は、高いメモリおよび計算負荷にもかかわらず、低解像度データでの学習と比較して顕著に高いモデル精度を達成できるか？
• RQ5 ハイブリッド並列学習およびI/Oを用いた場合、CosmoFlow や3D U-Net などの3次元CNNのパフォーマンス特性およびスケーリング効率はどのように変化するか？

主な発見

• 512³のフル解像度サンプルを用いてCosmoFlowを学習した結果、平均二乗誤差（MSE）が0.0018にまで低下し、128³サンプルで学習した場合（MSE = 0.0289）と比較して予測精度が10倍向上した。
• 64のミニバッチサイズで学習したCosmoFlowモデルは、512GPUから2048GPUにスケーリングした際、1.77倍の高速化を達成し、ミニバッチサイズを増加させずに強力なスケーリングを示した。
• 3D U-Netモデルは、256GPUから512GPUにスケーリングした際、1.42倍の高速化を示し、異なるアーキテクチャおよびワークロードにおいても強力なスケーリング特性を確認した。
• CosmoFlowおよび3D U-Netの両モデルが、2,048GPUまで良好な弱スケーリングおよび強スケーリングを示し、パフォーマンス解析により計算およびI/Oコンponentの両方で高いスケーリング効率が確認された。
• ハイブリッド並列I/OパイプラインはGPU数に比例してスケーリングに成功し、科学的3次元データで一般的なギガバイトサイズのサンプルに対しても、I/Oボトル neck を最小限に抑えることができた。