[moe] refactor mesh assignment

This commit is contained in:
hxwang
2024-07-25 06:19:54 +00:00
committed by Hongxin Liu
parent 034020bd04
commit cb01c0d5ce
10 changed files with 277 additions and 170 deletions

View File

@@ -42,7 +42,6 @@ def check_mixtral_moe_layer():
ep_group=plugin.ep_group,
tp_group=plugin.tp_group,
moe_dp_group=plugin.moe_dp_group,
moe_tp_group=plugin.moe_tp_group,
)
ep_output, ep_logits = model(x)
assert_close(orig_logits, ep_logits)