专家并行再升级，开源MoE模型训练效率提升9倍_moe 专家并行

ColossalAI-MoE

增大数据量、增大模型参数量是公认提升神经网络性能的最直接方法，GPT3将参数扩展到了千亿级别并且这种形势还会继续进行下去。极大的模型虽然带来了巨大的性能提升，但同时也带来了巨大的算力需求。Mixture of Experts（MoE）通过选择性地激活参数，能够高效地增大模型的参数量。最近，它的潜力吸引了越来越多的关注。

我们很高兴能够宣布ColossalAI-MoE模块进行了一次全面升级，它专门为MoE模型设计改进，以帮助用户更高效、更稳定地训练和使用专家模型，从训练到推理，我们提供了完整的解决方案，为您的项目注入全新的活力和创新力！以下是我们的权限特性：

1. OpenMoE支持：我们为OpenMoE（开源社区首个仅包含解码器的MoE transformer）提供了完善的训练与推理系统，让您能够更轻松地掌握MoE的娴熟应用。

2. 异构并行：针对MoE模型训练的独特特点，我们引入了创新的EZ-MoE技术，以更高效的方式进行异构并行，使您的工作更加便捷，更有成效。

3. 底层优化：我们不仅支持MoE的高效应用，还提出了一系列底层优化方法，确保您的MoE模型达到卓越的性能水平。

开源地址：

https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/openmoe

Mixture of Experts (MoE)  

在主流深度神经网络方法中，所有模型参数都会一同参与处理输入数据。这导致随着预训练模型参数的增加，计算资源需求和成本也会大幅上升。然而，MoE技术改变了这一状况。MoE允许我们有选择性地激活模型中的一部分参数来处理不同的输入数据，这意味着随着参数数量的增加，计算资源的需求可以保持相对稳定。这种方法能够有效地管理计算资源，降低大规模模型的计算负担。具体来说，在Transformer模型中，MoE的实现方式如下图所示。

​​​​​​

在每一层中，注意力部分的结构保持不变，但在多层感知器（MLP）的部分，有多组专业化的前馈神经网络，并通过一个门控网络（Router）来确定哪些专家应该被激活以处理特定的输入数据。这一创新让模型更智能地分配计算资源，提高了效率。

OpenMoE

OpenMoE 是开源社区中第一个仅包含解码器部分的 MoE Transformer 模型。OpenMoE的改进是基于ST-MoE，经过在 The RedPajama 和 The Stack Dedup 数据集上的训练而来。我们在 BigBench-Lite 上对 OpenMoE 进行了性能测试，采用 Few-Shot 学习方式，与 BIG-G、BIG-G-Sparse（由 Google 开发的 MoE 模型）以及 GPT-3 进行了比较，评估了相对成本下的模型性能。相对成本是通过激活的参数数量与训练过的标记数量之间的比率来计算的。每个数据点的大小表示激活的参数数量，而浅灰色的数据点表示了 MoE 模型的总参数数量。显而易见，在相同的成本条件下，OpenMoE 表现出更出色的性能。

OpenMoE采用Jax实现，而ColossalAI-MoE则首次为该模型提供了高效的PyTorch开源支持，从而使更多用户能够参与和使用该模型。我们在ColossalAI-MoE中采用了最新的异构并行方法和底层优化，这使我们的训练性能大幅领先于PyTorch FSDP和DeepSpeed-MoE。以OpenMoE-8B为例，我们的训练性能最高比PyTorch快7.9倍，比DeepSpeed快42%，如下图所示：

现有专家并行的局限性

通常情况下，在进行Mixture of Experts（MoE）的训练时，我们采用一种叫做专家并行（Expert Parallel）的方法。专家并行的思路是将不同的专家分配到不同的计算设备上，这有助于减少内存消耗并提高训练效率。在这个过程中，每个设备会根据MoE模型的路由规则，将自身的数据发送到相应专家所在的设备上，然后等待专家完成计算后再将结果返回到原设备。这个过程中，每两个专家之间需要进行通信，这个通信过程被称为AlltoAll通信。

然而，专家并行带来的一个主要问题是AlltoAll通信需要的时间较长，而且计算负载不够均衡。这两个问题的根本原因在于专家并行的规模太大，带来了以下两个主要挑战：1）AlltoAll通信的规模和频率呈指数级增长，导致通信时间增加，从而降低了整体训练效率。2）即使尽量进行负载均衡，每个计算设备上分配的专家数量仍然很少，这导致了单个批次内的负载不均衡问题，这一问题在训练过程中会引发显著的性能波动。

然而，当我们尝试降低专家并行的规模，或者在机器数量多于专家数量时，我们需要引入额外的并行维度，如数据并行，来将专家复制多份，或者使用模型并行来将单个专家分布到多个计算设备上。不过，这也带来了一些新问题，比如使用数据并行时，MoE模型的参数非常庞大，导致参数通信占用大量时间。另一方面，使用模型并行会降低计算效率并增加额外的通信时间。这些是在处理大规模MoE模型时需要面对的挑战。

异构专家并行

为了应对上述问题，我们提出了一种解决方案：异构专家并行，即EZ-MoE（HierachicalEP-ZeROMoE）。在EZ-MoE中，我们结合了异构ZeRO并行与专家并行，旨在以不损失性能的方式减少EP的并行规模，从而缓解了之前提到的问题。

传统的数据并行方法需要在不同设备之间同步参数，这导致了大量的通信时间开销。相比之下，我们提出的新型并行结构采用了一种不同的方式。对于MoE模块，我们在同一设备内复制多份拷贝，只在设备内进行参数通信，而将跨设备通信限定为AlltoAll通信。这种设计背后的理念是：1）尽管MoE模块的参数数量大幅增加，但张量的大小并没有改变，这导致AlltoAll通信的数据量通常明显小于参数通信的数据量，通常只占3%-10%。2）尽管AlltoAll通信的效率相对较低，但由于其受机器间通信带宽的限制较小，传输时间更短。3）专家并行规模得到了缩减。4）计算能够弥补AlltoAll通信的时间开销。综合考虑，AlltoAll通信成为了更为优化的选择。

然而，这种方法也导致了每张卡上专家所占用的显存大幅增加。为了解决这个问题，我们采用了异构的ZeRO并行策略。对于非MoE模块，我们采用全局ZeRO并行，因为这些模块的参数相对较少。在OpenMoE中，非MoE模块的参数通常仅占总参数量的20%左右，因此完全可以通过异步通信来处理这部分数据。对于MoE模块，由于机内互联速度快，而且每台设备上只有少数专家，因此通信量也不会很大，可以通过计算来弥补，从而实现无损的优化。此外，我们还引入了流水线并行，在大规模训练时有效减少了通信时间。

负载均衡

现有的负载均衡算法很少考虑到当下MoE模型的布局和通信开销。根据我们的并行方式，我们开发了一种专为现有专家排布而设计的负载均衡算法。其目标是以最小的通信代价，基于当前专家的分布，实现负载均衡。

我们设计了一款轻量级算法，基于beam search的方法。首先，我们在机内使用beam search策略进行专家的交换，逐个考虑不同的专家交换组合。我们根据交换后的负载均衡情况和相应的通信开销来评估每个交换组合的得分，然后选择前k个得分最高的组合，继续执行专家交换，直到实现负载均衡。相比于传统的动态规划方法，这种方式更加高效，并且能够应用于已有的专家排布，提供了更多分析的可能性。

我们采用这种方法来搜索机内最佳的交换可能性。如果在机内无法获得令人满意的结果，我们随后将扩展至机器间的交互，并且最小化机器间的通信代价。

底层优化

通过流水线重叠的通信和计算过程

在MoE中，AlltoAll通信模块通常被认为是最耗时的部分。为了解决这个问题，我们采用了流水线并行的方式，将MoE中的AlltoAll通信和计算过程巧妙地相互重叠，从而显著降低了通信所需的时间，这项技术在多机训练中能够减少大量的通信时间。此外，NCCL中AlltoAll算子的实现并未进行专门的优化，其仅通过多个并行的Send/Recv操作进行实现，因此无法很好地适配节点内和节点间不同的通信环境。为此，我们实现了层次化的AlltoAll算子，充分利用了混合网络的优势。该算子首先利用NVLink的带宽优势在节点内进行消息聚合和重排列，当消息聚合完成后才使用节点间网络进行通信，极大增加了单次的通信量，随后算子会对数据再次进行重排列并发送到不同的GPU上，保持算子原有的语义。

NCCL的AlltoAll实现示意图

层次化AlltoAll实现示意图

同时，对于OpenMoE中较为耗时的模块，例如MoE路由器、大型Embedding层以及FFN等模块，我们采用了一系列高效的优化策略，如kernel fusion和计算流程的优化。此外，我们还专门定制了高效的算子，将计算效率提升了30%以上。另外，我们引入了Flash Attention技术，以进一步提高计算效率。

这一系列底层优化策略使得ColossalAI-MoE在通信和计算方面均表现出色，从而大幅提高了整体性能。

开源地址：

https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/openmoe

潞晨科技和Colossal-AI

潞晨科技由新加坡国立大学尤洋教授创立，近期获得数亿元A轮融资，已在成立18个月内迅速连续完成三轮融资，投资方为创新工场、真格基金、蓝驰创投等。团队核心成员全部来自美国加州伯克利、斯坦福、清华、北大、新加坡国立、南洋理工大学等国际知名高校，具备谷歌、微软、英伟达、腾讯等知名厂商工作经历。

低成本AI大模型开发系统Colossal-AI (https://github.com/hpcaitech/ColossalAI)，基于PyTorch，可通过高效多维并行、异构内存等，降低AI大模型训练/微调/推理的开发与应用成本，提升模型任务表现，降低GPU需求等。作为当下最火热的开源AI大模型解决方案，Colossal-AI，目前已获得GitHub星数近4万颗，位列全球TOP400，细分赛道排名世界第一，并成功入选NeurIPS, SC, AAAI, PPoPP, CVPR, ISC, NVIDIA GTC等国际AI 与HPC 顶级会议的官方教程。

Colossal-AI已与多家世界/中国五百强、东南亚科技巨头、亚洲地区国家级研究机构、海外超算中心等开展合作，促进AI大模型商业化落地，合作涉及云计算、芯片设计、生物医药、自动驾驶、智能零售、互联网等领域，并已与多家知名厂商联合开发/优化千亿/百亿参数大模型或打造垂类模型。

参考链接

Fedus, William, Barret Zoph, and Noam Shazeer. "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity." The Journal of Machine Learning Research 23.1 (2022): 5232-5270.
Fuzhao Xue, Zian Zheng, Yao Fu. "OpenMoE: Open Mixture-of-Experts Language Models." https://github.com/XueFuzhao/OpenMoE.
Nie, Xiaonan, et al. "HetuMoE: An efficient trillion-scale mixture-of-expert distributed training system." arXiv preprint arXiv:2203.14685 (2022).
Zhai, Mingshu, et al. "{SmartMoE}: Efficiently Training {Sparsely-Activated} Models through Combining Offline and Online Parallelization." 2023 USENIX Annual Technical Conference (USENIX ATC 23). 2023.
Rajbhandari, Samyam, et al. "Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale." International Conference on Machine Learning. PMLR, 2022.

如果您有兴趣了解更多关于潞晨科技的最新动态，以及进行云平台测试，请扫描下方二维码关注我们的微信公众号。

扫码添加潞晨官方微信小助手，加入潞晨官方分享群提问讨论！我们有专业的工程师答疑。