github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 20:40:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

[moe] refactor mesh assignment

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

1
tests/test_moe/test_deepseek_layer.py
View File

@@ -47,7 +47,6 @@ def check_deepseek_moe_layer():
model,
ep_group=plugin.ep_group,
moe_dp_group=plugin.moe_dp_group,
moe_tp_group=plugin.moe_tp_group,
tp_group=plugin.tp_group,
)
ep_output = model(x)

1
tests/test_moe/test_mixtral_layer.py
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)