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.
jingyaogong
basil-k-aji-dev
jingyaogong
yuyu5333