19 lines
686 B
Python
19 lines
686 B
Python
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)
|