mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
30 lines
1.1 KiB
30 lines
1.1 KiB
import torch
|
|
|
|
from .tracer import register_leaf_module, register_leaf_module_impl
|
|
|
|
try:
|
|
import apex
|
|
register_leaf_module(apex.normalization.FusedLayerNorm)
|
|
register_leaf_module(apex.normalization.FusedRMSNorm)
|
|
register_leaf_module(apex.normalization.MixedFusedLayerNorm)
|
|
register_leaf_module(apex.normalization.MixedFusedRMSNorm)
|
|
|
|
@register_leaf_module_impl(apex.normalization.FusedLayerNorm)
|
|
@register_leaf_module_impl(apex.normalization.FusedRMSNorm)
|
|
@register_leaf_module_impl(apex.normalization.MixedFusedLayerNorm)
|
|
@register_leaf_module_impl(apex.normalization.MixedFusedRMSNorm)
|
|
def torch_nn_normalize(self, input: torch.Tensor):
|
|
# check shape
|
|
if isinstance(self, torch.nn.BatchNorm1d):
|
|
assert input.dim() in [2, 3]
|
|
elif isinstance(self, torch.nn.BatchNorm2d):
|
|
assert input.dim() == 4
|
|
elif isinstance(self, torch.nn.BatchNorm3d):
|
|
assert input.dim() == 5
|
|
|
|
# normalization maintain the same shape as the input
|
|
return input.clone()
|
|
|
|
except (ImportError, AttributeError):
|
|
pass
|