[論文レビュー] Scalable and Efficient MoE Training for Multitask Multilingual Models

Young Jin Kim, Ammar Ahmad Awan+7 more|arXiv (Cornell University)|Sep 22, 2021

Topic Modeling参考文献 26被引用数 33

ひとこと要約

本論文は DeepSpeed MoE を紹介し、マルチタスク多言語設定に跨る大規模 Mixture-of-Experts（MoE）モデルの訓練をスケーラブルに行えるシステムを提案する。訓練技術（RTS、AoE、剪定）とともに、Z-code M3 を用いた最大 10B パラメータまでの強力な MT および多言語生成結果を示す。

ABSTRACT

The Mixture of Experts (MoE) models are an emerging class of sparsely activated deep learning models that have sublinear compute costs with respect to their parameters. In contrast with dense models, the sparse architecture of MoE offers opportunities for drastically growing model size with significant accuracy gain while consuming much lower compute budget. However, supporting large scale MoE training also has its own set of system and modeling challenges. To overcome the challenges and embrace the opportunities of MoE, we first develop a system capable of scaling MoE models efficiently to trillions of parameters. It combines multi-dimensional parallelism and heterogeneous memory technologies harmoniously with MoE to empower 8x larger models on the same hardware compared with existing work. Besides boosting system efficiency, we also present new training methods to improve MoE sample efficiency and leverage expert pruning strategy to improve inference time efficiency. By combining the efficient system and training methods, we are able to significantly scale up large multitask multilingual models for language generation which results in a great improvement in model accuracy. A model trained with 10 billion parameters on 50 languages can achieve state-of-the-art performance in Machine Translation (MT) and multilingual natural language generation tasks. The system support of efficient MoE training has been implemented and open-sourced with the DeepSpeed library.

研究の動機と目的

• 億〜兆級パラメータに達する、多タスク多言語 Mixture-of-Experts モデルのスケーラブルな訓練を実現する。
• サンプル効率と実行時効率を改善する訓練およびシステム技術を開発する。
• 実用的なモデル（Z-code M3）を示し、50言語にわたる強力な機械翻訳と多言語生成を達成。

提案手法

• データ、エキスパート、モデル、テンソルスライシング、ZeRO の五つの parallelism forms (data, expert, model, tensor-slicing, ZeRO) と ZeRO-Offload を実現する DeepSpeed MoE システムを開発し、GPU メモリ制限を超える。
• MoE ルーティングのトークンバイアスを緩和するためにランダムトークン選択を導入する。
• チェックポイントを組み合わせて初期化時および訓練時により大きなエキスパートプールを作るための AoE（Aggregation of Experts）を提案する。
• 推論を高速化するためのエキスパート剪定戦略（ランダムおよび利用度ベース）を検討する。
• 複数タスクの多言語モデル（MT、DAE、ELECTRA、MLM）を、タスクリ losses の総和で単一の目的関数として訓練する。
• MoE 層を交互の層ごとに配置した Transformer encoder-decoder アーキテクチャと、250k の SentencePiece 語彙を使用する。

実験結果

リサーチクエスチョン

• RQ1 MoE アーキテクチャをどのようにスケールさせて、マルチタスク多言語訓練のために兆級パラメータへ拡張できるか？
• RQ2 DeepSpeed MoE は GPU メモリ制限を克服し、より大きなベースモデルとより多くのエキスパートを可能にできるか？
• RQ3 MoE ベースのマルチタスク多言語訓練は、密結合ベースと比較して下流の MT および NLG タスクを改善するか？
• RQ4 大規模 MoE モデルにおいて、どの訓練技術がサンプル効率と推論効率を最大化するか？
• RQ5 多タスク目的の影響は、多言語翻訳と生成品質にどのように表れるか？

主な発見

Model	X→English	English→X	Average
Non-MoE multilingual model	35.16	30.36	32.76
Individual bilingual models	35.70	30.40	33.05
Z-code M3 [10B]	42.09	32.22	37.15

• DeepSpeed MoE は GPU 間でほぼ線形のスループットスケーリングを実現し、ZeRO-Offload により GPU メモリを超えるモデルサイズをサポートする。
• RTS はトークン位置バイアスを低減し、MoE 訓練の収束速度と正則化を改善する。
• AoE はチェックポイントからパラメータを統合して、初期化時により大きなモデルを作るためのより大きな有効エキスパートプールを作成できる。
• エキスパート剪定は、ランダムまたは利用度ベースの選択により、性能を大きく犠牲にせずに推論を小型で効率的にする。
• Z-code M3 (10B パラメータ、64 エキスパート) は、MT および多言語タスクで密結合のベースラインや小型 MoE 構成を上回り、強力な多言語性能を示す。
• ファインチューニング済みの Z-code M3 モデルは、Wikilingua などの下流タスクやクロスリンガル生成で顕著な改善を達成。