前置知识
MOE(MixerOfExpert)
moe的主要原理是替换attention层后的MLP层, 通过将不同类型的token按照门控单元计算出的概率分配给最大概率处理的专家网络处理, 对比单一MLP更适合处理复杂多样化的数据集. 主要思想和集成学习感觉很像, 而且扩展性(遇到新的目标任务可以新增专家网络)和可解释性(每个专家分开调整)都比较强. MOE前向步骤(以最简单的top2 Expert为例):
- 门控网络, 输入是attention的输出, dim为(batch_size, tokens, emb_size), 输出dim为(batch_size, tokens, experts_num)
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计算进入各个专家网络的概率
- 稀疏化Experts
这部分每个expert的网络结构都可以根据场景设计的不一样, 因为在fp/bp计算的时候, 每个token都是只进入了topk的网络进行计算, 剩余的网络没有计算. 大部分的参数都没参与更新, 所以也被称为稀疏化的dense. 因为这个特性也给moe网络的并行化改造提供了基础.
expert负载不均匀问题
因为expert的概率纯粹是训练出来的参数决定的, 没法用LALB类似的负载均衡策略强制使每个expert接收到的token是均匀的, 极有可能出现某几个expert接收到了很多, 其他的基本没啥token的问题..主要有以下这么几个解决办法:
TokenBuffer: 给每个expert设置固定容量, 容量设置公式如下, 当这个expert收满token后就不再接受token
\[ 标签：expert,并行,网络,token,溢出,LLM,1st,logits,MoE From： https://www.cnblogs.com/sunstrikes/p/18310517