File size: 1,628 Bytes
5d2c747 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
"""
A collection of normalization layers.
"""
import torch
from torch.nn import functional as F
class LayerNorm(torch.nn.Module):
"""LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False"""
# taken from nanoGPT
def __init__(self, dim, bias):
super().__init__()
self.weight = torch.nn.Parameter(torch.ones(dim))
self.bias = torch.nn.Parameter(torch.zeros(dim)) if bias else None
def forward(self, x):
"""Apply Layer Norm"""
return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5)
class RMSNorm(torch.nn.Module):
"""
RMSNorm (https://arxiv.org/abs/1910.07467), implementation from
https://github.com/meta-llama/llama3/blob/main/llama/model.py
"""
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
self.eps = eps
self.weight = torch.nn.Parameter(torch.ones(dim))
def _norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
"""Apply RMSNorm"""
output = self._norm(x.float()).type_as(x)
return output * self.weight
NORMALIZATION_DICT = {
"rms_norm": lambda dim, bias: RMSNorm(dim=dim),
"layer_norm": lambda dim, bias: LayerNorm(dim=dim, bias=bias),
"none": lambda dim, bias: torch.nn.Identity(),
}
def build_normalization(normalization_name, dim, bias=None):
"""
Build the normalization layer
Available options: rmsnorm, layernorm
- Bias is ignored for RMSNorm
"""
return NORMALIZATION_DICT[normalization_name](dim=dim, bias=bias)
|