Normalising a single top-k weight by its own sum gives exactly 1.0 for every token, so the gate leaves the graph entirely: topk_idx is a discrete index and the weight is now a constant, which leaves the LM loss no differentiable path to the router. It then trains on aux_loss alone, whose optimum is a uniform distribution, and the experts never specialise. Measured on MOEFeedForward with router_aux_loss_coef=0, so only the LM path can contribute a gradient: num_experts_per_tok=1, norm_topk_prob=True 7.02e-10 <- current default num_experts_per_tok=1, norm_topk_prob=False 1.02e-02 num_experts_per_tok=2, norm_topk_prob=True 1.67e-02 Only normalise when there is more than one weight to normalise. Note this changes the forward pass for existing num_experts_per_tok=1 MoE checkpoints: the expert output was scaled by a constant 1.0 and is now scaled by the gate probability. Closes #858. |
||
|---|---|---|
| .. | ||
| __init__.py | ||
| model_lora.py | ||
| model_minimind.py | ||
| tokenizer.json | ||
| tokenizer_config.json | ||