github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-07 20:10:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into prefetch

Browse Source
This commit is contained in:
hxwang
2024-05-23 15:49:33 +00:00
parent 13c06d36a3 4647ec28c8
commit 15d21a077a
238 changed files with 20778 additions and 9787 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
examples/language/openmoe/model/modeling_openmoe.py
View File

@@ -35,7 +35,20 @@ from transformers.utils import (
replace_return_docstrings,
)
from colossalai.kernel.extensions.flash_attention import HAS_FLASH_ATTN
try:
# TODO: remove this after updating openmoe example
# NOTE(yuanheng-zhao): This is a temporary fix for the issue that
# the flash_attention module is not imported correctly for different CI tests.
# We replace the import path `colossalai.kernel.extensions.flash_attention`
# because in the current example test, colossalai version <= 0.3.6 is installed,
# where `colossalai.kernel.extensions.flash_attention` is still valid;
# however in unit test `test_moe_checkpoint`, the lastest version of colossalai is installed,
# where extension has been refactored and the path is not valid.
import flash_attention # noqa
HAS_FLASH_ATTN = True
except:
HAS_FLASH_ATTN = False
from colossalai.kernel.triton.llama_act_combine_kernel import HAS_TRITON
from colossalai.moe.layers import SparseMLP
from colossalai.moe.manager import MOE_MANAGER