| import torch | |
| import torch.nn as nn | |
| def fused_rms_norm(x: torch.Tensor, weight: nn.Parameter, eps: float): | |
| """ | |
| Applies a fused Root Mean Square (RMS) normalization to the input tensor. | |
| Args: | |
| x (torch.Tensor): The input tensor to be normalized. Expected to have | |
| at least one dimension. | |
| weight (nn.Parameter): A learnable parameter used to scale the normalized | |
| tensor. Its shape must be broadcastable to the shape of `x`. | |
| eps (float): A small constant added to the denominator for numerical | |
| stability during normalization. | |
| Returns: | |
| torch.Tensor: The normalized and scaled tensor with the same shape as `x`. | |
| """ | |
| x = x.float() | |
| return (x * torch.rsqrt((x * x).mean(-1, keepdim=True).add_(eps))) * weight | |
| class LayerNorm(nn.LayerNorm): | |
| def forward(self, input: torch.Tensor): | |
| """ | |
| Wrapper to ensure that the input tensor is cast to float before normalization. | |
| """ | |
| y = super().forward(input.float()) | |
| return y.type_as(input) | |
| class RMSNorm(nn.Module): | |
| def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine: bool = True): | |
| """ | |
| Initializes the normalization layer. | |
| Args: | |
| dim (int): The number of features in the input tensor. | |
| eps (float, optional): A small value added to the denominator for numerical stability. Defaults to 1e-5. | |
| elementwise_affine (bool, optional): If True, this layer will have learnable per-element affine parameters. Defaults to True. | |
| """ | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim), requires_grad=elementwise_affine) | |
| def forward(self, x): | |
| return fused_rms_norm(x, weight=self.weight, eps=self.eps).type_as(x) | |