MOE原理介绍和代码解读

Himon
65 阅读5分钟

MEGABLOCKS

Switch Transformers

  • 模型介绍

Swith Transformer的MoE设计是门控网络每次只route到1个专家。采用softmax来计算门控网络层的输出。

  • 可微的负载均衡loss（Load Balancing）：

    •   为了鼓励所有专家尽可能的负载均衡，拿到近似均匀分布的样本，本文加入了 负载均衡损失 ，这个loss在训练中直接加到总的loss上。当fi=pi=1/Nf_i= p_i= 1/N的时候损失是最小的，表示所有token完全平均分配到N个专家上。具体的，N（1,...,N）表示N个专家，B表示batch，T表示当前batch中token的数量，超参数α=10−2\alpha=10^{-2}是一个非常小的权重因子：
    •   loss=α∗N∗∑i=1Nfi∗Piloss = \alpha*N*\sum^N_{i=1}{f_i*P_i}
    •   其中，fif_i表示专家i被分配到token的数量的占比，比如有100个token，专家i分配到了10个，那么fi=1/10f_i=1/10：
    •   其中，PiP _i表示所有token分配到专家i 的概率占比，也就是求所有token经过门控层之后分配到专家i的平均概率**：

  • 数据并行/模型并行/专家并行说明

B每个batch中token的数量，N：core的数量，n：数据并行的分区数量，m：模型并行的分区数量，E：专家的数量，C：Expert capacity，每个专家的batch size。

  1. 数据并行：上图中的第一列。n=N，m=1，模型weight权重在所有core上复制一份（第一列上图），数据是切分成N份分配给所有core，计算时分别处理一部分（第一列上图）。这种模式，计算前、反向传播过程中不需要通信，当所有cores计算结束之后需要将梯度聚合起来，也就是只在优化器阶段同步所有core。

  2. 模型并行：上图第二列。n=1，m=M，模型weight权重被切分到16个core上，每个core都会处理全部的B个token。因为dffd_{ff}层需要切分到每个core上（按照维度切分），在计算FFN层第二个矩阵乘法时，每个core上都送过去 [B,dmodel][B,d_{model}]的数据，然后所有core上的结果需要重新sum起来，因此在前向、后向计算时都会使用 all-reduce操作。

  3. 数据并行+模型并行：上图第三列。是数据并行和模型并行的结合。一共有N=n*m个core，每个core负责计算 B/n个tokens以及分配到的dff/md_{ff} / m的模型权重和中间activation产出。在前向以及后向计算之间时的tensor大小是[B/n,dmodel][B/n,d_{model}]。这种模式cores被分为4组，每组共享一份全量的模型权重，4组模型权重做数据并行，4组模型组内切分做模型并行。数据切分，不同组模型处理不同的数据（不同颜色区分），在一个组内4个core做模型并行，数据被辅助成多份共用（如图中蓝色被复制成多份）。

  4. 专家并行+数据并行：上图第四列。每个专家分布到N个core上，互相独立，权重不同所以颜色不同。下方数据并行，数据是切分成N份分配给所有core，计算时分别处理一部分。

  5. 专家并行+数据并行+模型并行：从模型权重切分看，跟3模式相同，区别在于expert权重各不相同。从数据切分看，跟3模型相同。在设计模型的时候，我们需要平衡每个token的FLOPS和模型参数量，当我们增加了expert的数量，就会增加模型参数，但是并没有改变每个token的FLOPS。为了增加FLOPS，我们只能增加dffd_{ff}层的维度（这也意味着会增加参数量），这就引出一个权衡的问题：如果我们增加dffd_{ff}层的维度，会导致每个core计算OOM（out-of-memory），这个就不可避免的需要增加模型并行的分区数量m。但是，我们只有固定数量的core数量N，根据N=n∗mN=n*m ， 我们就必须减少n，也就是需要使用更小的batch-size。所以，在实际模型设计中，需要平衡FLOPS、通信消耗以及core的内存是一件非常复杂的事情，需要具体问题具体分析。

Mixtral 8x7B

Mixtral 8x7B 是在Mixtral 7B基础上发展的， “8”的意思是每层由8个专家组成。

如上图所示，MoE模块的输出是所有专家网络输出的加权和，每个专家网络就是前馈神经网络，权重是由门控网络生成，假设有n个专家网络，其中G(x)iG(x)_i是门控网络的输出向量（n维的概率向量），Ei(x)E_i(x)是第i个专家网络的输出向量，则MoE的输出可以表示成：

∑i=0n−1G(x)i∗Ei(x)\sum^{n-1}_{i=0}{G(x)_i}*E_i(x)

如果门控网络的输出概率向量是稀疏的，那么对应概率为零的专家网络就不需要参与计算。本文的实现方式非常简单，对于门控网络层的输出logits，先选取top-K，然后计算softmax进行归一化：

G(x):=Softmax(TopK(x∗Wg))G(x):=Softmax(TopK(x*W_g))

在本文中K设定为2，专家网络使用SwiGLU激活函数。则输入第x个token的最终输出y可以计算为：

y=∑i=0n−1Softmax(Top2(x∗Wg))i∗SwiGLUi(x)y=\sum^{n-1}_{i=0}{Softmax(Top2(x*W_g))_i * SwiGLU_i(x)}

DeepSeek-V3

参考：

  1. 全网首篇从tensorRT-LLM MoE CUDA kernel角度理解Mixtral-8x7b的推理加速及展望
  2. 混合专家模型 (MoE) 详解
  3. Mixtral Moe代码解读_moe 代码-CSDN博客
  4. 一文带你看清Mixtral内部结构及参数计算 翟泽鹏
  5. 图解大模型训练系列之：DeepSpeed-Megatron MoE并行训练（原理篇）
  6. 用通俗易懂的方式讲解大模型分布式训练并行技术:MOE并行_如何训练moe大模型-CSDN博客
  7. moe使用负载均衡的意义仅仅是为了方便训练吗？
  8. 混合专家模型 (MoE) 详解
  9. 分析transformer模型的参数量、计算量、中间激活、KV cache
avatar
文章
阅读
粉丝