[論文レビュー] HetuMoE: An Efficient Trillion-scale Mixture-of-Expert Distributed Training System

Xiaonan Nie, Pinxue Zhao+3 more|arXiv (Cornell University)|Mar 28, 2022

Domain Adaptation and Few-Shot Learning被引用数 11

ひとこと要約

HetuMoE は、複数のゲーティング戦略をサポートし、コンmodity GPU クラスタ上で訓練を高速化する、高性能で分散処理可能な混合専門家（MoE）訓練システムである。階層的 AllToAll 通信と最適化されたカーネルを導入し、バッチサイズ 32 のスイッチゲート条件下で DeepSpeed-MoE より最大 8.1× の高速化を達成した。

ABSTRACT

As giant dense models advance quality but require large amounts of GPU budgets for training, the sparsely gated Mixture-of-Experts (MoE), a kind of conditional computation architecture, is proposed to scale models while keeping their computation constant. Specifically, the input tokens are routed by the gate network and only activates part of the expert network. Existing MoE training systems only support part of mainstream MoE models (e.g. Top k) training under expensive high-bandwidth GPU clusters. In this paper, we present HetuMoE, a high-performance large-scale sparse MoE training system built on Hetu. HetuMoE provides multiple gating strategies and efficient GPU kernel implementations. To further improve the training efficiency on commodity GPU clusters (e.g, with only 1 NiC), we introduce the hierarchical AllToAll communication that combines hierarchical networks and aggregating messages. Compared with existing state-of-the-art MoE systems, HetuMoE obtains at least 15% speedup. Specifically, HetuMoE outperforms DeepSpeed-MoE up to 8.1x under the switch gate with a batch size of 32. Our code is available at: https://github.com/PKU-DAIR/Hetu.

研究の動機と目的

• コンmodity GPU クラスタ上での既存の MoE 訓練システムの非効率性とスケーラビリティの限界を解消すること。
• Switch、GShard、Top-k など、主流のゲーティング戦略を一つのフレームワークでサポートすること。
• 低帯域幅ネットワークにおける分散 MoE 訓練の通信ボトルネックを軽減すること。
• NVLink や InfiniBand といった高速インタコネクトに依存せずに、最先端の訓練パフォーマンスを達成すること。
• 標準的な GPU クラスタ上で、効率的でスケーラブルかつ実用的なトランティル・パラメータ MoE モデルの訓練を可能にすること。

提案手法

• ノード内およびノード間通信を統合した階層的 AllToAll 通信パターンを導入し、小さなメッセージを集約することで帯域幅の利用効率を向上させた。
• Top-k ゲーティング、レイアウト変換、AllToAll などの主要な MoE 操作に最適化された GPU カーネルを採用し、k=1 および k=2 の場合にアルゴリズム的改善を実施した。
• 二段階のデータレイアウト変換を採用：まずノード内でのエキスパートごとのトークングループ化を行い、次にノード間で AllToAll を実行することで、通信あたりのメッセージサイズを削減した。
• Hetu の分散ディープラーニングフレームワークを活用し、最小限のコード変更でネイティブに MoE 訓練をサポートし、高いパフォーマンスを実現した。
• Switch、GShard、M6、BASE Layer、Hash Layer、SAM、Dense-to-Sparse など、複数のゲーティング戦略をサポートし、広範なモデル互換性を実現した。
• ゲート計算、レイアウト変換、GPU 間通信、エキスパート処理、逆レイアウト変換を含む、MoE 訓練パイプライン全体を最適化した。

実験結果

リサーチクエスチョン

• RQ1 NVLink や InfiniBand といった高帯域幅インタコネクトを必要とせずに、コンmodity GPU クラスタ上で高性能な MoE 訓練システムを実現できるか？
• RQ2 階層的 AllToAll 通信は、マルチノード MoE 訓練における通信オーバーヘッドをどの程度軽減できるか？
• RQ3 Top-k ゲーティングおよびレイアウト変換のカーネルレベル最適化は、MoE 訓練効率をどの程度向上できるか？
• RQ4 さまざまなゲーティング戦略およびバッチサイズにおいて、HetuMoE は DeepSpeed-MoE や FastMoE、Tutel といった最先端システムと比較してどの程度のパフォーマンスを示すか？
• RQ5 統合されたシステムとして、幅広い MoE ゲーティング戦略を効率的にサポートしつつ、高い訓練スループットを維持できるか？

主な発見

• HetuMoE は、さまざまな構成とゲーティング戦略において、既存の MoE システムと比較して最低でも 15% の高速化を達成した。
• バッチサイズ 32 のスイッチゲート条件下では、HetuMoE は DeepSpeed-MoE より最大 8.1× の訓練速度向上を達成した。
• 階層的 AllToAll 最適化は、4×8 GPU 環境で通常の AllToAll と比較して 1.66× の高速化を、8×8 GPU 環境では 2× の高速化を達成した。
• HetuMoE の最適化された Top-k カーネルは、k=1 および k=2 の場合に PyTorch 実装と比較して平均で 25% のパフォーマンス向上を実現した。
• HetuMoE のレイアウト変換カーネルは、最先端の実装と比較して 26% 以上の改善を達成した。
• HetuMoE は、Switch、GShard、M6、BASE Layer、Hash Layer、SAM、Dense-to-Sparse など、多様な MoE ゲーティング戦略を効果的にサポートし、広範なモデル互換性を実現した。