Skip to content

Commit aeec295

Browse files
authored
[bugfix] fix deepseek bug when tp_size == 1 (#1755)
### What this PR does / why we need it? fix deepseek bug when tp_size == 1 ### Does this PR introduce _any_ user-facing change? ### How was this patch tested? Signed-off-by: zzzzwwjj <1183291235@qq.com>
1 parent 9d6808e commit aeec295

File tree

1 file changed

+7
-5
lines changed

1 file changed

+7
-5
lines changed

vllm_ascend/ops/fused_moe.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1242,13 +1242,15 @@ def forward(self,
12421242
if isinstance(e_hidden_states, tuple):
12431243
e_hidden_states, shared_hidden_states = e_hidden_states
12441244

1245-
if tp_size > 1 and fused_moe_state != FusedMoEState.AllGather:
1246-
dist.all_gather(list(chunk_hidden_states), e_hidden_states,
1247-
self.tp_group)
1248-
final_hidden_states = torch.cat(chunk_hidden_states, dim=0)
1245+
if fused_moe_state != FusedMoEState.AllGather:
1246+
if tp_size > 1:
1247+
dist.all_gather(list(chunk_hidden_states), e_hidden_states,
1248+
self.tp_group)
1249+
final_hidden_states = torch.cat(chunk_hidden_states, dim=0)
1250+
else:
1251+
final_hidden_states = e_hidden_states
12491252
if num_tokens < forward_context.padded_num_tokens:
12501253
final_hidden_states = final_hidden_states[:num_tokens]
1251-
dispose_tensor(e_hidden_states)
12521254
elif self.dp_size > 1 and fused_moe_state == FusedMoEState.AllGather:
12531255
final_hidden_states = dist._functional_collectives.reduce_scatter_tensor(
12541256
e_hidden_states,

0 commit comments

Comments
 (0)