deepseekmirror/DeepSeek-V3
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepSeek-V3.git synced 2025-02-22 05:38:59 -05:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

fix scores mask

Browse Source
This commit is contained in:
Xingkai Yu 2025-02-14 20:26:45 +08:00 committed by GitHub
parent 2f7b80eece
commit 1398800ebf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
inference/model.py
View File

@ -585,8 +585,8 @@ class Gate(nn.Module):
else: else:
group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1) group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1)
indices = group_scores.topk(self.topk_groups, dim=-1)[1] indices = group_scores.topk(self.topk_groups, dim=-1)[1]
mask = torch.zeros_like(scores[..., 0]).scatter_(1, indices, True) mask = scores.new_ones(x.size(0), self.n_groups, dtype=bool).scatter_(1, indices, False)
scores = (scores * mask.unsqueeze(-1)).flatten(1) scores = scores.masked_fill_(mask.unsqueeze(-1), float("-inf")).flatten(1)
indices = torch.topk(scores, self.topk, dim=-1)[1] indices = torch.topk(scores, self.topk, dim=-1)[1]
weights = original_scores.gather(1, indices) weights = original_scores.gather(1, indices)
if self.score_func == "sigmoid": if self.score_func == "sigmoid":