| 274 | # Norms & Blocks |
| 275 | # ---------------------------- |
| 276 | class RMSNorm(nn.Module): |
| 277 | def __init__(self, dim, eps=1e-5): |
| 278 | super().__init__() |
| 279 | self.eps = eps |
| 280 | self.weight = nn.Parameter(torch.ones(dim)) |
| 281 | |
| 282 | def norm(self, x): |
| 283 | return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) |
| 284 | |
| 285 | def forward(self, x): |
| 286 | dtype = x.dtype |
| 287 | return self.norm(x.float()).to(dtype) * self.weight |
| 288 | |
| 289 | |
| 290 | class AttentionModule(nn.Module): |