From fa36e4861d08088992db3eda2636cf991ab323da Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Fri, 18 Sep 2026 16:34:37 +0800 Subject: [PATCH] [fix] top-1 router gradient --- model/model_minimind.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/model/model_minimind.py b/model/model_minimind.py index dc080e5..0f0c3ef 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -158,7 +158,12 @@ 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) + if self.config.norm_topk_prob: + if self.config.num_experts_per_tok > 1: + topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20) + else: + top1 = torch.topk(F.softmax(self.gate(x_flat.detach()), dim=-1), k=1, dim=-1, sorted=False)[0] + topk_weight = top1 - top1.detach() + 1.0 # k=1: 1.0 forward on purpose, gradient via straight-through y = torch.zeros_like(x_flat) for i, expert in enumerate(self.experts): mask = (topk_idx == i)