[論文レビュー] Tutel: Adaptive Mixture-of-Experts at Scale

Chang Ho Hwang, Wei Cui+13 more|arXiv (Cornell University)|Jun 7, 2022

Advanced Neural Network Applications被引用数 31

ひとこと要約

Tutelは、疎な MoE モデルのゼロコスト切替適応並列化とパイプライニングフレームワークを導入し、既存のMoE定式とアルゴリズム的等価性を維持しつつ、規模拡大時の大幅な速度向上を実現する。

ABSTRACT

Sparsely-gated mixture-of-experts (MoE) has been widely adopted to scale deep learning models to trillion-plus parameters with fixed computational cost. The algorithmic performance of MoE relies on its token routing mechanism that forwards each input token to the right sub-models or experts. While token routing dynamically determines the amount of expert workload at runtime, existing systems suffer inefficient computation due to their static execution, namely static parallelism and pipelining, which does not adapt to the dynamic workload. We present Flex, a highly scalable stack design and implementation for MoE with dynamically adaptive parallelism and pipelining. Flex designs an identical layout for distributing MoE model parameters and input data, which can be leveraged by all possible parallelism or pipelining methods without any mathematical inequivalence or tensor migration overhead. This enables adaptive parallelism/pipelining optimization at zero cost during runtime. Based on this key design, Flex also implements various MoE acceleration techniques. Aggregating all techniques, Flex finally delivers huge speedup at any scale -- 4.96x and 5.75x speedup of a single MoE layer over 16 and 2,048 A100 GPUs, respectively, over the previous state-of-the-art. Our evaluation shows that Flex efficiently and effectively runs a real-world MoE-based model named SwinV2-MoE, built upon Swin Transformer V2, a state-of-the-art computer vision architecture. On efficiency, Flex accelerates SwinV2-MoE, achieving up to 1.55x and 2.11x speedup in training and inference over Fairseq, respectively. On effectiveness, the SwinV2-MoE model achieves superior accuracy in both pre-training and down-stream computer vision tasks such as COCO object detection than the counterpart dense model, indicating the readiness of Flex for end-to-end real-world model training and inference.

研究の動機と目的

• MoEワークロードがエクサスケールで動的かつ課題が多い理由を、トークンルーティングと不均等なエキスパートワークロードのために説明する。
• DP、EP+DP+MP、および関連構成を統一し切替えるゼロオーバーヘッドの適応並列切替機構を提案する。
• 適応パイプライニングとAll-to-All戦略（2DHを含む）を開発し、計算と通信を共同で最適化する。
• 実用的でフレームワークに依存しないMoEシステム（Tutel）を、動的キャパシティファクターと動的top-kルーティング対応とともに提供する。
• 大規模なGPUクラスター上でSwinV2-MoEを用いて、スケーラビリティと実世界での有効性を示す。

提案手法

• MoE層の動的なワークロード特性と、それが並列性の選択に与える影響を分析する。
• 最小限でありながら完全な並列戦略の集合を導出する（DPとEP+DP+MP、構成全体で最適な性能に必要なもの）。
• 遷移時にO(1)オーバーヘッドとなるようデータレイアウトを再構築するゼロコスト切替可能な並列性を導入する。
• ワークロードに基づいてパイプライニング度とAll-to-Allアルゴリズム（Linear vs 2DH）を共同で選択する適応パイプライニングを開発する。
• 容量レンジに跨る最適な並列性/パイプライニング設定のディクショナリ（ハッシュマップ）を構築し、実行時の高速な意思決定を可能にする。
• 疎計算とFlexible All-to-Allを用いた高速エンコード/デコードを実装し、メモリと通信のオーバーヘッドを低減する。

実験結果

リサーチクエスチョン

• RQ1 MoEワークロードはトレーニングのダイナミクスとキャパシティファクターによってどう変化し、静的な並列性と動的な並列性にはどんな意味があるか？
• RQ2 同一のデータ配置とアルゴリズム的正確性を保ちながら、並列戦略間のゼロコスト切替を達成できるか？
• RQ3 動的MoEワークロード下で最良のスループットをもたらすデータ並列、モデル並列、エキスパート並列戦略の組み合わせは何か、そしてそれを実行時にどう適応させるか？
• RQ4 適応パイプライニングとAll-to-Allアルゴリズム（Linear vs 2DH）を共同最適化して、MoEモデルを大規模にスループット最大化するにはどうすればよいか？
• RQ5 SwinV2-MoEのような実世界のMoEモデルに対するTutelのトレーニングと推論における実用的な性能向上はどの程度か？

主な発見

• Tutelは、前処理の最先端と比較して、16 GPUでは単一MoE層で最大4.96x、2048 GPUでは最大5.75xの速度アップを達成。
• エンドツーエンドのSwinV2-MoEトレーニングは、Fairseqと比較してトレーニングで最大1.55x、推論で最大2.11xの速度向上を示す。
• 適応並列切替（DPとEP+DP+MPを介して）は、容量ファクターfに依存したスループット利得を提供し、多様な設定で静的戦略よりも上回る。
• 適応パイプライニングは、243のMoEモデル設定で平均的に9%〜100%以上の改善を生み、最悪ケースでは最大599%の改善となる。
• 柔軟なAll-to-Allと高速エンコード/デコード技術により、メモリ使用量を削減しスループットを向上させ、大規模GPUクラスタ上でのスケーラブルな疎MoE実装を可能にする。
• Tutelを用いたSwinV2-MoEは、ダウンサブCVタスク（例：COCO）でデンスなベースラインよりも高い精度を達成し、アプローチの実用性を示している。