fix: correct usage of minimax-m2 deepep moe forward (#13892)
Co-authored-by: Dash <dash@minimaxi.com> Co-authored-by: 赵晨阳 <zhaochen20@outlook.com>
This commit is contained in:
@@ -222,7 +222,7 @@ class MiniMaxM2MoE(nn.Module):
|
||||
if hidden_states.shape[0] > 0:
|
||||
# router_logits: (num_tokens, n_experts)
|
||||
router_logits, _ = self.gate(hidden_states.to(torch.float32))
|
||||
topk_weights, topk_idx, _ = self.topk(
|
||||
topk_output = self.topk(
|
||||
hidden_states,
|
||||
router_logits,
|
||||
num_token_non_padded=forward_batch.num_token_non_padded,
|
||||
@@ -231,14 +231,10 @@ class MiniMaxM2MoE(nn.Module):
|
||||
),
|
||||
)
|
||||
else:
|
||||
topk_weights, topk_idx, _ = self.topk.empty_topk_output(
|
||||
hidden_states.shape[0], self.top_k
|
||||
)
|
||||
topk_output = self.topk.empty_topk_output(device=hidden_states.device)
|
||||
final_hidden_states = self.experts(
|
||||
hidden_states=hidden_states,
|
||||
topk_idx=topk_idx,
|
||||
topk_weights=topk_weights,
|
||||
forward_batch=forward_batch,
|
||||
topk_output=topk_output,
|
||||
)
|
||||
|
||||
return final_hidden_states
|
||||
|
||||
Reference in New Issue
Block a user