[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

@@ -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)