kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-14 06:27:45 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[Qwen3] chore: fix bug of fused_moe on tp > 1 (#4093)

Browse Source 
* fix bug of fused_moe on tp > 1

Signed-off-by: bhsueh <11360707+byshiue@users.noreply.github.com>

* refine codes

Signed-off-by: bhsueh <11360707+byshiue@users.noreply.github.com>

---------

Signed-off-by: bhsueh <11360707+byshiue@users.noreply.github.com>
This commit is contained in:
bhsueh_NV 2025-05-07 11:06:37 +08:00 committed by GitHub
parent c28b90984f
commit f670a036df
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorrt_llm/_torch/modules/fused_moe.py
View File

@ -821,10 +821,7 @@ class FusedMoE(nn.Module):
final_hidden_states = final_hidden_states[0]
if not self.enable_alltoall:
if self.reduce_results and self.parallel_size > 1:
return self.all_reduce(final_hidden_states)
else:
return final_hidden_states
return final_hidden_states
else:
return self.alltoall_combine(final_hidden_states, alltoall_info,
token_count)