|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
import math |
|
|
from huggingface_hub import PyTorchModelHubMixin |
|
|
|
|
|
|
|
|
class RoPEPositionalEncoding(nn.Module): |
|
|
def __init__(self, dim, max_len=2048): |
|
|
super().__init__() |
|
|
self.dim = dim |
|
|
|
|
|
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) |
|
|
self.register_buffer("inv_freq", inv_freq) |
|
|
|
|
|
self._cached_cos = None |
|
|
self._cached_sin = None |
|
|
self._cached_len = 0 |
|
|
|
|
|
def _compute_cache(self, seq_len, device): |
|
|
if seq_len > self._cached_len or ( |
|
|
self._cached_cos is not None and self._cached_cos.device != device |
|
|
): |
|
|
t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) |
|
|
inv_freq = self.inv_freq.to(device) |
|
|
freqs = torch.outer(t, inv_freq) |
|
|
emb = torch.cat((freqs, freqs), dim=-1) |
|
|
|
|
|
self._cached_cos = emb.cos() |
|
|
self._cached_sin = emb.sin() |
|
|
self._cached_len = seq_len |
|
|
|
|
|
return ( |
|
|
self._cached_cos[:seq_len].to(device), |
|
|
self._cached_sin[:seq_len].to(device), |
|
|
) |
|
|
|
|
|
def rotate_half(self, x): |
|
|
x1 = x[..., : x.shape[-1] // 2] |
|
|
x2 = x[..., x.shape[-1] // 2 :] |
|
|
return torch.cat((-x2, x1), dim=-1) |
|
|
|
|
|
def apply_rope(self, q, k, seq_len): |
|
|
cos, sin = self._compute_cache(seq_len, q.device) |
|
|
cos = cos.unsqueeze(0).unsqueeze(0) |
|
|
sin = sin.unsqueeze(0).unsqueeze(0) |
|
|
|
|
|
q = (q * cos) + (self.rotate_half(q) * sin) |
|
|
k = (k * cos) + (self.rotate_half(k) * sin) |
|
|
|
|
|
return q, k |
|
|
|
|
|
|
|
|
class BitLinearFunction(torch.autograd.Function): |
|
|
@staticmethod |
|
|
def forward(ctx, input, weight, bias=None): |
|
|
scale = 127.0 / input.abs().max(dim=-1, keepdim=True).values.clamp(min=1e-5) |
|
|
x_quant = (input * scale).round().clamp(-128, 127) / scale |
|
|
|
|
|
w_scale = weight.abs().mean().clamp(min=1e-5) |
|
|
w_quant = (weight / w_scale).round().clamp(-1, 1) * w_scale |
|
|
|
|
|
ctx.save_for_backward(input, weight) |
|
|
ctx.w_quant = w_quant |
|
|
|
|
|
return F.linear(x_quant, w_quant, bias) |
|
|
|
|
|
@staticmethod |
|
|
def backward(ctx, grad_output): |
|
|
input, weight = ctx.saved_tensors |
|
|
w_quant = ctx.w_quant |
|
|
|
|
|
grad_input = grad_output.matmul(w_quant) |
|
|
|
|
|
grad_output_flat = grad_output.view(-1, grad_output.shape[-1]) |
|
|
input_flat = input.view(-1, input.shape[-1]) |
|
|
grad_weight = grad_output_flat.t().mm(input_flat) |
|
|
|
|
|
grad_bias = None |
|
|
if ctx.needs_input_grad[2]: |
|
|
grad_bias = grad_output_flat.sum(0) |
|
|
|
|
|
return grad_input, grad_weight, grad_bias |
|
|
|
|
|
|
|
|
class RigorousBitLinear(nn.Module): |
|
|
def __init__(self, in_features, out_features, bias=False): |
|
|
super().__init__() |
|
|
self.weight = nn.Parameter(torch.randn(out_features, in_features)) |
|
|
self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None |
|
|
|
|
|
def forward(self, x): |
|
|
return BitLinearFunction.apply(x, self.weight, self.bias) |
|
|
|
|
|
|
|
|
class RMSNorm(nn.Module): |
|
|
def __init__(self, dim, eps=1e-6): |
|
|
super().__init__() |
|
|
self.eps = eps |
|
|
self.weight = nn.Parameter(torch.ones(dim)) |
|
|
|
|
|
def forward(self, x): |
|
|
normed = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
|
|
return normed * self.weight |
|
|
|
|
|
|
|
|
class ImprovedBitAttention(nn.Module): |
|
|
def __init__(self, dim, heads=8, dropout=0.1, max_len=2048): |
|
|
super().__init__() |
|
|
self.heads = heads |
|
|
self.head_dim = dim // heads |
|
|
self.scale = self.head_dim ** -0.5 |
|
|
|
|
|
self.q_proj = RigorousBitLinear(dim, dim) |
|
|
self.k_proj = RigorousBitLinear(dim, dim) |
|
|
self.v_proj = RigorousBitLinear(dim, dim) |
|
|
self.out_proj = RigorousBitLinear(dim, dim) |
|
|
|
|
|
self.rope = RoPEPositionalEncoding(self.head_dim, max_len) |
|
|
self.dropout = nn.Dropout(dropout) |
|
|
|
|
|
def forward(self, x): |
|
|
B, L, D = x.shape |
|
|
|
|
|
q = self.q_proj(x).view(B, L, self.heads, self.head_dim).transpose(1, 2) |
|
|
k = self.k_proj(x).view(B, L, self.heads, self.head_dim).transpose(1, 2) |
|
|
v = self.v_proj(x).view(B, L, self.heads, self.head_dim).transpose(1, 2) |
|
|
|
|
|
q, k = self.rope.apply_rope(q, k, L) |
|
|
|
|
|
attn = (q @ k.transpose(-2, -1)) * self.scale |
|
|
|
|
|
mask = torch.tril(torch.ones(L, L, device=x.device, dtype=torch.bool)) |
|
|
attn = attn.masked_fill(~mask, float("-inf")) |
|
|
|
|
|
attn = F.softmax(attn, dim=-1) |
|
|
attn = self.dropout(attn) |
|
|
|
|
|
out = (attn @ v).transpose(1, 2).contiguous().view(B, L, D) |
|
|
return self.out_proj(out) |
|
|
|
|
|
|
|
|
|
|
|
class SwiGLUMLP(nn.Module): |
|
|
def __init__(self, dim, expansion=2.67, dropout=0.1): |
|
|
super().__init__() |
|
|
hidden = int(dim * expansion) |
|
|
|
|
|
|
|
|
self.gate_proj = RigorousBitLinear(dim, hidden) |
|
|
self.up_proj = RigorousBitLinear(dim, hidden) |
|
|
self.down_proj = RigorousBitLinear(hidden, dim) |
|
|
|
|
|
self.dropout = nn.Dropout(dropout) |
|
|
|
|
|
def forward(self, x): |
|
|
gate = F.silu(self.gate_proj(x)) |
|
|
up = self.up_proj(x) |
|
|
return self.down_proj(self.dropout(gate * up)) |
|
|
|
|
|
|
|
|
|
|
|
class ImprovedBitBlock(nn.Module): |
|
|
def __init__(self, dim, heads=8, dropout=0.1, max_len=2048): |
|
|
super().__init__() |
|
|
self.norm1 = RMSNorm(dim) |
|
|
self.attn = ImprovedBitAttention(dim, heads, dropout, max_len) |
|
|
self.norm2 = RMSNorm(dim) |
|
|
self.mlp = SwiGLUMLP(dim, dropout=dropout) |
|
|
|
|
|
def forward(self, x): |
|
|
x = x + self.attn(self.norm1(x)) |
|
|
x = x + self.mlp(self.norm2(x)) |
|
|
return x |
|
|
|
|
|
|
|
|
class ImprovedBitNet(nn.Module, PyTorchModelHubMixin): |
|
|
def __init__( |
|
|
self, |
|
|
vocab_size: int = 30522, |
|
|
dim: int = 768, |
|
|
depth: int = 12, |
|
|
heads: int = 12, |
|
|
max_len: int = 512, |
|
|
dropout: float = 0.05, |
|
|
): |
|
|
super().__init__() |
|
|
|
|
|
self.vocab_size = vocab_size |
|
|
self.dim = dim |
|
|
self.depth = depth |
|
|
|
|
|
|
|
|
self.token_emb = nn.Embedding(vocab_size, dim) |
|
|
|
|
|
|
|
|
self.blocks = nn.ModuleList( |
|
|
[ |
|
|
ImprovedBitBlock( |
|
|
dim=dim, |
|
|
heads=heads, |
|
|
dropout=dropout, |
|
|
max_len=max_len, |
|
|
) |
|
|
for _ in range(depth) |
|
|
] |
|
|
) |
|
|
|
|
|
|
|
|
self.norm = RMSNorm(dim) |
|
|
self.head = nn.Linear(dim, vocab_size) |
|
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
|
x = self.token_emb(x) |
|
|
|
|
|
for block in self.blocks: |
|
|
x = block(x) |
|
|
|
|
|
x = self.norm(x) |
|
|
logits = self.head(x) |
|
|
return logits |
|
|
|
