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

Optimized MoE layer and fixed some bugs;

Browse Source 
Decreased moe tests;

Added FFNExperts and ViTMoE model
This commit is contained in:
1SAA
2022-02-18 20:42:31 +08:00
committed by Frank Lee
parent 3dba070580
commit 219df6e685
15 changed files with 1552 additions and 203 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
setup.py
View File

@@ -162,6 +162,10 @@ if build_cuda_ext:
['scaled_masked_softmax.cpp', 'scaled_masked_softmax_cuda.cu'],
extra_cuda_flags + cc_flag))
ext_modules.append(cuda_ext_helper('colossal_moe_cuda',
['moe_cuda.cpp', 'moe_cuda_kernel.cu'],
extra_cuda_flags + cc_flag))
extra_cuda_flags = ['-maxrregcount=50']
ext_modules.append(cuda_ext_helper('colossal_layer_norm_cuda',