06. MoE Router代码笔记

背景TopKRouter是混合专家模型Mixture of Experts, MoE中的门控路由器。它接收一个批次中所有 Token 的隐藏状态为每个 Token 选出最合适的K 个专家并计算对应的权重。路由包含三个关键步骤用线性层产生每个专家的 logits对 logits 做 softmax 得到概率分布选取概率最大的 K 个专家并重归一化这些权重。三个 TODO 分别对应第 2、3 步以及重归一化。TODO 1对全量 Logits 进行 Softmax 获取所有专家的概率分布routing_probsF.softmax(router_logits.float(),dim-1)作用将门控层输出的原始分数router_logits形状[total_tokens, num_experts]转换为概率分布使每个 Token 在所有专家上的权重和为 1。细节与必要性.float()强制将router_logits转换为float32再进行 softmax。在混合精度训练如bfloat16或float16中logits 可能数值范围大直接 softmax 容易溢出或损失精度。提升到 FP32 能保证数值稳定性。F.softmax(..., dim-1)在最后一个维度专家维度上计算 softmax。结果routing_probs的每一行是概率向量表示该 Token 与各专家的匹配程度。这一步是路由的基础为后续 Top-K 选择提供标准的概率分布。TODO 2从概率分布中截取 Top-K 最大的概率及其索引routing_weights,selected_expertstorch.topk(routing_probs,self.top_k,dim-1)作用对于每个 Token从routing_probs中选出概率最大的self.top_k个专家得到它们的权重和专家索引。细节torch.topk(routing_probs, self.top_k, dim-1)返回两个张量routing_weights形状[total_tokens, top_k]每个 Token 选出的 K 个最高的概率值。selected_experts形状[total_tokens, top_k]这些概率对应的专家索引long类型。这一步体现了 MoE 的“稀疏性”每个 Token 只激活 K 个专家其余专家不参与计算。大大节省了计算量。选取后的routing_weights不再是一个完整的概率分布因为舍弃了其余专家需要下一步的处理。TODO 3对截取后的 routing_weights 进行重归一化Re-normalizerouting_weightsrouting_weights/routing_weights.sum(dim-1,keepdimTrue)作用让选出的 K 个专家的权重按比例放大使其总和为 1变成一个合法的概率分布仅在这 K 个专家上。为何需要重归一化Top-K 截取后被选中的 K 个概率之和通常小于 1除非这 K 个恰好就是概率最大的全部且剩余概率为 0。后续在SparseMoEBlock中每个选中的专家输出将乘以routing_weights再累加到最终结果。如果直接使用截取后的概率累加的总贡献会变小相当于多了一个衰减因子这不符合“这 K 个专家完全负责该 Token”的预期会改变梯度的尺度。重归一化保证了“被激活的这 K 个专家共同解释该 Token 的全部输出”使得输出幅值稳定。计算方式routing_weights / routing_weights.sum(dim-1, keepdimTrue)对每个 Token 在top_k维度上除以各自的和。最后代码通过routing_weights routing_weights.to(hidden_states.dtype)将权重转回输入的精度如bfloat16以便与后续计算保持一致。额外说明SparseMoEBlock 如何使用这些结果虽然这不是 TODO但理解上下文有助于理解路由的目的在SparseMoEBlock.forward中对每个专家使用torch.where(selected_experts expert_idx)找到所有选中该专家的 Token分别计算专家输出然后乘以对应的重归一化权重routing_weights[token_idx, kth_expert]累加到最终输出中。这体现了 MoE 的“稀疏加权汇聚”每个 Token 只被自己选中的 K 个专家处理最后按权重合并。路由器返回的routing_weights和selected_experts正是驱动这一过程的核心数据结构。总结TODO 1用 FP32 执行 softmax得到每个 Token 对所有专家的完整概率分布。TODO 2通过torch.topk从中选出 K 个最大的概率和对应的专家索引实现稀疏选择。TODO 3对这 K 个权重进行重归一化使它们和为 1保证激活专家的输出幅值稳定。这三步共同构成了 MoE 路由的标准流程平衡了计算效率和模型容量。