model.py

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