Spaces:
Runtime error
Runtime error
| """ | |
| RMSNorm - root-mean-square layer norm (Llama-style) | |
| No mean-subtraction and no bias like LayerNorm: just scale each vector to | |
| unit RMS, then multiply by a learned per-channel gain. Cheaper, and works | |
| as well for transformers. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| class AbyssRMSNorm(nn.Module): | |
| def __init__(self, dim, eps=1e-5): | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x): | |
| # normalize in float 32 for stability, then cast back to x's dtype | |
| dtype = x.dtype | |
| x = x.float() | |
| rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) | |
| return (self.weight * (x * rms)).to(dtype) |