| 6 | |
| 7 | |
| 8 | class RMSNorm(torch.nn.Module): |
| 9 | def __init__(self, dim, eps, elementwise_affine=True): |
| 10 | super().__init__() |
| 11 | self.eps = eps |
| 12 | if elementwise_affine: |
| 13 | self.weight = torch.nn.Parameter(torch.ones((dim,))) |
| 14 | else: |
| 15 | self.weight = None |
| 16 | |
| 17 | def forward(self, hidden_states): |
| 18 | input_dtype = hidden_states.dtype |
| 19 | variance = hidden_states.to(torch.float32).square().mean(-1, keepdim=True) |
| 20 | hidden_states = hidden_states * torch.rsqrt(variance + self.eps) |
| 21 | hidden_states = hidden_states.to(input_dtype) |
| 22 | if self.weight is not None: |
| 23 | hidden_states = hidden_states * self.weight |
| 24 | return hidden_states |
| 25 | |
| 26 | |
| 27 | |