Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -667,8 +667,6 @@ class GemmoeSparseMoeBlock(nn.Module):
         topk_weight, topk_idx = torch.topk(routing_weights, self.top_k, dim=-1, sorted=False)
         topk_weight /= topk_weight.sum(dim=-1, keepdim=True)
 
-        hidden_states = hidden_states.repeat_interleave(self.top_k, dim=0)
-
         expert_outputs = []
         for i in range(self.num_experts):
             expert_input = hidden_states[topk_idx[:, i]]
@@ -676,9 +674,10 @@ class GemmoeSparseMoeBlock(nn.Module):
             expert_outputs.append(expert_output)
 
         expert_outputs = torch.stack(expert_outputs, dim=1)
-        expert_outputs = expert_outputs.view(batch_size * sequence_length, self.top_k, -1)
+        expert_outputs = expert_outputs.view(batch_size, sequence_length, self.top_k, -1)
+        topk_weight = topk_weight.view(batch_size, sequence_length, self.top_k, 1)
         
-        final_hidden_states = torch.einsum("bke,bkd->bed", topk_weight, expert_outputs)
+        final_hidden_states = (expert_outputs * topk_weight).sum(dim=2)
         final_hidden_states = final_hidden_states.view(batch_size, sequence_length, hidden_dim)
         
         return final_hidden_states.to(hidden_states.dtype), router_logits.to(hidden_states.dtype)