FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models

📄 arXiv: 2501.10714v1 📥 PDF

作者: Xinglin Pan, Wenxiang Lin, Lin Zhang, Shaohuai Shi, Zhenheng Tang, Rui Wang, Bo Li, Xiaowen Chu

分类: cs.LG

发布日期: 2025-01-18

DOI: 10.1145/3669940.3707272

💡 一句话要点

FSMoE：一种灵活可扩展的稀疏混合专家模型训练系统

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 稀疏激活 分布式训练 任务调度 通信优化

📋 核心要点

  1. 现有MoE训练系统在任务调度、通信优化和梯度聚合等方面存在不足，限制了模型训练效率。
  2. FSMoE通过统一抽象MoE模块、协同调度通信与计算、自适应梯度划分等技术，优化任务调度，提升训练效率。
  3. 实验表明，FSMoE在多种MoE配置和真实模型上均优于现有系统，加速比最高可达3.01倍。

📝 摘要（中文）

本文提出了一种灵活的稀疏混合专家模型（MoE）训练系统FSMoE，旨在优化任务调度，从而提高训练效率和模型质量。MoE包含token路由、token通信、专家计算和专家并行四个模块。FSMoE通过以下三种创新技术实现优化：1) 统一抽象和在线分析MoE模块，以支持跨多种MoE实现的任务调度；2) 协同调度节点内和节点间通信与计算，以最小化通信开销；3) 设计自适应梯度划分方法用于梯度聚合，并设计自适应流水线调度通信和计算，以支持近乎最优的任务调度。在两个GPU集群上，使用配置的MoE层和真实MoE模型进行了大量实验。结果表明，FSMoE支持四种流行的MoE路由函数，并且比现有实现更有效（加速高达1.42倍），并且在1458个MoE层上，FSMoE优于最先进的MoE训练系统（DeepSpeed-MoE和Tutel）1.18倍-1.22倍，在使用流行路由函数的基于GPT-2和Mixtral的真实MoE模型上，FSMoE优于最先进的MoE训练系统1.19倍-3.01倍。

🔬 方法详解

问题定义：论文旨在解决稀疏混合专家模型（MoE）训练过程中任务调度效率低下的问题。现有的MoE训练系统在处理不同MoE实现、优化通信开销以及进行高效梯度聚合方面存在局限性，导致训练速度慢，资源利用率不高。

核心思路：论文的核心思路是通过统一抽象MoE模块，实现跨不同MoE实现的灵活任务调度；通过协同调度节点内和节点间通信与计算，减少通信开销；以及通过自适应梯度划分和流水线调度，优化梯度聚合过程，从而提升整体训练效率。

技术框架：FSMoE训练系统主要包含以下几个关键模块：1) MoE模块抽象层：对不同的MoE实现进行统一抽象，提供一致的任务调度接口。2) 在线性能分析器：实时监控MoE模块的性能，为任务调度提供依据。3) 协同调度器：根据性能分析结果，协同调度节点内和节点间通信与计算任务。4) 自适应梯度划分器：根据通信带宽和计算能力，自适应地划分梯度，以优化梯度聚合过程。5) 流水线调度器：自适应地调整通信和计算的流水线执行方式，以最大化资源利用率。

关键创新：FSMoE的关键创新在于其统一的MoE模块抽象、协同调度机制和自适应梯度划分方法。与现有方法相比，FSMoE能够更好地适应不同的MoE实现，更有效地减少通信开销，并更智能地进行梯度聚合，从而实现更高的训练效率。

关键设计：在MoE模块抽象层，论文定义了一组通用的接口，用于描述不同MoE模块的行为和性能特征。协同调度器采用了一种基于优先级的调度算法，优先调度通信量大的任务。自适应梯度划分器根据通信带宽和计算能力的比例，动态调整梯度划分的大小。流水线调度器则根据任务的依赖关系和资源可用性，动态调整流水线的深度和宽度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，FSMoE在多种MoE配置和真实模型上均优于现有系统。具体来说，FSMoE支持四种流行的MoE路由函数，并且比现有实现更有效（加速高达1.42倍），并且在1458个MoE层上，FSMoE优于最先进的MoE训练系统（DeepSpeed-MoE和Tutel）1.18倍-1.22倍，在使用流行路由函数的基于GPT-2和Mixtral的真实MoE模型上，FSMoE优于最先进的MoE训练系统1.19倍-3.01倍。

🎯 应用场景

FSMoE可应用于各种需要大规模MoE模型的场景，例如自然语言处理、计算机视觉和推荐系统。通过提高MoE模型的训练效率，FSMoE能够加速新模型的开发和部署，并降低训练成本，从而推动人工智能技术的进步。

📄 摘要（原文）

Recent large language models (LLMs) have tended to leverage sparsity to reduce computations, employing the sparsely activated mixture-of-experts (MoE) technique. MoE introduces four modules, including token routing, token communication, expert computation, and expert parallelism, that impact model quality and training efficiency. To enable versatile usage of MoE models, we introduce FSMoE, a flexible training system optimizing task scheduling with three novel techniques: 1) Unified abstraction and online profiling of MoE modules for task scheduling across various MoE implementations. 2) Co-scheduling intra-node and inter-node communications with computations to minimize communication overheads. 3) To support near-optimal task scheduling, we design an adaptive gradient partitioning method for gradient aggregation and a schedule to adaptively pipeline communications and computations. We conduct extensive experiments with configured MoE layers and real-world MoE models on two GPU clusters. Experimental results show that 1) our FSMoE supports four popular types of MoE routing functions and is more efficient than existing implementations (with up to a 1.42$\times$ speedup), and 2) FSMoE outperforms the state-of-the-art MoE training systems (DeepSpeed-MoE and Tutel) by 1.18$\times$-1.22$\times$ on 1458 MoE layers and 1.19$\times$-3.01$\times$ on real-world MoE models based on GPT-2 and Mixtral using a popular routing function.