Mixtral Moe代码解读

原创

已于 2024-08-06 18:12:37 修改 · 1.2k 阅读

10 ·

CC 4.0 BY-SA版权

文章标签：

#MOE

于 2024-01-11 02:51:13 首次发布

本文深入解析了Mixtral Moe模型的工作原理与代码实现，重点介绍了如何通过稀疏专家网络分配不同专家权重，确保每个专家都能在训练过程中发挥作用，即使部分专家在某些批次未被选中。

一直对稀疏专家网络好奇，有些专家没被选中，那么梯度是否为0，这一轮被选中有梯度，下一轮没被选中无梯度，模型可以训练收敛吗？

由于每个token都会选择topk个专家，所以在每一轮epoch中，所有专家都参与了前向传播，所以梯度都能得到更新
即使真有专家一直没被选中，那么其梯度保持不变，没有参与更新而已

self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)

# 获取到每个token的mlp层输入特征 
batch_size, sequence_length, hidden_dim = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_dim)

# 得到每个专家的打分，维度是batch * sequence, num_experts，取topk个专家
router_logits = self.gate(hidden_states)
routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)

# 取到topk个专家的打分，需要计算在归一化一下，用于对后面的expert计算出来的结果进行加权
routing_weights /= routing_weights.sum(dim=-