github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 11:02:05 +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
tests/test_moe/test_kernel.py
View File

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