| 57 | |
| 58 | |
| 59 | class LPLayerNorm(torch.nn.LayerNorm): |
| 60 | def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True): |
| 61 | super().__init__( |
| 62 | normalized_shape=normalized_shape, |
| 63 | eps=eps, |
| 64 | elementwise_affine=elementwise_affine, |
| 65 | ) |
| 66 | |
| 67 | def forward(self, x): |
| 68 | module_device = x.device |
| 69 | downcast_x = _cast_if_autocast_enabled(x) |
| 70 | downcast_weight = ( |
| 71 | _cast_if_autocast_enabled(self.weight) |
| 72 | if self.weight is not None |
| 73 | else self.weight |
| 74 | ) |
| 75 | downcast_bias = ( |
| 76 | _cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias |
| 77 | ) |
| 78 | with torch.autocast(enabled=False, device_type=module_device.type): |
| 79 | return torch.nn.functional.layer_norm( |
| 80 | downcast_x, |
| 81 | self.normalized_shape, |
| 82 | downcast_weight, |
| 83 | downcast_bias, |
| 84 | self.eps, |
| 85 | ) |
| 86 | |
| 87 | |
| 88 | class SharedEmbedding(nn.Embedding): |