github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[hotfix] fix OpenMOE example import path (#5697)

pull/5689/head
Yuanheng Zhao 2024-05-08 15:48:47 +08:00 committed by GitHub
parent 55cc7f3df7
commit 12e7c28d5e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 14 additions and 1 deletions
Show Stats 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