GPipe: 大規模ニューラルネットワークの効率的な学習を実現するパイプライン並列処理

深層ニューラルネットワークの容量を拡大することは、さまざまな機械学習タスクにおいてモデル品質を向上させる効果的なアプローチとして知られています。多くの場合、単一のアクセラレータのメモリ制限を超えてモデル容量を増加させるには、特別なアルゴリズムやインフラストラクチャを開発する必要がありました。これらの解決策はしばしばアーキテクチャに特化しており、他のタスクには転用できません。効率的かつタスクに依存しないモデル並列性のニーズに対応するために、我々はGPipe（パイプライン並列性ライブラリ）を導入します。このライブラリは、層のシーケンスとして表現できる任意のネットワークをスケーリングすることを可能にします。異なる層のサブシーケンスを個別のアクセラレータ上でパイプライン化することで、GPipeは多様なネットワークを効率的に巨大なサイズまで拡大する柔軟性を提供します。さらに、GPipeは新しいバッチ分割パイプラインアルゴリズムを利用しており、モデルが複数のアクセラレータに分割された場合にほぼ線形の速度向上が得られます。我々は2つの異なるタスクとネットワークアーキテクチャで大規模なニューラルネットワークの訓練を行い、GPipeの利点を示しています。(i) 画像分類: 5億5700万パラメータを持つAmoebaNetモデルを訓練し、ImageNet-2012でトップ1精度84.4% を達成しました。(ii) 多言語ニューラル機械翻訳: 100以上の言語をカバーするコーパス上で60億パラメータ、128層からなる単一のTransformerモデルを訓練し、すべての二言語モデルよりも高い品質を達成しました。