github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-24 14:33:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Hotfix] hotfix normalization (#6163)

Browse Source 
* [fix] hotfix normalization

* [hotfix] force doc ci test

* [hotfix] fallback doc
This commit is contained in:
duanjunwen 2024-12-23 16:29:48 +08:00 committed by GitHub
parent 130229fdcb
commit fa9d0318e4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 15 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/shardformer/layer/normalization.py
View File

@ -76,6 +76,7 @@ if SUPPORT_NPU:
FusedRMSNormWithHook = NPUFusedRMSNormWithHook FusedRMSNormWithHook = NPUFusedRMSNormWithHook
else: else:
try:
from apex.normalization import FusedRMSNorm as ApexFusedRMSNorm from apex.normalization import FusedRMSNorm as ApexFusedRMSNorm
class CUDAFusedRMSNormWithHook(ApexFusedRMSNorm): class CUDAFusedRMSNormWithHook(ApexFusedRMSNorm):
@ -88,6 +89,11 @@ else:
return output return output
FusedRMSNormWithHook = CUDAFusedRMSNormWithHook FusedRMSNormWithHook = CUDAFusedRMSNormWithHook
except ImportError:
warnings.warn(
"Please install apex from source (https://github.com/NVIDIA/apex) to use the fused RMSNorm kernel"
)
FAST_LAYERNORM_SUPPORTED_SIZE = [ FAST_LAYERNORM_SUPPORTED_SIZE = [
1024, 1024,