dit/norm.py

import torch

def rms_norm(x, weight=None, eps=1e-05):
    output = x / torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
    return output * weight if weight is not None else output

class RMSNorm(torch.nn.Module):

    def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True, dtype=None, device=None):
        super().__init__()
        self.eps = eps
        if elementwise_affine:
            self.weight = torch.nn.Parameter(torch.ones(normalized_shape, dtype=dtype, device=device))
        else:
            self.register_parameter('weight', None)

    def forward(self, x):
        return rms_norm(x.float(), self.weight, self.eps).to(dtype=x.dtype)