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.

翻译：稀疏门控混合专家（MoE）技术已在极端规模深度神经网络的扩展中展现出显著效能。尽管已有大量研究从模型设计或系统优化角度提升MoE性能，现有MoE调度模式仍无法充分挖掘异构网络环境的潜力。本文从模型-系统协同设计视角提出TA-MoE——一种面向大规模MoE训练的拓扑感知路由策略，可根据网络拓扑动态调整MoE调度模式。基于通信建模，我们将调度问题抽象为优化目标，并获取不同拓扑下的近似调度模式。在此基础上，我们设计了拓扑感知辅助损失函数，该函数能自适应路由数据以适应底层拓扑结构且不牺牲模型精度。实验表明，TA-MoE在各种硬件和模型配置下均显著优于同类方法，相比于主流DeepSpeed-MoE、FastMoE及FasterMoE，性能提升幅度分别约为1.01-1.61倍、1.01-4.77倍和1.25-1.54倍。