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

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

Domain Adaptation and Few-Shot Learning参考文献 41被引用数 236

ひとこと要約

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.

研究の動機と目的

• タスクやアーキテクチャに依存しないスケーラブルなモデル並列性を動機づける
• GPipeとして、ネットワークを加速器間で連続的なセルに分割するパイプライン並列ライブラリを導入する
• 利用率を最大化するためのバッチ分割マイクロバッチスケジューリングアルゴリズムを開発する
• 再材料化を用いてアクティベーションメモリを削減し、より大きなモデルを適合させる
• 画像分類と多言語機械翻訳でのスケーリングをデモンストレーションする

提案手法

• 任意のネットワークを L 層の連続として表現し、加速器間で K 個のセルに分割する
• 前向き伝搬でマイクロバッチ(M)をセルを通して流し、単一の同期更新前にミニバッチ全体で勾配を蓄積する
• パイプラインのバランスを取り、バブルオーバーヘッドを減らすためにセルコストの分散を最小化する
• バックワードパスで前方のアクティベーションを再計算する再材料化を使用し、ピークメモリを制限する
• トレーニングをさらにスケールするためにオプションのデータ並列性を許可する
• マイクロバッチを使用する場合、マイクロバッチ間のバッチ正規化統計を扱う

実験結果

リサーチクエスチョン

• RQ1 加速器を跨いでスケールする際、バッチ分割を伴うパイプライン並列性はほぼ線形のスピードアップを達成可能か？
• RQ2 GPipeは巨大モデルで安定した学習を維持するために、メモリと同期をどのように管理するか？
• RQ3 GPU/TPU上で非常に大規模なモデル（例：十億パラメータ）を訓練する際のメモリとスループットのトレードオフは何か？
• RQ4 さまざまなアーキテクチャ（ConvNetsとTransformerなど）とハードウェア・インターコネクトでGPipeはどのように性能を発揮するか？
• RQ5 マイクロバッチサイズと分割数が訓練の効率と安定性に与える影響は何か？

主な発見

• GPipeは、アクセラレータ間で分割し、マイクロバッチパイプライン並列性を使用することで、AmoebaNetは最大1.8Bパラメータ、Transformerは最大83.9Bパラメータに達するなど、非常に大規模なモデルの訓練を可能にする。
• 再材料化はピークアクティベーションメモリをO(N·L)から概ねO(N + (L/K)·(N/M))へ削減し、固定ハードウェア上でより大きなモデルを可能にする。
• M ≥ 4·K のとき、バブルオーバーヘッドは無視できるレベルとなり、トランスフォーマーではアクセラレータ数に対してスケーラビリティはほぼ線形になる。
• GPipeはアクセラレータ全体でほぼ線形のスループット速度アップを達成し、Transformerアーキテクチャでは線形スケーリング、AmoebaNetのような不均一なネットワークでは部分的に線形である。
• 実験はモデルサイズと翻訳品質の大幅な向上を示し、ImageNetで557Mパラメータでトップ1が最大84.4%、多言語Transformer 6Bパラメータは言語間のバイリンガルベースラインを上回る