RMSNorm layer.
| 63 | return x * torch.sigmoid(1.702 * x) |
| 64 | |
| 65 | class LayerNorm(nn.RMSNorm): |
| 66 | """ |
| 67 | RMSNorm layer. |
| 68 | """ |
| 69 | def __init__(self, dim, eps=1e-5): |
| 70 | super().__init__(dim, eps=eps) |
| 71 | def forward(self, x): |
| 72 | type_ = x.dtype |
| 73 | ret = super().forward(x.type(torch.float32)) |
| 74 | return ret.type(type_) |
| 75 | |
| 76 | class MLP(nn.Module): |
| 77 | def __init__( |