| 18 | |
| 19 | |
| 20 | class RMSNorm(nn.Module): |
| 21 | def __init__( |
| 22 | self, |
| 23 | dim: int, |
| 24 | elementwise_affine=True, |
| 25 | eps: float = 1e-6, |
| 26 | device=None, |
| 27 | dtype=None, |
| 28 | ): |
| 29 | """ |
| 30 | Initialize the RMSNorm normalization layer. |
| 31 | |
| 32 | Args: |
| 33 | dim (int): The dimension of the input tensor. |
| 34 | eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. |
| 35 | |
| 36 | Attributes: |
| 37 | eps (float): A small value added to the denominator for numerical stability. |
| 38 | weight (nn.Parameter): Learnable scaling parameter. |
| 39 | |
| 40 | """ |
| 41 | factory_kwargs = {"device": device, "dtype": dtype} |
| 42 | super().__init__() |
| 43 | self.eps = eps |
| 44 | if elementwise_affine: |
| 45 | self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) |
| 46 | |
| 47 | def _norm(self, x): |
| 48 | """ |
| 49 | Apply the RMSNorm normalization to the input tensor. |
| 50 | |
| 51 | Args: |
| 52 | x (torch.Tensor): The input tensor. |
| 53 | |
| 54 | Returns: |
| 55 | torch.Tensor: The normalized tensor. |
| 56 | |
| 57 | """ |
| 58 | return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
| 59 | |
| 60 | def forward(self, x): |
| 61 | """ |
| 62 | Forward pass through the RMSNorm layer. |
| 63 | |
| 64 | Args: |
| 65 | x (torch.Tensor): The input tensor. |
| 66 | |
| 67 | Returns: |
| 68 | torch.Tensor: The output tensor after applying RMSNorm. |
| 69 | |
| 70 | """ |
| 71 | output = self._norm(x.float()).type_as(x) |
| 72 | if hasattr(self, "weight"): |
| 73 | output = output * self.weight |
| 74 | return output |
| 75 | |
| 76 | |
| 77 | ACTIVATION_FUNCTIONS = { |