SMARTMOE：通过结合离线与在线并行化高效训练稀疏激活模型

本文链接：https://blog.youkuaiyun.com/m0_49448331/article/details/145895830

SMARTMOE: Efficiently Training Sparsely-Activated Models through Combining Offline and Online Parallelization

研究背景

背景介绍:
这篇文章的研究背景是深度神经网络（DNN）的规模不断增大，导致训练这些模型需要消耗大量的计算资源。为了减少训练成本并扩大模型规模，稀疏激活模型（如Mixture-of-Experts, MoE）被越来越多地提出和应用。然而，现有的自动并行化方法难以应用于这些稀疏模型，因为稀疏模型对数据敏感且工作负载不均衡。
研究内容:
该问题的研究内容包括设计一个自动并行化系统SMARTMoE，用于稀疏激活模型的分布式训练。SMARTMoE通过结合离线和在线并行化来优化训练过程，以应对稀疏模型的动态和不平衡特性。
文献综述:
该问题的相关工作包括针对密集神经网络的自动并行化系统，但这些系统通常假设工作负载是静态的，无法有效处理稀疏模型的动态特性。此外，现有的MoE模型训练系统主要采用特定的专家并行性策略，缺乏对混合并行性的支持。

研究方法

这篇论文提出了SMARTMoE系统，用于解决稀疏激活模型的并行化问题。具体来说，

扩展的混合并行性空间:
SMARTMoE在混合并行性空间中寻找优化机会，考虑了数据敏感模型的工作负载。通过将空间分解为离线静态池，并在运行时在线选择，SMARTMoE能够在动态和不平衡的MoE训练过程中保持高效率。
两阶段自适应自动并行化:
- 离线阶段: 创建一个策略池，保证良好的内在性能和低切换开销。设计了一个工作负载感知的性能模型，用于估计数据敏感模型的性能，以便在训练前构建最优池。
- 在线阶段: 开发轻量级算法，在选定的池内找到当前工作负载的最优并行策略。算法定期在运行时执行，以确定是否应采用新的并行策略。
工作负载感知的性能建模:
通过估计门控网络的选择分布来实现工作负载感知的性能建模。使用最大可能的工作负载来准确预测池的性能，避免了实际训练数据的统计。

实验设计

实验平台:
在三个不同的集群上评估SMARTMoE，集群配置包括不同数量的GPU和网络带宽。
模型选择:
使用GPT-MoE和Swin-MoE模型进行评估，分别用于自然语言处理和计算机视觉任务。模型参数随着GPU数量的增加而增加。
基线系统:
将SMARTMoE与四个强大的训练系统进行比较，包括DeepSpeed-MoE、Tutel、FasterMoE和Alpa。

结果与分析

端到端加速:
SMARTMoE在三个集群上实现了高达1.88倍的端到端训练加速。在inky集群上，GPT-MoE模型平均加速1.53倍，最高加速1.88倍；在pinky集群上，平均加速1.17倍。
离线并行化消融研究:
SMARTMoE的数据感知方法相比数据不敏感的方法实现了更高的加速效果，验证了SMARTMoE能够找到更好的池。
在线并行化消融研究:
SMARTMoE在每个层面上平均实现了1.16倍的加速，最多在某一层实现了1.43倍的加速。
细粒度性能分析:
SMARTMoE在通信和计算开销方面优于Alpa，验证了工作负载感知的优势。

结论

SMARTMoE通过扩展混合并行性空间并引入两阶段自适应自动并行化方法，解决了稀疏激活模型的并行化问题。与现有系统相比，SMARTMoE在端到端训练中实现了高达1.88倍的加速。SMARTMoE的设计为稀疏激活模型的分布式训练提供了有效的解决方案。
这篇论文通过创新的方法和实验验证，展示了SMARTMoE在稀疏模型训练中的高效性和实用性。