中文

FSMoE:稀疏混合专家模型的灵活且可扩展训练系统

机器学习 2025-01-22 v1

摘要

近年来,大型语言模型(LLM)倾向于利用稀疏性来降低计算量,采用稀疏激活的混合专家(MoE)技术。MoE 引入四个模块,包括 token 路由、token 通信、专家计算和专家并行,这些模块影响模型质量和训练效率。为实现 MoE 模型的灵活用途,我们引入 FSMoE,这是一个灵活的训练系统,通过三项新技术优化任务调度:1)对 MoE 模块进行统一抽象和在线分析,以实现跨各种 MoE 实现的任务调度;2)协调节点内外通信与计算,以最小化通信开销;3)为实现近最优任务调度,我们设计了用于梯度聚合的自适应分区方法以及一种自适应流水线化通信和计算的调度方案。我们在两个 GPU 集群上对配置的 MoE 层和真实世界的 MoE 模型进行了广泛实验。实验结果表明:1)我们的 FSMoE 支持四种流行的 MoE 路由函数类型,并且比现有实现更高效(最高可实现 1.42×\times 加速),且 2)FSMoE 在 1458 个 MoE 层上优于最先进的 MoE 训练系统(DeepSpeed-MoE 和 Tutel)实现了 1.18×\times-1.22×\times 加速,在基于 GPT-2 和 Mixtral 的真实世界 MoE 模型上实现了 1.19×\times-3.01×\times 加速。

关键词

引用

@article{arxiv.2501.10714,
  title  = {FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models},
  author = {Xinglin Pan and Wenxiang Lin and Lin Zhang and Shaohuai Shi and Zhenheng Tang and Rui Wang and Bo Li and Xiaowen Chu},
  journal= {arXiv preprint arXiv:2501.10714},
  year   = {2025}
}