github
/
minimind-v
mirror of https://github.com/jingyaogong/minimind-v.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

moe note

pull/45/merge
jingyaogong 2025-04-09 17:37:51 +08:00
parent 139cc98bb9
commit 5f130d142d
1 changed files with 0 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
model/model.py
View File

@ -221,7 +221,6 @@ class MOEFeedForward(nn.Module):
x = x.view(-1, x.shape[-1]) x = x.view(-1, x.shape[-1])
flat_topk_idx = topk_idx.view(-1) flat_topk_idx = topk_idx.view(-1)
if self.training: if self.training:
# 训练模式下，重复输入数据
x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0) x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
y = torch.empty_like(x, dtype=torch.float16) y = torch.empty_like(x, dtype=torch.float16)
for i, expert in enumerate(self.experts): for i, expert in enumerate(self.experts):
@ -229,7 +228,6 @@ class MOEFeedForward(nn.Module):
y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1) y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
y = y.view(*orig_shape) y = y.view(*orig_shape)
else: else:
# 推理模式下，只选择最优专家
y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape) y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
if self.config.n_shared_experts is not None: if self.config.n_shared_experts is not None:
y = y + self.shared_experts(identity) y = y + self.shared_experts(identity)