[論文レビュー] GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

Yanping Huang, Youlong Cheng+9 more|arXiv (Cornell University)|Nov 16, 2018

Digital Imaging for Blood Diseases被引用数 881

ひとこと要約

GPipeは、同期勾配を用いて複数の加速器全体で巨大なニューラルネットワークを訓練するための、バッチ分割パイプライン並列性を導入します。これにより、ほぼ線形の速度向上とスケーラブルなメモリ効率が実現します。

ABSTRACT

Scaling up deep neural network capacity has been known as an effective approach to improving model quality for several different machine learning tasks. In many cases, increasing model capacity beyond the memory limit of a single accelerator has required developing special algorithms or infrastructure. These solutions are often architecture-specific and do not transfer to other tasks. To address the need for efficient and task-independent model parallelism, we introduce GPipe, a pipeline parallelism library that allows scaling any network that can be expressed as a sequence of layers. By pipelining different sub-sequences of layers on separate accelerators, GPipe provides the flexibility of scaling a variety of different networks to gigantic sizes efficiently. Moreover, GPipe utilizes a novel batch-splitting pipelining algorithm, resulting in almost linear speedup when a model is partitioned across multiple accelerators. We demonstrate the advantages of GPipe by training large-scale neural networks on two different tasks with distinct network architectures: (i) Image Classification: We train a 557-million-parameter AmoebaNet model and attain a top-1 accuracy of 84.4% on ImageNet-2012, (ii) Multilingual Neural Machine Translation: We train a single 6-billion-parameter, 128-layer Transformer model on a corpus spanning over 100 languages and achieve better quality than all bilingual models.

研究の動機と目的

• 単一の加速器のメモリを超えるサイズのネットワークのスケーラブルな訓練を可能にする。
• タスクおよびアーキテクチャに依存しない柔軟なモデル並列性ライブラリを提供する。
• 安定性のための同期勾配更新を備えたバッチ分割パイプラインアルゴリズムを開発する。
• 再材料化によって活性化メモリを削減し、巨大モデルを適合させる。
• 画像分類と多言語機械翻訳タスクでスケーラビリティを実証する。

提案手法

• L層からなる逐次的なネットワークをK個のセルに分割し、それぞれを別の加速器に配置する。
• サイズNの各ミニバッチをM個のマイクロバッチに分割し、K個の加速器にまたがってパイプライン化する。
• パラメータを更新する前に、M個のマイクロバッチ全体で勾配を蓄積して、同期ミニバッチ勾配降下を用いる。
• バックワードパス中に前方活性化の再計算（再材料化）を適用し、ピークメモリを O(N + (L/K) * (N/M)) に削減する。
• パイプラインの効率を均衡させ、バブルオーバーヘッドを低減するため、推定セルコストの分散を最小化する（K-1)/(M+K-1)。
• さらなるスケーリングのために、GPipeをデータ並列と組み合わせることもできる。

実験結果

リサーチクエスチョン

• RQ1 GPipeは、単一の加速器のメモリ制限を超える、レイヤーの列として表現できる任意の深いネットワークをスケールさせることができるか？
• RQ2 マイクロバッチ分割（M）と分割数（K）が訓練効率、メモリ、スケーラビリティに与える影響は？
• RQ3 再材料化が巨大モデルのメモリ使用量と訓練の安定性に与える影響は？
• RQ4 畳み込み（AmoebaNet）およびトランスフォーマーをベースとした（多言語対応の Transformer）アーキテクチャで、どの程度の性能とスループットの向上が達成できるか？
• RQ5 層分布が非均衡なネットワークアーキテクチャはGPipe下でのスピードアップにどのように影響するか？

主な発見

• GPipeは8つの加速器で最大1.8B AmoebaNetパラメータの訓練を可能にし、単一加速器の制限を大幅に超える。
• 6Bパラメータの多言語Transformerを128の加速器で訓練でき、100言語以上でバイリンガルベースラインを上回る改善を達成。
• Mが少なくとも4×Kであれば、バブルオーバーヘッドは無視できるレベルとなり、Transformerワークロードでほぼ線形のスピードアップを達成。
• 再材料化による大幅なメモリ削減をGPipeは実現し、与えられたハードウェアではるかに大きなモデルを可能にする。
• 実験は557MパラメータのAmoebaNetと、データセット間で競合的な転移学習を示し、ImageNetで良好な結果。
• 多言語MTでは深さと幅のトレードオフが存在し、深いモデルは資源の少ない言語で優れ、非常に大きなモデルは資源の多い言語で改善を促す。