| import torch | |
| class LlamaRMSNorm(torch.nn.Module): | |
| def __init__(self, hidden_size, eps=1e-6): | |
| """ | |
| LlamaRMSNorm is equivalent to T5LayerNorm | |
| """ | |
| super().__init__() | |
| self.weight = torch.nn.Parameter(torch.ones(hidden_size)) | |
| self.variance_epsilon = eps | |
| def forward(self, hidden_states: torch.Tensor): | |
| input_dtype = hidden_states.dtype | |
| hidden_states = hidden_states.to(torch.float32) | |
| variance = hidden_states.pow(2).mean(-1, keepdim=True) | |
| hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) | |
| return self.weight.to(hidden_states.device) * hidden_states.to(input_dtype) | |