| 517 | return torch.nn.functional.linear(input,weight,bias) |
| 518 | |
| 519 | class RMSNorm(torch.nn.Module): |
| 520 | def __init__(self, module): |
| 521 | super().__init__() |
| 522 | self.module = module |
| 523 | |
| 524 | def forward(self,hidden_states,**kwargs): |
| 525 | weight= cast_weight(self.module,hidden_states) |
| 526 | input_dtype = hidden_states.dtype |
| 527 | variance = hidden_states.to(torch.float32).square().mean(-1, keepdim=True) |
| 528 | hidden_states = hidden_states * torch.rsqrt(variance + self.module.eps) |
| 529 | hidden_states = hidden_states.to(input_dtype) * weight |
| 530 | return hidden_states |
| 531 | |
| 532 | def replace_layer(model): |
| 533 | for name, module in model.named_children(): |