github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 02:51:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

[test] add mixtral transformer test

Browse Source
This commit is contained in:
hxwang
2024-07-02 09:08:41 +00:00
committed by Hongxin Liu
parent f9b6fcf81f
commit 0b76b57cd6
6 changed files with 281 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_moe/moe_utils.py
View File

@@ -10,8 +10,6 @@ from colossalai.legacy.engine.gradient_handler.utils import bucket_allreduce
from colossalai.legacy.registry import GRADIENT_HANDLER
from colossalai.moe.manager import MOE_MANAGER
from colossalai.moe.utils import get_moe_epsize_param_dict
# from colossalai.shardformer.layer.moe import SparseMLP
from colossalai.tensor.moe_tensor.api import get_ep_group, get_ep_size, set_moe_tensor_ep_group