TA-MoE: Topology-Aware Large Scale Mixture-of-Expert Training

要約

まばらにゲートされた Mixture-of-Expert (MoE) は、ディープ ニューラル ネットワークを極端な規模にスケールアップする際の有効性を実証しています。
モデル設計またはシステム最適化の観点から MoE のパフォーマンスを改善するために多くの努力が払われてきたにもかかわらず、既存の MoE ディスパッチ パターンは依然として、基盤となる異種ネットワーク環境を十分に活用できていません。
このホワイト ペーパーでは、ネットワーク トポロジに応じて MoE ディスパッチ パターンを動的に調整できる、モデル システム共同設計の観点から、大規模な MoE トレーニングのためのトポロジ対応ルーティング戦略である TA-MoE を提案します。
通信モデリングに基づいて、ディスパッチ問題を最適化目標に抽象化し、さまざまなトポロジーの下でのおおよそのディスパッチ パターンを取得します。
その上で、モデルの精度を犠牲にすることなく、基本的なトポロジに適合するようにデータを適応的にルーティングできる、トポロジを意識した補助損失を設計します。
実験によると、TA-MoE はさまざまなハードウェアおよびモデル構成で対応するものよりも大幅に優れており、一般的な DeepSpeed-MoE、FastMoE、および FasterMoE よりも約 1.01x-1.61x、1.01x-4.77x、1.25x-1.54x の改善が見られます。

要約(オリジナル)

Sparsely gated Mixture-of-Expert (MoE) has demonstrated its effectiveness in scaling up deep neural networks to an extreme scale. Despite that numerous efforts have been made to improve the performance of MoE from the model design or system optimization perspective, existing MoE dispatch patterns are still not able to fully exploit the underlying heterogeneous network environments. In this paper, we propose TA-MoE, a topology-aware routing strategy for large-scale MoE trainging, from a model-system co-design perspective, which can dynamically adjust the MoE dispatch pattern according to the network topology. Based on communication modeling, we abstract the dispatch problem into an optimization objective and obtain the approximate dispatch pattern under different topologies. On top of that, we design a topology-aware auxiliary loss, which can adaptively route the data to fit in the underlying topology without sacrificing the model accuracy. Experiments show that TA-MoE can substantially outperform its counterparts on various hardware and model configurations, with roughly 1.01x-1.61x, 1.01x-4.77x, 1.25x-1.54x improvements over the popular DeepSpeed-MoE, FastMoE and FasterMoE.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google