Tutel: Adaptive Mixture-of-Experts at Scale

要約

Sparsely-gated mixture-of-experts (MoE)は、一定の計算コストで深層学習モデルを兆を超えるパラメータに拡張するために広く採用されています。MoEのアルゴリズム性能は、各入力トークンを適切なサブモデルまたはエキスパートに転送するトークン・ルーティング機構に依存している。トークン・ルーティングは実行時にエキスパートの作業量を動的に決定しますが、既存のシステムは静的な実行、すなわち静的な並列化とパイプラインにより、動的な作業量に適応せず、非効率な計算を強いられています。我々は、動的に適応する並列性とパイプラインを持つMoE用の拡張性の高いスタック設計と実装であるFlexを提示する。Flexは、MoEモデルのパラメータと入力データを分散するための同一のレイアウトを設計し、数学的な不等価性やテンソル移行のオーバーヘッドなしに、あらゆる並列化またはパイプライン化手法で利用することができる。これにより、実行時にゼロコストで適応的な並列化/パイプライン化の最適化が可能になります。この重要な設計に基づき、Flexは様々なMoEアクセラレーション技術も実装しています。すべての技術を統合することで、Flexは最終的にあらゆるスケールで圧倒的なスピードアップを実現しました。A100 GPUを16個、2,048個使用した場合の1つのMoEレイヤーのスピードアップは、従来の最先端と比べてそれぞれ4.96倍と5.75倍です。我々の評価では、Flexが、最先端のコンピュータビジョンアーキテクチャであるSwin Transformer V2上に構築されたSwinV2-MoEという実世界のMoEベースのモデルを効率的かつ効果的に実行することを示しています。効率面では、FlexはSwinV2-MoEを高速化し、Fairseqと比較して学習と推論でそれぞれ最大1.55倍と2.11倍の高速化を達成しました。また、SwinV2-MoEモデルは、事前学習とCOCO物体検出などのダウンストリームコンピュータビジョンタスクの両方で、対応するdenseモデルよりも優れた精度を達成し、Flexがエンドツーエンドの実世界モデルの学習と推論に適していることを示しています。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL