"""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)