| # SPDX-License-Identifier: BUSL-1.1 | |
| # Copyright (c) 2024-2026 Lucas Ricardo Mella Chillemi | |
| """RMS Normalization — Llama-style, más eficiente que LayerNorm.""" | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| class RMSNorm(nn.Module): | |
| """ | |
| Root Mean Square Layer Normalization. | |
| Más eficiente que LayerNorm: omite el centrado, solo normaliza por RMS. | |
| Usado en Llama, Qwen, Mistral. | |
| """ | |
| def __init__(self, dim: int, eps: float = 1e-6): | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| rms = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) | |
| return (x.float() * rms).type_as(x) * self.weight | |