【元来如此】第四章——MOE专家并行再升级
文 | 软件生态中心-模型应用部-思成
正文2000字 阅读5分钟
“
在之前的文章中，我们介绍了MoE模型推理相关的探索内容，还没有了解的同学可以点击下方图片链接跳转阅读。接下来为大家带来MoE大模型并行训练的内容。
MoE架构
▲图0 Switch Transformers[1]论文中的MoE Layer
如图0所示，MoE网络主要的改进是将Transformer模型中的每个FFN层替换为MoE结构，所谓的MoE结构主要由一个Router和N个Expert组成，其中Router负责将Token路由到对应的Expert。之后在对应的专家完成对应的计算操作，最后将多个专家的结果通过Router路由概率计算加权和。
后续几年，陆续有很多关于MoE的研究。比如对于上述提到的Router，演进出了很多Token- Choice、Expert-Choice、Token Merge、Expert Merge算法。针对Expert网络结构，有在FNN和Attention两个主要方面的工作，比较有代表性的比如GShard和MoA。针对Share Expert算法衍生出了比如DeepseekMoE、Qwen1.5-MoE等方面的工作。
本篇内容会从基础的MoE结构入手，更多的集中在分布式方案设计上，并不会过多关注MoE相关算法相关内容（如果有感兴趣的同学，可以额外开章节进行介绍）。
下文内容我们做如下约定和假设：
各层MoE网络结构是同构的，图示中只展示其中一层网络
为方便分析，我们图中只展示Top1-Expert的情况
各个Router是负载均衡的，且不考虑溢出或者丢弃等
同时定义以下符号：
DP(Data Parallelism)
TP(Tensor Parallelism)
EP(Expert Parallelism)
方案0
从中易得，前向计算过程中Attention部分不存在通信，在MoE部分需要每一个DP维度内的EP4之间的All2All操作，经过MoE的计算之后，再进行一次All2All操作还原数据。 在模型反向计算过程中GPU0-GPU4，GPU1-GPU5，GPU2-GPU6，GPU3-GPU7的Expert做Allreduce操作，在Attention部分做8卡的Allreduce操作。
▲ 图1 Attention-DP8-MoE-EP4-DP2
在这样的方案设计中，我们不难得到这样的结论：非MoE部分，即Attention部分在DP维度完成Allreduce操作。在MoE部分，专家之间在EP维度完成All2All操作，在DP维度完成Allreduce操作。
方案1
接下来让我们稍微进阶一下。假设方案0中GPU卡的显存并不够。其中一种方案是将Expert维度改变成EP4-TP2，如图2所示。单机8卡，4个专家，Attention部分（蓝色）采用DP8，MoE部分采用EP4-TP2。
可以非常直观的看到，前向过程中Attention部分保持不变。MoE部分（为展示清晰，图中只画出了GPU3中的数据在MoE结构中的流程，其余7张卡同理）首先进行了一次变种的All2All通信，经过MoE计算之后，通过另一个正常的All2All操作还原数据，注意，经过还原的数据需要在各自的rank内再额外完成一次reduce操作，当然这个reduce操作并不产生额外的通信。
反向计算过程中，Attention分布保持8卡做Allreduce操作。MoE部分没有额外的通信。
▲ 图2 Attention-DP8-MoE-EP4-TP2
在这个方案中。不管是哪个All2All操作，对比方案1，通信量都增加了1倍，同时这个All2All操作会的范围扩大了1倍，这种通信模式在非两两互联的link中对性能的挑战都会进一步增大。另外，在第一次All2All的时候存在一个变种的All2All操作（虽然通过一些trick方案可以绕过，感兴趣的同学可以想想看）对于代码可读性和维护性上也存在一定的破坏。
方案2
在前向计算的时候， Attention部分和之前2个方案不同，GPU0-GPU1，GPU2-GPU3，GPU4-GPU5，GPU6-GPU7之间需要进行Allreduce通信。之后看起来在MoE部分完成2次All2All操作。细心的同学可能已经发现，如果按照常规的All2All操作，第一次会多发送1倍的重复数据。第二次进行的是一个类似与上一个方案的All2All操作。 在反向计算的时候，GPU0-GPU4，GPU1-GPU5，GPU2-GPU6，GPU3-GPU7的Expert做Allreduce操作，在Attention部分做GPU0-GPU2-GPU4-GPU6和GPU1-GPU3-GPU5-GPU7的Allreduce操作。
▲ 图3 Attention-TP2-DP4-MoE-EP4-DP2
▲ 图4 Attention-TP2-DP4-MoE-EP4-DP2
▲ 图5 Attention-TP2-DP4-MoE-EP4-DP2
方案3
前向过程中，Attention部分和之前介绍的一致。在MoE部分，首先需要一次All2All（2组，分别是GPU0-GPU2-GPU4-GPU6和GPU1-GPU3-GPU5-GPU7）操作。之后在GPU0-GPU1、GPU2-GPU3等各自完成Allreduce计算，最后再通过一次All2All操作（同样也是2组）。 这里不进一步赘述有关反向的操作流程了，相信聪明的你一定可以根据上面几种情况的学习推理出这里发生了什么。
▲ 图6 Attention-TP2-DP4-MoE-EP4-TP2
从上述分析来看，为了高效训练一个确定的MoE模型，针对不同的硬件特点，设计贴合硬件特点的分布式训练方案是一种可以有效提升性能的手段。我们以上一篇提到的Mixtral-8x7B模型为例，可以得到单层Attention参数量42million，而8个Expert对应的参数量为1409million，两者参数差异在33倍以上。从显存容量角度看，尽可能的分割专家就已经可以很好的释放显存压力，非必要的时候并不需要额外切分Attention部分的参数。
免费试用申请
如果您对我们的产品感兴趣，
可点击底部“阅读原文”，进行试用申请。
参考文献
[1] https://arxiv.org/abs/2101.03961
