diff --git a/model/model.py b/model/model.py index 8e873cb..848cf8a 100644 --- a/model/model.py +++ b/model/model.py @@ -221,7 +221,6 @@ class MOEFeedForward(nn.Module): x = x.view(-1, x.shape[-1]) flat_topk_idx = topk_idx.view(-1) if self.training: - # 训练模式下,重复输入数据 x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0) y = torch.empty_like(x, dtype=torch.float16) 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(*orig_shape) else: - # 推理模式下,只选择最优专家 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: y = y + self.shared_experts(identity)