mirror of https://github.com/hpcaitech/ColossalAI
40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
import torch
|
|
from colossalai.tensor.op_wrapper import colo_op_impl
|
|
from colossalai.tensor import ColoTensor, dist_spec
|
|
|
|
|
|
@colo_op_impl(torch.nn.functional.layer_norm)
|
|
def colo_layernorm(types, args=(), kwargs=None, pg=None):
|
|
arg_num = len(args)
|
|
if arg_num > 0:
|
|
input_tensor = args[0]
|
|
if arg_num > 1:
|
|
normalized_shape = args[1]
|
|
if arg_num > 2:
|
|
weight = args[3]
|
|
if arg_num > 3:
|
|
bias = args[4]
|
|
if arg_num > 4:
|
|
eps = args[5]
|
|
|
|
if 'input' in kwargs:
|
|
input_tensor = kwargs['input']
|
|
if 'weight' in kwargs:
|
|
weight = kwargs['weight']
|
|
if 'bias' in kwargs:
|
|
bias = kwargs['bias']
|
|
if 'eps' in kwargs:
|
|
eps = kwargs['eps']
|
|
|
|
if isinstance(input_tensor, ColoTensor):
|
|
# TODO (ver217): check input dist spec
|
|
input_tensor.to_dist_spec(dist_spec.replicate(input_tensor.spec.get_process_group()))
|
|
input_tensor = input_tensor.torch_tensor()
|
|
if isinstance(weight, ColoTensor):
|
|
weight = weight.torch_tensor()
|
|
if isinstance(bias, ColoTensor):
|
|
bias = bias.torch_tensor()
|
|
|
|
return ColoTensor.init_from_torch_tensor(
|
|
torch.nn.functional.layer_norm(input_tensor, normalized_shape, weight, bias, eps))
|