| 210 | |
| 211 | |
| 212 | class RMSNorm(nn.Module): |
| 213 | def __init__(self, dim, eps: float, elementwise_affine: bool = True): |
| 214 | super().__init__() |
| 215 | |
| 216 | self.eps = eps |
| 217 | |
| 218 | if isinstance(dim, numbers.Integral): |
| 219 | dim = (dim,) |
| 220 | |
| 221 | self.dim = torch.Size(dim) |
| 222 | |
| 223 | if elementwise_affine: |
| 224 | self.weight = nn.Parameter(torch.ones(dim)) |
| 225 | else: |
| 226 | self.weight = None |
| 227 | |
| 228 | def forward(self, hidden_states): |
| 229 | input_dtype = hidden_states.dtype |
| 230 | variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) |
| 231 | hidden_states = hidden_states * torch.rsqrt(variance + self.eps) |
| 232 | |
| 233 | if self.weight is not None: |
| 234 | # convert into half-precision if necessary |
| 235 | if self.weight.dtype in [torch.float16, torch.bfloat16]: |
| 236 | hidden_states = hidden_states.to(self.weight.dtype) |
| 237 | hidden_states = hidden_states * self.weight |
| 238 | else: |
| 239 | hidden_states = hidden_states.to(input_dtype) |
| 240 | |
| 241 | return hidden_states |
| 242 | |
| 243 | |
| 244 | class GlobalResponseNorm(nn.Module): |