SE-MoE: A Scalable and Efficient Mixture-of-Experts Distributed Training and Inference System

要約

現在、ML インフラストラクチャの多様性が増しており、大きなモデルの作成を容易にするために、異種コンピューティング システム上での分散トレーニングが求められています。
専門家混合 (MoE) モデルは、分割統治方式でのゲーティングと並列処理を通じて、モデル/データの全体的なサイズに応じてトレーニングのコストを削減するために提案されています。
DeepSpeed は、異種インフラストラクチャ上で大規模な MoE トレーニングを実行することに取り組んできましたが、トレーニングと推論の効率は、負荷分散、通信/計算効率、メモリ フットプリント制限などのいくつかのシステム側面からさらに改善される可能性があります。
この研究では、さまざまなタイプで効率的な並列処理を楽しむために、2D プリフェッチと階層ストレージ上の Fusion 通信を使用した Elastic MoE トレーニングを提案する SE-MoE を紹介します。
単一ノードでのスケーラブルな推論の場合、特にモデル サイズが GPU メモリよりも大きい場合、SE-MoE は CPU と GPU メモリを結合してセクションのリングに形成し、モデルをロードし、メモリ セクション全体で計算タスクを実行します。
効率的な推論のためのラウンドロビン方式。
SE-MoE を評価するために広範な実験を実施しました。SE-MoE は、48 枚の A100 GPU カード上で 8 日間で 12B パラメーターのスパースリーゲート専門家混合モデルを使用して統合機能最適化 (UFO) モデルをトレーニングすることに成功しました。
最先端技術との比較では、SE-MoE が DeepSpeed よりも優れており、トレーニングではスループット (1 秒あたりのトークン数) が 33% 高く、推論では全体的に 13% 高いスループットを示しています。
特に、UFO などの不均衡な MoE タスクの下では、SE-MoE は 18% 低いメモリ フットプリントで 64% 高いスループットを達成しました。
フレームワークのコードは https://github.com/PaddlePaddle/Paddle でリリースされます。

要約(オリジナル)

With the increasing diversity of ML infrastructures nowadays, distributed training over heterogeneous computing systems is desired to facilitate the production of big models. Mixture-of-Experts (MoE) models have been proposed 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 SE-MoE that proposes Elastic MoE training with 2D prefetch and Fusion communication over Hierarchical storage, so as to enjoy efficient parallelisms in various types. For scalable inference in a single node, especially when the model size is larger than GPU memory, SE-MoE forms 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 SE-MoE, where SE-MoE 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 SE-MoE 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, SE-MoE achieved 64% higher throughput with 18% lower memory footprints. The code of the framework will be released on: https://github.com/PaddlePaddle/Paddle.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google