LLM并行训练-moe并行学习笔记

前置知识

MOE(MixerOfExpert)

moe的主要原理是替换attention层后的MLP层, 通过将不同类型的token按照门控单元计算出的概率分配给最大概率处理的专家网络处理, 对比单一MLP更适合处理复杂多样化的数据集. 主要思想和集成学习感觉很像, 而且扩展性(遇到新的目标任务可以新增专家网络)和可解释性(每个专家分开调整)都比较强. MOE前向步骤(以最简单的top2 Expert为例):

1. 门控网络, 输入是attention的输出, dim为(batch_size, tokens, emb_size), 输出dim为(batch_size, tokens, experts_num)

<code>topkgate_linear = nn.Linear(n_embed, num_experts) # 从emb_size->专家个数的映射, 根据这个线性层计算每个token进入各个专家网络的概率

logits = topkgate_linear(mh_output)

top_k_logits, top_k_indices = logits.topk(top_k, dim=-1) # 从4个专家网络里取top2,

zeros = torch.full_like(logits, float('-inf'))

sparse_logits = zeros.scatter(-1, top_k_indices, top_k_logits) #把除了top2剩余的位置置-inf

gating_output= F.softmax(sparse_logits, dim=-1) #softmax计算进入各个专家网络的概率

2. 稀疏化Experts

这部分每个expert的网络结构都可以根据场景设计的不一样, 因为在fp/bp计算的时候, 每个token都是只进入了topk的网络进行计算, 剩余的网络没有计算. 大部分的参数都没参与更新, 所以也被称为稀疏化的dense. 因为这个特性也给moe网络的并行化改造提供了基础.

expert负载不均匀问题

因为expert的概率纯粹是训练出来的参数决定的, 没法用LALB类似的负载均衡策略强制使每个expert接收到的token是均匀的, 极有可能出现某几个expert接收到了很多, 其他的基本没啥token的问题..主要有以下这么几个解决办法:

TokenBuffer: 给每个expert设置固定容量, 容量设置公式如下, 当这个expert收满token后就不再接受token

\[𝑐𝑎𝑝𝑎𝑐𝑖𝑡𝑦=𝑚𝑎𝑥(\frac{𝑆}{𝐸}∗𝐾∗𝑐𝑎𝑝𝑎𝑐𝑖𝑡𝑦\_𝑓𝑎𝑐𝑡𝑜𝑟,𝑚𝑖𝑛\_𝑐𝑎𝑝𝑎𝑐𝑖𝑡𝑦)

\]

\(E\)：expert_num \(S\) ：token数 \(K\): topK数

在deepspeed的实现里处理token溢出的方法:

• 这个token在top2的两个expert只溢出了一个, 那么把另一个没溢出的expert的softmax权重设成1放到那个expert里(这里有点像hash线性探测那种方法haha)
• 在两个expert里都溢出了, 那么把这个token直接跳过expert通过残差的方式直连到上层

Random Routing:

随机路由的方法主要针对的是2nd的expert, 1st的直接发出去. deepspeed的2nd随机选择策略:

1. 从随机分布中采样expert_num个随机数作为噪声
2. 把噪声加到softmax的结果上, 另外把1st的mask掉(因为1st是必发的, 只需要再选一个最高的就够了)
3. 在剩下的里面找一个最高的作为2nd expert.
4. 因为随机后的不一定两个expert概率加和为1, 所以需要进行重新normalize \(P_0' = \frac{P_0}{P_0 + P_1}\) \(P_1' = \frac{P_1}{P_0 + P_1}\) 这里算的概率用于把经过2个对应expert后的token结果进行加权平均.

辅助损失函数:

\[l_{\text {aux }}=\frac{1}{E} \sum_{e=1}^E \frac{c_e}{S} * m_e

\]

• \(C_e\) ：某expert的buffer中已经存下的token数量（该expert作为1st时接收到的token数）
• \(m_e\) ：某expert的buffer中已经存下的token在该专家上的avg(weight)（token考虑范围也是那些将该专家作为1st专家的token）, 加这个参数主要是为了让这部分可导, 能够进行bp

把这个loss加到主loss后面, 目标也是最小化这个辅助loss. 因为我们最理想的情况是每个expert作为1/2的1st和1/2的2nd, 而如果某个expert溢出时他作为1st的概率远高于其他expert, 而溢出后根据上面处理溢出的方法会把2nd转成1st, 使得1st变多. 所以我们minimize 辅助loss的时候其实就是让网络学习时尽量避免溢出.

MOE并行(gshard)

施工中

参考

gshard论文: https://arxiv.org/pdf/2006.16668

Deepspeed-moe: https://github.com/microsoft/Megatron-DeepSpeed/blob/main/megatron/model/transformer.py