[論文レビュー] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

Dmitry Lepikhin, HyoukJoong Lee+7 more|arXiv (Cornell University)|Jun 30, 2020

Topic Modeling参考文献 75被引用数 348

ひとこと要約

GShardは一般的な注釈ベースのアプローチとXLAベースのSPMDコンパイラを導入し、Sparsely-Gated MoE層を持つ巨大Transformerモデルを訓練することで、計算成長をサブリニアにし、2048 TPUsで4日間に600Bパラメータの多言語翻訳を実現します。

ABSTRACT

Neural network scaling has been critical for improving the model quality in many real-world machine learning applications with vast amounts of training data and compute. Although this trend of scaling is affirmed to be a sure-fire approach for better model quality, there are challenges on the path such as the computation cost, ease of programming, and efficient implementation on parallel devices. GShard is a module composed of a set of lightweight annotation APIs and an extension to the XLA compiler. It provides an elegant way to express a wide range of parallel computation patterns with minimal changes to the existing model code. GShard enabled us to scale up multilingual neural machine translation Transformer model with Sparsely-Gated Mixture-of-Experts beyond 600 billion parameters using automatic sharding. We demonstrate that such a giant model can efficiently be trained on 2048 TPU v3 accelerators in 4 days to achieve far superior quality for translation from 100 languages to English compared to the prior art.

研究の動機と目的

• モデル品質を向上させるためのニューラルネットワークのスケーリングの必要性を動機づけつつ、計算、プログラミングの容易さ、並列展開といった実践的な課題に対処する。

提案手法

• TransformerをPosition-wise Sparsely-Gated Mixture-of-Experts (MoE)層に拡張して計算規模をサブリニアにスケールさせる。
• GShardを、軽量な注釈APIのモジュールと自動並列化のためのXLAコンパイラ拡張として導入する。
• デバイス数に関係なくコンパイル時間を一定に保つSPMD（Single Program Multiple Data）分割戦略を採用する。
• モデル開発者が巨大な単一デバイス上で作業しているかのように設計を提供し、コンパイラによって自動的に分割を適用する。
• エキスパート容量制約と補助損失を用いたゲーティング機構と、千を超えるエキスパート間の負荷を均等化する仕組みを使用する。
• 100言語対の多言語機械翻訳でのエンドツーエンドの訓練とスケーリングを実証する。

実験結果

リサーチクエスチョン

• RQ1 非常に大規模なTransformerモデルを、煩雑なコンパイルや通信オーバーヘッドなく、千を超えるデバイスで効率的に訓練するにはどうすればよいか。
• RQ2 Sparsely-Gated MoE層による条件付き計算は、モデル容量が増加しても計算のサブリニア成長を提供できるか。
• RQ3 注釈駆動型のGShardアプローチは、XLA上で自動的でスケーラブルな分割を可能にする一方で、モデル開発を簡素化できるか。
• RQ4 多言語設定で数百億パラメータへスケールする際の翻訳品質の実用的な向上は何か。

主な発見

• 2048 TPU v3デバイスを用い4日間で訓練した600Bパラメータの sparsely-gated MoE Transformerは、100言語から英語への翻訳品質で先行技術を上回った。
• モデルサイズの増加に伴い訓練コストがサブリニアに増加し、容量の増加に対する計算のサブリニアスケーリングを示した。
• 密結合のベースラインTransformer（2.3Bパラメータ）は235.5 TPU v3コア年を要し、MoEアプローチの効率向上を示した。
• GShardは自動分割を可能にし、SPMDアプローチでデバイス数が増えてもO(1)のコンパイル時間でスケールする。
• MoEのゲーティングにはエキスパート容量制約と負荷を均衡させる補助損失、キャパシティを効果的に活用するための第二候補エキスパートのランダムルーティングが含まれる。