[論文レビュー] MoESys: A Distributed and Efficient Mixture-of-Experts Training and Inference System for Internet Services

Dianhai Yu, Liang Shen+6 more|arXiv (Cornell University)|May 20, 2022

Advanced Neural Network Applications被引用数 12

ひとこと要約

SE-MoE は、2次元プリフェッチ、階層的ストレージ統合、CPU-GPU メモリリングベースの推論を活用して、負荷バランス、通信効率、メモリ利用効率を向上させる分散型で効率的な混合専門家（MoE）トレーニングおよび推論システムである。DeepSpeed よりも 33% の高いトレーニングスループットと 13% の高い推論スループットを達成しており、非均一な MoE ワークロード下ではスループットが 64% 高く、メモリ使用量が 18% 低い。48 個の A100 GPU を用いて 8 日間で 120億パラメータの UFO モデルを正常にトレーニングした。

ABSTRACT

While modern internet services, such as chatbots, search engines, and online advertising, demand the use of large-scale deep neural networks (DNNs), distributed training and inference over heterogeneous computing systems are desired to facilitate these DNN models. Mixture-of-Experts (MoE) is one the most common strategies to lower the cost of training subject to the overall size of models/data through gating and parallelism in a divide-and-conquer fashion. While DeepSpeed has made efforts in carrying out large-scale MoE training over heterogeneous infrastructures, the efficiency of training and inference could be further improved from several system aspects, including load balancing, communication/computation efficiency, and memory footprint limits. In this work, we present a novel MoESys that boosts efficiency in both large-scale training and inference. Specifically, in the training procedure, the proposed MoESys adopts an Elastic MoE training strategy with 2D prefetch and Fusion communication over Hierarchical storage, so as to enjoy efficient parallelisms. For scalable inference in a single node, especially when the model size is larger than GPU memory, MoESys builds the CPU-GPU memory jointly into a ring of sections to load the model, and executes the computation tasks across the memory sections in a round-robin manner for efficient inference. We carried out extensive experiments to evaluate MoESys, where MoESys successfully trains a Unified Feature Optimization (UFO) model with a Sparsely-Gated Mixture-of-Experts model of 12B parameters in 8 days on 48 A100 GPU cards. The comparison against the state-of-the-art shows that MoESys outperformed DeepSpeed with 33% higher throughput (tokens per second) in training and 13% higher throughput in inference in general. Particularly, under unbalanced MoE Tasks, e.g., UFO, MoESys achieved 64% higher throughput with 18% lower memory footprints.

研究の動機と目的

• 負荷の不均衡、通信のボトル neck、高いメモリ使用量によって引き起こされる分散 MoE トレーニングおよび推論の非効率性を是正すること。
• 異種の GPU-CPU クラスタ上でスケーラブルで高スループットの MoE トレーニングおよび推論を可能にすること。
• GPU メモリ容量を超える MoE モデルをデプロイする際のメモリ圧力を軽減し、システムの利用効率を向上させること。
• 革新的なシステムレベルの技術を用いて、MoE システムにおける通信および計算ワークロードを最適化すること。
• インターネットサービスにおける大規模 MoE モデルのプロダクションスケールでのデプロイを、最小限のリソースオーバーヘッドで可能にすること。

提案手法

• 2次元プリフェッチを用いたエラスティック MoE トレーニングを導入し、デバイス間でのデータ並列性およびモデル並列性の効率を向上させる。
• 階層的ストレージ上での統合通信を採用し、ノード間通信のオーバーヘッドを低減し、帯域幅の利用効率を向上させる。
• CPU-GPU メモリリング構造を用いて、GPU メモリ容量を超える MoE モデルの効率的でラウンドロビン方式の推論を実現する。
• データ並列性における埋め込み分割を適用し、GPU メモリ消費量を削減し、トレーニングスループットを向上させる。
• 階層的ストレージおよびメモリ管理を活用し、データ移動と D2H/H2D オーバーヘッドを最小限に抑える。
• 最適化されたルーティングおよび負荷分散戦略を適用し、MoE トレーニング中のエキスパートの不均衡を緩和する。

実験結果

リサーチクエスチョン

• RQ1 分散かつ異種のコンピューティング環境において、MoE トレーニングおよび推論をどのようにしてより効率化できるか？
• RQ2 MoE トレーニングにおける通信およびメモリのボトル neck を軽減するためのシステムレベル最適化は何か？
• RQ3 GPU メモリ容量を超える MoE モデルを推論時に効率的に実行するにはどうすればよいか？
• RQ4 システムレベルの設計によって、MoE ルーティングにおける負荷の不均衡はどの程度軽減可能か？
• RQ5 統合されたシステムは、MoE トレーニングおよび推論において、高スループットと低メモリフットプリントの両方を達成できるか？

主な発見

• SE-MoE は、標準的な MoE ワークロードにおいて DeepSpeed よりも 33% の高いトレーニングスループット（1秒あたりのトークン数）を達成した。
• SE-MoE は、一般の MoE タスクにおいて DeepSpeed よりも 13% の高い推論スループットを達成した。
• UFO モデルのような非均一な MoE ワークロード下では、SE-MoE は DeepSpeed よりも 64% の高いスループットと 18% の低いメモリフットプリントを達成した。
• 本システムは、48 個の A100 GPU を用いて 120億パラメータの統合特徴最適化（UFO）MoE モデルを 8 日間で正常にトレーニングした。
• データ並列性における埋め込み分割により、V100 GPU 上で GPU メモリ使用量が最大 26.3% 減少し、スループットが最大 15.6% 向上した。
• CPU-GPU メモリリング推論機構により、ラウンドロビンタスクスケジューリングを用いて、GPU メモリ容量を超える MoE モデルの効率的実行が可能になった。