from typing import List, Optional import torch.nn.functional as F from colossalai.tensor.op_wrapper import colo_op_impl from colossalai.tensor import ColoTensor, distspec, ColoTensorSpec, ReplicaSpec from ._utils import GeneralTensor, convert_to_colo_tensor @colo_op_impl(F.layer_norm) def colo_layernorm( input_tensor: GeneralTensor, normalized_shape: List[int], weight: Optional[GeneralTensor] = None, bias: Optional[GeneralTensor] = None, eps: float = 1e-5, ): assert isinstance(weight, ColoTensor) input_tensor = convert_to_colo_tensor(input_tensor, weight.get_process_group()) bias = convert_to_colo_tensor(bias, weight.get_process_group()) input_tensor = input_tensor.redistribute(ReplicaSpec()) output = F.layer_norm(input_tensor, normalized_shape, weight=weight, bias=bias, eps=eps) output = ColoTensor.from_torch_tensor( tensor=output, spec=ColoTensorSpec( pg=input_tensor.get_process_group(), dist_attr=input_tensor.dist_spec)) return output