当前位置：首页 > news >正文

06. MoE Router代码笔记

news 2026/6/25 13:29:00

背景

TopKRouter是混合专家模型（Mixture of Experts, MoE）中的门控路由器。它接收一个批次中所有 Token 的隐藏状态，为每个 Token 选出最合适的K 个专家，并计算对应的权重。路由包含三个关键步骤：

用线性层产生每个专家的 logits；
对 logits 做 softmax 得到概率分布；
选取概率最大的 K 个专家，并重归一化这些权重。

三个 TODO 分别对应第 2、3 步以及重归一化。

TODO 1：对全量 Logits 进行 Softmax 获取所有专家的概率分布

routing_probs=F.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_experts=torch.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-normalize）

routing_weights=routing_weights/routing_weights.sum(dim=-1,keepdim=True)

作用：让选出的 K 个专家的权重按比例放大，使其总和为 1，变成一个合法的概率分布（仅在这 K 个专家上）。

为何需要重归一化：

Top-K 截取后，被选中的 K 个概率之和通常小于 1（除非这 K 个恰好就是概率最大的全部且剩余概率为 0）。
后续在SparseMoEBlock中，每个选中的专家输出将乘以routing_weights再累加到最终结果。如果直接使用截取后的概率，累加的总贡献会变小（相当于多了一个衰减因子），这不符合“这 K 个专家完全负责该 Token”的预期，会改变梯度的尺度。
重归一化保证了“被激活的这 K 个专家共同解释该 Token 的全部输出”，使得输出幅值稳定。
计算方式：routing_weights / routing_weights.sum(dim=-1, keepdim=True)，对每个 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正是驱动这一过程的核心数据结构。