[Bug] num_experts_per_tok=1 时 MoE 路由器收不到梯度（topk 权重归一化后退化为常数）

环境

• commit：a3c7b01
• 配置：MiniMindConfig(hidden_size=768, num_hidden_layers=8, use_moe=True) —— 即默认的 num_experts=4、num_experts_per_tok=1、norm_topk_prob=True、router_aux_loss_coef=5e-4
• 数据：pretrain_t2t_mini.jsonl（1.24GB / 0.33B token）
• 超参：epochs=2, bs=32, accumulation_steps=8, lr=5e-4, max_seq_len=340

1. 现象

用完全相同的数据和超参分别训练 dense 与 MoE：

模型	参数量	ppl（60 条留出样本）	最终 logits_loss
dense	63.9M	5.25	1.6442
MoE	198.4M	5.45	1.7053
官方 pretrain_768_moe.pth	198.4M	3.98	—

MoE 参数量多 3 倍，效果反而更差。检查路由分布：

检查项	我训练的 MoE	官方 MoE
专家负载分布	[0.241, 0.263, 0.247, 0.249]	[0.274, 0.168, 0.396, 0.162]
偏离均匀程度	0.026（几乎完美均匀）	0.340（明显分化）
aux_loss	0.004005（理论下界 0.004000）	0.011545（高出 188%）

专家完全没有分化，aux_loss 全程贴着理论最小值。

2. 根因

model/model_minimind.py:160-161：

topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False)
if self.config.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)

路由器影响输出的路径只有两条：

1. 选哪个专家 —— 由 topk_idx 决定，是 torch.topk 的离散索引，不可导；
2. 选中后的权重 —— 由 topk_weight 决定，这是唯一的可导通路。

而 num_experts_per_tok=1 时 topk_weight 形状为 [T, 1]， topk_weight / topk_weight.sum(-1, keepdim=True) 恒等于 1.0 —— 与 gate 无关的常数。

⇒ 两条路都断，LM loss 对路由器的梯度恒为 0，路由器只能被 aux_loss 训练， 而 aux_loss 的目标是"均匀"，于是路由被锁死在均匀分布上。

3. 最小复现（10 行，零依赖）

import torch, torch.nn.functional as F
torch.manual_seed(0)

B, H, E = 32, 768, 4
x = torch.randn(B, H)
gate_w = torch.randn(E, H, requires_grad=True)
scores = F.softmax(x @ gate_w.T, dim=-1)                 # [B, E]

for k, norm in [(1, True), (2, True), (1, False)]:
    w, idx = torch.topk(scores, k=k, dim=-1)             # [B, k]
    if norm:
        w = w / (w.sum(dim=-1, keepdim=True) + 1e-20)    # 仓库里的归一化
    t = torch.randn(B, k)                                # 固定下游目标，避免输出退化成常数
    g = torch.autograd.grad((w * t).sum(), gate_w, retain_graph=True)[0]
    print(f'k={k} norm={norm}  →  梯度范数 {g.norm():.3e}')

输出：

k=1 norm=True   →  梯度范数 1.249e-07      ← ≈ 0
k=2 norm=True   →  梯度范数 1.223e+01
k=1 norm=False  →  梯度范数 1.292e+01

4. 完整模型上的实测

torch.autograd.grad(r.loss, gate_params) 与 torch.autograd.grad(r.aux_loss, gate_params)， 模型处于 train() 模式：

配置	LM 梯度	aux 梯度	谁在训练路由器
k=1, norm_topk_prob=True（当前默认）	1.44e-08	4.79e-03	aux 独占（差 33 万倍）
k=1, norm=True, coef=0（关掉 aux）	1.52e-08	0	LM 梯度纹丝不动
k=1, norm_topk_prob=False	2.09e-01	1.27e-02	LM 主导（反超 16 倍）
k=2, norm=True	3.91e-01	1.01e-02	LM 主导（39 倍）

补充两点：

• aux_loss 只占总损失的 0.049%（随机初始化）/ 约 0.24%（训练后期）。 所以问题不是"aux 太强"，而是 LM 梯度路径断了：把 coef 设为 0，LM 梯度仍是 1.52e-08， 路由器会被完全冻住，依然不分化。
• 这也解释了为什么没出现常见的"专家坍缩"，反而出现"过度均匀"—— aux_loss 是路由器唯一的训练信号。

4.5 修复验证：A/B 对照（同一份数据、同样步数，只改 norm_topk_prob）

pretrain_100k.jsonl（10 万条切片），各训练 3,125 步（约 8 分钟）。 两组唯一的差别是 norm_topk_prob：A 用默认的 True， B 把 model/model_minimind.py:44 的默认值临时改成 False（其余参数、数据顺序完全一致）。

	专家负载分布	偏离均匀	aux_loss	ppl
A：默认（norm=True）	[0.252, 0.229, 0.269, 0.250]	0.042	0.004019	33.75
B：修复（norm=False）	[0.188, 0.233, 0.257, 0.322]	0.158	0.004258	24.70
理论下界	[0.25, 0.25, 0.25, 0.25]	0	0.004000	—

同一检查点上的梯度：

	LM 梯度	aux 梯度	谁主导
A	1.887e-08	1.223e-03	aux 主导
B	1.985e-01	1.926e-03	LM 主导（105 倍）

修掉归一化后：LM 梯度放大 1000 万倍 → 路由偏离度 ×3.8 → ppl 降低 27%。

还有一组很有说服力的对照：

模型	训练步数	偏离均匀度
A（默认配置）	3,125	0.042
我之前的完整 MoE（默认配置）	79,390	0.026 ← 训得越久越均匀
B（修复配置）	3,125	0.158
官方 MoE 权重	未知	0.340

默认配置下训练 79,390 步（25 倍于 A）后的偏离度反而比 A 更低， 说明"均匀化"就是这条路径的收敛方向——不是训练不足，而是被锁死了。

5. 沿革：这个守卫原本存在过（供参考）

git log -S "norm_topk_prob" -- model/model_minimind.py 显示 MoE 模块在 101d7df（2026-03-24，[update] minimind-3）这次重构中改写成 Qwen3-MoE 风格。 重写前的实现（a62faf3，2025-04-26）里有一处守卫：

# a62faf3（2025-04-26）
topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)

if self.top_k > 1 and self.norm_topk_prob:                      # ← 守卫
    denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
    topk_weight = topk_weight / denominator

值得注意的是：当时的默认值是 num_experts_per_tok = 2，所以这个守卫在默认配置下 永远不会被触发（top_k > 1 恒成立）——它只对"手动把 top_k 设成 1"的用户起作用。 换句话说，它在旧版本里存在感很低，重写时没有保留下来是很自然的。

但重写同时把默认值改成了 1，两个变化叠加才产生了这个缺陷：

配置项	重写前（a62faf3）	当前（a3c7b01）
num_experts_per_tok	2	1
归一化守卫	top_k > 1 and ...	无
归一化后的权重	依赖输入（可导）	恒为常数 1.0（不可导）
退化情形	需手动设置 top_k=1，且有守卫保护	默认即触发

这一点对修复的意义：恢复守卫等价于恢复项目自己原本就有的设计， 且 k≥2（含 Qwen3-MoE 风格的 top-8）时的前向与梯度行为完全不变——属于低风险改动。

参考：HF 的 qwen2_moe / qwen3_moe / mixtral 里是没有这个守卫的， 但它们的默认 num_experts_per_tok 分别是 8 / 8 / 2，都不会走到 k=1 这个分支。

6. 建议的修复

恢复原来的守卫即可（一行）：

if self.config.norm_topk_prob and self.config.num_experts_per_tok > 1:
    topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)

这一行不改变 k≥2（含 Qwen3-MoE 风格的 top-8）时的任何前向与梯度行为， 也不影响已发布权重的加载；只影响 k=1，而 k=1 目前的语义是"权重恒为 1"。

其他可选方案：

• 默认 norm_topk_prob=False（model_minimind.py:44），对齐 Switch Transformer 的 top-1 语义（该论文用 gate 概率本身作为专家权重，不做归一化）；
• 默认 num_experts_per_tok=2（对齐 Mixtral），代价是每步计算量翻倍。

7. 影响

默认配置下，README 中 minimind-3-moe 的设计意图（"以更低激活参数获得更高容量"）无法实现： 路由器学不到任何"该把 token 给谁"的信息，专家不分化，多出来的 ~134M 专家参数 处于"各自只训练了 1/4 数据"的低效状态。实测 MoE（ppl 5.45）不如同数据同轮数的 dense（ppl 5.25）。

8. 一个我还没搞清楚的点

官方发布的 pretrain_768_moe.pth 里专家明显分化了（负载偏离 0.340、aux 高出下界 188%）， 但在当前代码的默认配置下这不该发生。想请问这个权重是用什么配置/代码版本训练的？ （我实测过 norm_topk_prob=True/False 两种评估方式，True 都更好，所以不确定差异来源。）