| import torch | |
| import torch.nn as nn | |
| class RMSNorm(nn.Module): | |
| """RMSNorm (Section 4.7): cheaper alternative to LayerNorm. | |
| Rescales by root-mean-square of the activations instead of full | |
| mean/variance normalization. No bias, single learnable scale per dim. | |
| """ | |
| def __init__(self, dim: int, eps: float = 1e-5): | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| # compute in float32 for stability regardless of input dtype (bf16 etc.) | |
| dtype = x.dtype | |
| x = x.float() | |
| rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) | |
| out = x * rms | |
| return (out.to(dtype)) * self.weight | |