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] clean legacy code

Browse Source
This commit is contained in:
hxwang
2024-07-16 09:08:31 +00:00
committed by Hongxin Liu
parent 74eccac0db
commit 3e2b6132b7
39 changed files with 163 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_moe/test_kernel.py
View File

@@ -4,9 +4,7 @@ import pytest
import torch
from colossalai.accelerator import get_accelerator
# from colossalai.moe import SparseMLP
from colossalai.moe._operation import MoeCombine, MoeDispatch, moe_cumsum
from colossalai.moe.operators import MoeCombine, MoeDispatch, moe_cumsum
NUM_EXPERTS = 4
BATCH_SIZE = 4