[論文レビュー] M6-T: Exploring Sparse Expert Models and Beyond

Yang An, Junyang Lin+13 more|arXiv (Cornell University)|May 31, 2021

Recommender Systems and Techniques参考文献 34被引用数 14

ひとこと要約

この論文では、スパースエキスパートモデルにおけるパフォーマンスを向上させつつ計算コストを一定に保ちながら、エキスパートをk個のプロトタイプに分割し、k個のトップ-1ルーティングを用いるエキスパートプロトタイピングという手法を提案している。このアプローチにより、1.1兆パラメータのモデルを480個のV100 GPUでのみ使用して、同規模のベースラインと比較して5倍速く学習可能となり、2048 TPUコアで学習された先行SOTAモデルを上回る性能を達成した。

ABSTRACT

Mixture-of-Experts (MoE) models can achieve promising results with outrageous large amount of parameters but constant computation cost, and thus it has become a trend in model scaling. Still it is a mystery how MoE layers bring quality gains by leveraging the parameters with sparse activation. In this work, we investigate several key factors in sparse expert models. We observe that load imbalance may not be a significant problem affecting model quality, contrary to the perspectives of recent studies, while the number of sparsely activated experts $k$ and expert capacity $C$ in top-$k$ routing can significantly make a difference in this context. Furthermore, we take a step forward to propose a simple method called expert prototyping that splits experts into different prototypes and applies $k$ top-$1$ routing. This strategy improves the model quality but maintains constant computational costs, and our further exploration on extremely large-scale models reflects that it is more effective in training larger models. We push the model scale to over $1$ trillion parameters and implement it on solely $480$ NVIDIA V100-32GB GPUs, in comparison with the recent SOTAs on $2048$ TPU cores. The proposed giant model achieves substantial speedup in convergence over the same-size baseline.

研究の動機と目的

• 混合エキスパート（MoE）モデルにおけるモデル品質と学習効率に影響を与える要因を調査すること。
• 1000億パラメータを超えるスケーリングを、スパースエキスパートモデルで効率的に行う課題に取り組むこと。
• 極めて大規模なMoEアーキテクチャにおいて、計算コストを削減しながらモデルパフォーマンスを向上させること。
• 消費者向けGPUクラスタを用いて、1兆パラメータのMoEモデルをTPUベースのシステムではなく効率的に学習可能であることを実証すること。

提案手法

• エキスパートプロトタイピングは、各エキスパートをk個の異なるプロトタイプに分割し、標準的なトップ-kルーティングではなくk個のトップ-1ルーティングを可能にする。
• 各入力トークンは、ゲーティング関数を用いてk個の別々のプロトタイプにルーティングされ、各プロトタイプごとに独立してルーティング意思決定がなされる。
• 最終出力は、トップ-1選択結果に対するソフトマックスゲーティングアテンションを用いた、k個のプロトタイプ出力の重み付き和として計算される。
• この手法は、全トップ-k計算を避けることで、FLOPsを一定に保つ。代わりに、k個の独立したトップ-1ルーティング処理を実行する。
• 安定性を確保するため、学習率を低減（0.005）し、重み初期化をスケーリング（σ=0.002）したAdafactor最適化法を用いてモデルを学習する。
• RDMAインターコネクトを備えたWhale分散フレームワークを用い、480個のNVIDIA V100-32GB GPUから構成されるクラスタ上で学習を実施する。

実験結果

リサーチクエスチョン

• RQ1 負荷の不均衡は、MoEモデルにおけるモデル品質に顕著に悪影響を及えるのか？
• RQ2 アクティブ化されるエキスパート数（k）とエキスパート容量（C）の変化が、MoEパフォーマンスに与える影響は何か？
• RQ3 計算コストを増加させることなく、エキスパートプロトタイピングはモデル品質を向上させることができるのか？
• RQ4 エキスパートプロトタイピングは、極めて大規模なモデル（例：1000億パラメータ以上）の学習に対してどれほど有効なのか？
• RQ5 1兆パラメータのMoEモデルをTPUポッドではなくGPUクラスタで効率的に学習可能なのか？

主な発見

• 負荷の不均衡は、従来の仮定とは異なり、MoEモデルにおけるモデル品質に顕著な悪影響を及えないことが判明した。
• アクティブ化されるエキスパート数（k）とエキスパート容量（C）を増加させることで、モデルパフォーマンスが顕著に向上することが分かった。
• エキスパートプロトタイピングにより、1.1兆パラメータのMoEモデルの収束が、同規模のベースラインと比較して5倍速くなった。
• 1.1兆パラメータのエキスパートプロトタイピングモデルは、480個のV100-32GB GPU上で学習されたが、2048 TPUコアで学習された先行SOTAモデルを上回る性能を達成した。
• この手法は、計算コストを一定に保ちながらモデル品質を向上させることができ、特に大規模な設定において顕著な改善が見られた。
• FLOPsが同一であるにもかかわらず、モデルの収束が著しく速くなったことから、学習ダイナミクスにおける効率性の向上が示された。