From 1c876435ef9f66919b3c415b2124c2146ed33fec Mon Sep 17 00:00:00 2001 From: basil-k-aji-dev <70605804+basil-k-aji-dev@users.noreply.github.com> Date: Thu, 17 Sep 2026 21:31:41 +0530 Subject: [PATCH] [fix] keep the MoE router trainable at num_experts_per_tok=1 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. --- model/model_minimind.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/model/model_minimind.py b/model/model_minimind.py index dc080e5..96dd8fb 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -158,7 +158,14 @@ class MOEFeedForward(nn.Module): x_flat = x.view(-1, hidden_dim) scores = F.softmax(self.gate(x_flat), dim=-1) 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) + # Normalising a single top-k weight by its own sum yields 1.0 for every + # token, so the gate drops out of the graph: topk_idx is discrete and the + # weight is now a constant, leaving the LM loss no path to the router. It + # then trains on aux_loss alone, whose optimum is a uniform distribution, + # and the experts never specialise. Only normalise when there is more than + # one weight to normalise. + 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) y = torch.zeros_like(x_flat) for i, expert in enumerate(self.experts): mask = (topk_idx == i)