Root-mean-square layer normalisation without a learnable scale/bias, as used inside BitLinear per the training manuscript.
| 65 | # ────────────────────────────────────────────────────────────────────────────── |
| 66 | |
| 67 | class _RMSNorm(nn.Module): |
| 68 | """ |
| 69 | Root-mean-square layer normalisation without a learnable scale/bias, |
| 70 | as used inside BitLinear per the training manuscript. |
| 71 | """ |
| 72 | |
| 73 | def __init__(self, dim: int, eps: float = 1e-6) -> None: |
| 74 | super().__init__() |
| 75 | self.eps = eps |
| 76 | self.dim = dim |
| 77 | |
| 78 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 79 | # x: [..., dim] |
| 80 | rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).sqrt() |
| 81 | return x / rms |
| 82 | |
| 83 | |
| 84 | # ────────────────────────────────────────────────────────────────────────────── |