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

[moe] full test for deepseek and mixtral (pp + sp to fix)

Browse Source
This commit is contained in:
hxwang
2024-07-19 06:11:11 +00:00
committed by Hongxin Liu
parent 7077d38d5a
commit 803878b2fd
17 changed files with 430 additions and 517 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/legacy/moe/layer/layers.py
View File

@@ -9,7 +9,7 @@ import torch.nn.functional as F
from colossalai.legacy.moe.load_balance import LoadBalancer
from colossalai.legacy.moe.utils import create_ep_hierarchical_group, get_noise_generator
from colossalai.moe.operators import AllGather, AllToAll, HierarchicalAllToAll, MoeCombine, MoeDispatch, ReduceScatter
from colossalai.moe._operation import AllGather, AllToAll, HierarchicalAllToAll, MoeCombine, MoeDispatch, ReduceScatter
from colossalai.shardformer.layer.moe import MLPExperts
from colossalai.tensor.moe_tensor.api import get_dp_group, get_ep_group, get_ep_group_ranks, get_ep_size