Initialize the RMSNorm normalization layer. Args: dim (int): The dimension of the input tensor. eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. Attributes: eps (float): A small value added to the denominat
| 113 | |
| 114 | |
| 115 | class RMSNorm(nn.Module): |
| 116 | """ |
| 117 | Initialize the RMSNorm normalization layer. |
| 118 | |
| 119 | Args: |
| 120 | dim (int): The dimension of the input tensor. |
| 121 | eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. |
| 122 | |
| 123 | Attributes: |
| 124 | eps (float): A small value added to the denominator for numerical stability. |
| 125 | weight (nn.Parameter): Learnable scaling parameter. |
| 126 | |
| 127 | """ |
| 128 | |
| 129 | def __init__(self, dim: int, eps: float = 1e-6): |
| 130 | super().__init__() |
| 131 | self.eps = eps |
| 132 | self.weight = nn.Parameter(torch.ones(dim)) |
| 133 | |
| 134 | def _norm(self, x: torch.Tensor): |
| 135 | return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
| 136 | |
| 137 | def forward(self, x: torch.Tensor): |
| 138 | output = self._norm(x.float()).type_as(x) |
| 139 | return output * self.weight |
| 140 | |
| 141 | def reset_parameters(self): |
| 142 | torch.nn.init.ones_(self.weight) # type: ignore |
| 143 | |
| 144 | |
| 145 | class Attention(nn.Module): |