| 139 | |
| 140 | |
| 141 | class MixedFusedRMSNorm(nn.Module): |
| 142 | # Extracted from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py |
| 143 | def __init__(self, hidden_size, eps=1e-6): |
| 144 | super().__init__() |
| 145 | self.weight = nn.Parameter(torch.ones(hidden_size)) |
| 146 | self.variance_epsilon = eps |
| 147 | |
| 148 | def forward(self, hidden_states): |
| 149 | input_dtype = hidden_states.dtype |
| 150 | hidden_states = hidden_states.to(torch.float32) |
| 151 | variance = hidden_states.pow(2).mean(-1, keepdim=True) |
| 152 | hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) |
| 153 | return self.weight * hidden_states.to(input_dtype) |
| 154 | |
| 155 | |
| 156 | class FlashSelfAttention(torch.nn.Module): |