| 407 | |
| 408 | |
| 409 | class RMSNorm(nn.Module): |
| 410 | def __init__(self, dim, eps: float, elementwise_affine: bool = True): |
| 411 | super().__init__() |
| 412 | |
| 413 | self.eps = eps |
| 414 | |
| 415 | if isinstance(dim, numbers.Integral): |
| 416 | dim = (dim,) |
| 417 | |
| 418 | self.dim = torch.Size(dim) |
| 419 | |
| 420 | if elementwise_affine: |
| 421 | self.weight = nn.Parameter(torch.ones(dim)) |
| 422 | else: |
| 423 | self.weight = None |
| 424 | |
| 425 | def forward(self, hidden_states): |
| 426 | input_dtype = hidden_states.dtype |
| 427 | variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) |
| 428 | hidden_states = hidden_states * torch.rsqrt(variance + self.eps) |
| 429 | |
| 430 | if self.weight is not None: |
| 431 | # convert into half-precision if necessary |
| 432 | if self.weight.dtype in [torch.float16, torch.bfloat16]: |
| 433 | hidden_states = hidden_states.to(self.weight.dtype) |
| 434 | hidden_states = hidden_states * self.weight |
| 435 | else: |
| 436 | hidden_states = hidden_states.to(input_dtype) |
| 437 | |
| 438 | return hidden_states |
| 439 | |
| 440 | |
| 441 | class GlobalResponseNorm(nn.Module): |