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

[Hotfix] Avoid fused RMSnorm import error without apex (#5985) 

Co-authored-by: Edenzzzz <wtan45@wisc.edu>
pull/5991/head
Edenzzzz 2024-08-09 18:17:09 +08:00 committed by GitHub
parent ad3fa4f49c
commit b4d2377d4c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 13 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -42,7 +42,7 @@ try:
return output
except ImportError:
warnings.warn("Please install apex from source (https://github.com/NVIDIA/apex) to use the fused layernorm kernel")
warnings.warn("Please install apex from source (https://github.com/NVIDIA/apex) to use the fused RMSNorm kernel")
FAST_LAYERNORM_SUPPORTED_SIZE = [
1024,
@ -270,12 +270,6 @@ class FusedRMSNorm(BaseLayerNorm):
Returns:
nn.Module: FusedRMSNorm module.
"""
try:
pass
except ImportError:
raise ImportError(
"Please install apex from source (https://github.com/NVIDIA/apex) to use the fused RMS normalization kernel"
)
LazyInitContext.materialize(module)
@ -284,11 +278,18 @@ class FusedRMSNorm(BaseLayerNorm):
eps = module.variance_epsilon if hasattr(module, "variance_epsilon") else module.eps
elementwise_affine = getattr(module, "elementwise_affine", True)
try:
rmsnorm = FusedRMSNormWithHook(
normalized_shape=normalized_shape,
eps=eps,
elementwise_affine=elementwise_affine,
)
except ImportError:
warnings.warn(
"Module replacement failed.\
Please install apex from source (https://github.com/NVIDIA/apex) to use the fused RMS normalization kernel"
)
return module
rmsnorm.weight = module.weight