[論文レビュー] Maximizing Parallelism in Distributed Training for Huge Neural Networks

Zhengda Bian, Qifan Xu+2 more|arXiv (Cornell University)|May 30, 2021

Advanced Neural Network Applications参考文献 14被引用数 20

ひとこと要約

本稿では、特にTransformer言語モデルを対象として、線形層の計算を3次元（テンソル、シーケンス、ヘッド）に分散させることで、完全な負荷バランスを達成する、巨大ニューラルネットワークのトレーニングに向けた新規な3次元モデル並列技術を提案する。メモリと通信コストを低減し、64枚のV100 GPU上での1-Dおよび2-D並列化と比較して、それぞれ2.32倍および1.57倍の高速化を達成した。

ABSTRACT

The recent Natural Language Processing techniques have been refreshing the state-of-the-art performance at an incredible speed. Training huge language models is therefore an imperative demand in both industry and academy. However, huge language models impose challenges to both hardware and software. Graphical processing units (GPUs) are iterated frequently to meet the exploding demand, and a variety of ASICs like TPUs are spawned. However, there is still a tension between the fast growth of the extremely huge models and the fact that Moore's law is approaching the end. To this end, many model parallelism techniques are proposed to distribute the model parameters to multiple devices, so as to alleviate the tension on both memory and computation. Our work is the first to introduce a 3-dimensional model parallelism for expediting huge language models. By reaching a perfect load balance, our approach presents smaller memory and communication cost than existing state-of-the-art 1-D and 2-D model parallelism. Our experiments on 64 TACC's V100 GPUs show that our 3-D parallelism outperforms the 1-D and 2-D parallelism with 2.32x and 1.57x speedup, respectively.

研究の動機と目的

• 単一デバイスのメモリおよび計算能力を超える巨大言語モデルのトレーニングという、ますます深刻化する課題に対処すること。
• 線形層の3次元分解を導入することで、分散モデル並列化における通信およびメモリオーヘッドを低減すること。
• 巨大なTransformerモデルの分散トレーニングにおいて、GPU間で完全な負荷バランスを達成すること。
• 弱スケーリングおよび強スケーリングの両状況において、従来の1-Dおよび2-Dモデル並列化技術を上回ること。

提案手法

• 線形層をテンソル、シーケンス、ヘッドの3次元に沿って分解することで、計算と通信のバランスを取る3次元モデル並列化アルゴリズムを提案する。
• アイドル時間を最小限に抑え、すべてのデバイスのGPU利用率を最大化するための負荷バランスの取れたデータレイアウトを採用する。
• 3Dタイリング戦略を適用して、重み行列と活性化値を3次元GPUグリッドに分割し、通信量を削減する。
• PyTorchの分散通信バックエンドと統合し、混合精度トレーニングをサポートすることで、メモリフットプリントを低減する。
• 計算と通信をオーバーラップするカスタム通信スケジューリングを採用し、レイテンシを隠蔽する。
• 既存のPyTorch実装を拡張して3次元テンソル並列を実装し、3次元並列Transformerを実現する。

実験結果

リサーチクエスチョン

• RQ1 3次元モデル並列化設計は、分散トレーニングにおける大規模なTransformerモデルに対して、1-Dまたは2-D手法よりも優れた負荷バランスを達成できるか？
• RQ2 従来の2-Dおよび1-D手法と比較して、3-D並列化は通信コストおよびメモリコストを低減できるか？
• RQ3 大規模なGPUクラスタ上での弱スケーリングおよび強スケーリングの両状況において、3-Dアプローチはどのようにスケーリングするか？
• RQ4 3-D並列化は、モデルの精度を維持したまま、より高いトレーニングスループットを達成できるか？

主な発見

• 64枚のV100 GPU上では、3次元並列化は強スケーリングにおいて1-Dモデル並列化と比較して2.32倍の高速化を達成した。
• 同じハードウェア環境下で、3次元アプローチは2-Dモデル並列化と比較して1.57倍の高速化を達成した。
• 弱スケーリングにおいて、3次元手法は平均ステップ時間の増加が最も遅く、優れたスケーラビリティと低い通信オーバーヘッドを示した。
• 3次元並列化は、すべてのGPU数において最小の平均ステップ時間を維持しており、最適な負荷バランスを示した。
• 計算とデータのGPU間でのより均等な分散を実現することで、メモリおよび通信コストを低減した。
• 3次元手法は、特にスケーリングが進む状況において、前方伝搬および逆伝搬の両時間において1-Dおよび2-D手法を上回った。