[論文レビュー] FastMoE: A Fast Mixture-of-Expert Training System

Jiaao He, Jiezhong Qiu+4 more|arXiv (Cornell University)|Mar 24, 2021

Topic Modeling参考文献 16被引用数 39

ひとこと要約

FastMoE は PyTorch ベースの分散 Mixture-of-Experts (MoE) トレーニングシステムで、複数のエキスパートを GPU およびノード全体に配置できるようにし、最適化された CUDA カーネルと柔軟な Megatron-LM 連携により、汎用ハードウェア上でスケーラブルな MoE トレーニングを実現します。

ABSTRACT

Mixture-of-Expert (MoE) presents a strong potential in enlarging the size of language model to trillions of parameters. However, training trillion-scale MoE requires algorithm and system co-design for a well-tuned high performance distributed training system. Unfortunately, the only existing platform that meets the requirements strongly depends on Google's hardware (TPU) and software (Mesh Tensorflow) stack, and is not open and available to the public, especially GPU and PyTorch communities. In this paper, we present FastMoE, a distributed MoE training system based on PyTorch with common accelerators. The system provides a hierarchical interface for both flexible model design and easy adaption to different applications, such as Transformer-XL and Megatron-LM. Different from direct implementation of MoE models using PyTorch, the training speed is highly optimized in FastMoE by sophisticated high-performance acceleration skills. The system supports placing different experts on multiple GPUs across multiple nodes, enabling enlarging the number of experts linearly against the number of GPUs. The source of FastMoE is available at https://github.com/laekov/fastmoe under Apache-2 license.

研究の動機と目的

• TPU/Mesh TensorFlow スタックを超えた、公開可能なハードウェア（GPU）上でのスケーラブルな MoE トレーニングの動機づけ。
• PyTorch と Megatron-LM に統合された、使いやすく柔軟で効率的な MoE フレームワークを提供する。
• 複数の GPU およびノードにまたがって多数のエキスパートを分散配置し、モデル容量を拡張できるようにする。

提案手法

• 任意のニューラルネットワークをエキスパートとして再現でき、高度な MoE 動作のためのカスタム expert_fn をサポートする柔軟な FMoE インターフェースを導入する。
• トランスフォーマーベースの MoE 向けに最適化された高性能 FFN（FMoETransformerMLP）と、並列エキスパート実行のための専用 FMoELinear モジュールを提供する。
• 2 行の変換（fmoefy）により Megatron-LM とのプラグイン風統合を実装し、FFN を MoE ネットワークと入れ替える。
• 複数のワーカーとノードにまたがるエキスパートの分散配置（モデル並列 MoE）を、隠れたグローバルデータ交換操作と不均一な同期モジュールでサポートする。
• エキスパートごとに入力をバッチ化するためのカスタム scatter/gather メモリレイアウトを使用し、GPU 利用率を最大化しデータ移動を削減する。

実験結果

リサーチクエスチョン

• RQ1 FastMoE は汎用 GPU を用いたマルチ GPU・マルチノードクラスターでスケーラブルな MoE トレーニングを実現できるか。
• RQ2 専用 FFN と並列エキスパート実行を備えた最適化された MoE 実装は、単一 GPU 上の素朴な PyTorch MoE 実装を上回るか。
• RQ3 Megatron-LM との統合を前提としたエンドツーエンドのトレーニングで、FastMoE は非 MoE のベースラインと比較してどのように性能を示すか。
• RQ4 GPUs やエキスパート数を増やしたときのスケーラビリティの限界とボトルネック（例：通信）は何か。

主な発見

• FastMoE は類似タスクに対して、単一 GPU でのベースライン PyTorch MoE 実装よりもフォワード/バックワードが速い。
• クラスター上の 8 台の NVIDIA Tesla V100 GPU で、FastMoE はスループットを約 10 TFLOPs から 25 TFLOPs にスケールさせ、通信オーバーヘッドによる部分的な線形性の欠如を示す。
• 8 GPU に跨る 12 層の GPT モデルを、層あたり 96 エキスパート（GPU 当たり 12 エキスパート）でエンドツーエンドのトレーニングは、反復ごとの速度は遅くなるものの、同じトレーニング時間でより低い損失を達成することで MoE によるモデル容量の利得を示す。
• Megatron-LM を用いた FastMoE による、レイヤーあたり 96 エキスパートの MoE GPT モデルを 70 時間訓練すると、同等の計算量を用いた非 MoE ベースラインよりモデル性能で MoE の利得を示す。
• FFN を置換するだけで MoE への簡易転換を可能にし、2 行のコード変更（fmoefy）で Megatron-LM との実用的な統合を実証する。
• このシステムはオープンソース化の利用可能性と、将来の改善として負荷分散機構を追加する計画を強調している。