mirror of https://github.com/hpcaitech/ColossalAI
35 lines
936 B
Python
35 lines
936 B
Python
import torch
|
|
from colossalai.tensor import ColoTensor, ColoParameter
|
|
from colossalai.utils import get_current_device
|
|
from torch.nn import Parameter
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def test_layernorm():
|
|
ln_op = torch.nn.LayerNorm(2, 3, device=get_current_device())
|
|
|
|
input_t = torch.randn(3, 2, device=get_current_device())
|
|
input_t_colo = ColoTensor.from_torch_tensor(input_t.clone().detach())
|
|
|
|
# prepare colossalai LN
|
|
weight = ColoTensor(Parameter(ln_op.weight.detach()))
|
|
bias = ColoTensor(Parameter(ln_op.bias.detach()))
|
|
|
|
output = ln_op(input_t)
|
|
output_colo = F.layer_norm(input_t_colo, ln_op.normalized_shape, weight, bias, ln_op.eps)
|
|
|
|
assert torch.allclose(output_colo, output)
|
|
|
|
torch.mean(output).backward()
|
|
torch.mean(output_colo).backward()
|
|
|
|
assert torch.allclose(ln_op.weight.grad, weight.grad)
|
|
|
|
|
|
def check_all():
|
|
test_layernorm()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
check_all()
|