File size: 6,407 Bytes
d8b3c96 | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 | """Boopit: ~28M BitNet b1.58 transformer with RoPE and 4096 context."""
from __future__ import annotations
import math
from dataclasses import asdict, dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F
@dataclass
class BoopitConfig:
vocab_size: int = 16384
block_size: int = 4096
n_layer: int = 6
n_head: int = 8
n_embd: int = 512
# BitNet-style: ternary weights from step 0, 8-bit absmean activations.
bitnet: bool = True
def to_dict(self) -> dict:
return asdict(self)
@classmethod
def from_dict(cls, raw: dict) -> "BoopitConfig":
known = {k: raw[k] for k in cls.__dataclass_fields__ if k in raw}
return cls(**known)
def activation_quant(x: torch.Tensor) -> torch.Tensor:
scale = 127.0 / x.abs().mean(dim=-1, keepdim=True).clamp(min=1e-5)
y = (x * scale).round().clamp(-128, 127) / scale
return x + (y - x).detach()
def weight_quant(w: torch.Tensor) -> torch.Tensor:
scale = w.abs().mean().clamp(min=1e-5)
y = (w / scale).round().clamp(-1, 1) * scale
return w + (y - w).detach()
def ternary_and_scale(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
scale = w.abs().mean().clamp(min=1e-5)
t = (w / scale).round().clamp(-1, 1).to(torch.int8)
return t, scale.detach().to(torch.float16)
class BitLinear(nn.Module):
def __init__(self, in_features: int, out_features: int) -> None:
super().__init__()
self.weight = nn.Parameter(torch.empty(out_features, in_features))
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
def forward(self, x: torch.Tensor) -> torch.Tensor:
return F.linear(activation_quant(x), weight_quant(self.weight))
class BitEmbedding(nn.Module):
def __init__(self, num_embeddings: int, embedding_dim: int) -> None:
super().__init__()
self.weight = nn.Parameter(torch.empty(num_embeddings, embedding_dim))
nn.init.normal_(self.weight, mean=0.0, std=0.02)
def forward(self, idx: torch.Tensor) -> torch.Tensor:
return F.embedding(idx, weight_quant(self.weight))
class RMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-5) -> None:
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x.float()
rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
return (x * rms * self.weight.float()).to(self.weight.dtype)
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
x1, x2 = x[..., ::2], x[..., 1::2]
return torch.stack((-x2, x1), dim=-1).flatten(-2)
class Rotary(nn.Module):
def __init__(self, head_dim: int, max_seq: int, base: float = 10000.0) -> None:
super().__init__()
inv = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
t = torch.arange(max_seq).float()
freqs = torch.outer(t, inv)
self.register_buffer("cos", torch.cos(freqs), persistent=False)
self.register_buffer("sin", torch.sin(freqs), persistent=False)
def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
t = q.size(-2)
cos = self.cos[:t].to(dtype=q.dtype)
sin = self.sin[:t].to(dtype=q.dtype)
cos = cos.repeat_interleave(2, dim=-1)[None, None, :, :]
sin = sin.repeat_interleave(2, dim=-1)[None, None, :, :]
return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin
class Block(nn.Module):
def __init__(self, cfg: BoopitConfig, rope: Rotary) -> None:
super().__init__()
self.n_head = cfg.n_head
self.head_dim = cfg.n_embd // cfg.n_head
self.rope = rope
self.ln_1 = RMSNorm(cfg.n_embd)
self.qkv = BitLinear(cfg.n_embd, 3 * cfg.n_embd)
self.proj = BitLinear(cfg.n_embd, cfg.n_embd)
self.ln_2 = RMSNorm(cfg.n_embd)
self.fc = BitLinear(cfg.n_embd, 4 * cfg.n_embd)
self.up = BitLinear(4 * cfg.n_embd, cfg.n_embd)
def forward(self, x: torch.Tensor) -> torch.Tensor:
b, t, c = x.shape
h = self.ln_1(x)
qkv = self.qkv(h).view(b, t, 3, self.n_head, self.head_dim)
q, k, v = qkv.unbind(2)
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
q, k = self.rope(q, k)
attn = F.scaled_dot_product_attention(q, k, v, is_causal=True)
attn = attn.transpose(1, 2).contiguous().view(b, t, c)
x = x + self.proj(attn)
h = self.ln_2(x)
x = x + self.up(F.gelu(self.fc(h), approximate="tanh"))
return x
class Boopit(nn.Module):
def __init__(self, cfg: BoopitConfig | None = None) -> None:
super().__init__()
self.config = cfg or BoopitConfig()
c = self.config
if c.n_embd % c.n_head:
raise ValueError("n_embd must divide n_head")
self.tok_emb = BitEmbedding(c.vocab_size, c.n_embd)
self.rope = Rotary(c.n_embd // c.n_head, c.block_size)
self.blocks = nn.ModuleList(Block(c, self.rope) for _ in range(c.n_layer))
self.ln_f = RMSNorm(c.n_embd)
self.lm_head = BitLinear(c.n_embd, c.vocab_size)
self.lm_head.weight = self.tok_emb.weight # tied, still ternary
def forward(self, idx: torch.Tensor) -> torch.Tensor:
t = idx.size(1)
if t > self.config.block_size:
raise ValueError(f"sequence {t} exceeds block_size {self.config.block_size}")
x = self.tok_emb(idx)
for block in self.blocks:
x = block(x)
return self.lm_head(self.ln_f(x))
def num_params(self) -> int:
seen: dict[int, int] = {}
total = 0
for p in self.parameters():
if id(p) not in seen:
seen[id(p)] = p.numel()
total += p.numel()
return total
def sequence_loss(model: Boopit, tokens: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
logits = model(tokens[:, :-1])
targets = tokens[:, 1:]
if mask is None:
return F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
per = F.cross_entropy(
logits.reshape(-1, logits.size(-1)),
targets.reshape(-1),
reduction="none",
).view_as(targets)
scale = mask[:, 1:].to(per.dtype)
return (per * scale).sum() / scale.sum().clamp(min=1.0)
|