TA-MoE：面向大规模混合专家训练的拓扑感知方法
机器学习
2023-02-21 v1 人工智能
摘要
稀疏门控混合专家 (MoE) 已证明其将深度神经网络扩展到极大规模的有效性。尽管已从模型设计或系统优化角度做出大量努力以改进 MoE 性能，现有 MoE 调度模式仍无法充分挖掘底层异构网络环境。本文中，我们从模型-系统协同设计角度提出 TA-MoE，一种面向大规模 MoE 训练的拓扑感知路由策略，其可根据网络拓扑动态调整 MoE 调度模式。基于通信建模，我们将调度问题抽象为优化目标，并在不同拓扑下获得近似调度模式。在此基础上，我们设计了一种拓扑感知辅助损失，可在不牺牲模型精度的前提下自适应路由数据以适配底层拓扑。实验表明，TA-MoE 在各种硬件与模型配置上均大幅优于同类方法，相较流行的 DeepSpeed-MoE、FastMoE 和 FasterMoE 分别有约 1.01x-1.61x、1.01x-4.77x、1.25x-1.54x 的提升。
引用
@article{arxiv.2302.09915,
title = {TA-MoE: Topology-Aware Large Scale Mixture-of-Expert Training},
author = {Chang Chen and Min Li and Zhihua Wu and Dianhai Yu and Chao Yang},
journal= {arXiv preprint arXiv:2302.09915},
year = {2023}
}