Upload Boopit 1.58-bit packed weights

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model.boopit filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,40 @@
+---
+license: apache-2.0
+language:
+- en
+library_name: pytorch
+pipeline_tag: text-generation
+datasets:
+- openbmb/Ultra-FineWeb-L1
+tags:
+- bitnet
+- ternary
+- 1.58-bit
+- boopit
+
+---
+
+This model was asked to be published under my account, not the creators. The compute came from https://huggingface.co/posts/ProCreations/855858308074329
+
+# Boopit 1
+
+A **27.3M**-parameter language model with a **4096** context window, trained from scratch with **BitNet b1.58 ternary weights** (weights in {-1,0,1} from step 0) on **4B tokens** of [openbmb/Ultra-FineWeb-L1](https://huggingface.co/datasets/openbmb/Ultra-FineWeb-L1).
+
+Packed weights (`model.boopit`) are **under 7MB**.
+
+Architecture is a small RoPE transformer in the same family as [babble / booper](https://github.com/kowo-co/babble): 6 layers, 512-wide, 8 heads, tied embeddings, RMSNorm, GELU MLP. Tokenizer is the byte-level BPE from [ProCreations/booper-pretrain](https://huggingface.co/ProCreations/booper-pretrain).
+
+Chat fine-tune: [ProCreations/boopit-1-chat](https://huggingface.co/ProCreations/boopit-1-chat).
+
+## Files
+
+- `model.boopit` — packed 1.58-bit ternary weights
+- `tokenizer.json` — BPE merges
+- `config.json` — architecture
+- `pack.py` / `modeling_boopit.py` — load + generate
+
+## Training
+
+- 1,369,964,544 tokens, 5,226 steps, batch 64 × 4096
+- NVIDIA RTX PRO 6000 Blackwell (96GB)
+

config.json

@@ -0,0 +1,8 @@
+{
+  "vocab_size": 16384,
+  "block_size": 4096,
+  "n_layer": 6,
+  "n_head": 8,
+  "n_embd": 512,
+  "bitnet": true
+}

loss.jsonl

@@ -0,0 +1,27 @@
+{"stage": "pretrain", "step": 200, "tokens": 52428800, "val_loss": 6.92905, "lr": 5.9999999999999995e-05, "elapsed_s": 174.1, "tokens_per_s": 301064.6, "samples": {"the cat": "the cat                             you \n            \n   ", "Hello": "HelloceanHighLevel\u0016Wu\u001dminisprowcommerSCOBYtournamKiel\ufffd\u0015HighLevelceanMaintenembarr\ufffdpropertpurpopurpouseumprowcean\u001d\ufffd\ufffdtournamceancommerprincLumierecommer\ufffd\ufffdtournam\u001d\u001d\ufffdceanhttps:\u0001tournamMaintenpropertBloomsizzatournam"}}
+{"stage": "pretrain", "step": 400, "tokens": 104857600, "val_loss": 5.888227, "lr": 0.00011999999999999999, "elapsed_s": 315.1, "tokens_per_s": 332758.8, "samples": {"the cat": "the cat  and that from so this so of  it and  on is  the  a a In and  , of has be  ", "Hello": "Helloat this   thecan   arenot  from    a at  In  and    in\n this     the     "}}
+{"stage": "pretrain", "step": 600, "tokens": 157286400, "val_loss": 5.120083, "lr": 0.00017999999999999998, "elapsed_s": 456.4, "tokens_per_s": 344590.7, "samples": {"the cat": "the cat their the of a to in with of for of to of\nthe  all the is to of from the for the the ", "Hello": "Hello of The a that the that for\nand  your the to this you and of and  of - to the to\nmore to"}}
+{"stage": "pretrain", "step": 800, "tokens": 209715200, "val_loss": 4.721462, "lr": 0.00023999999999999998, "elapsed_s": 598.2, "tokens_per_s": 350553.4, "samples": {"the cat": "the cat it and the the of the of the the your The the this is the are the in the from can to their in", "Hello": "Hello The bewith the your the other the the not The the to and are a a be the the new the was not "}}
+{"stage": "pretrain", "step": 1000, "tokens": 262144000, "val_loss": 4.487182, "lr": 0.0003, "elapsed_s": 742.6, "tokens_per_s": 352995.4, "samples": {"the cat": "the cat of of the it or is  a it to  a and an a to a of your the the the of an a", "Hello": "Hello in to for\u201ctheit that a this  not you to the and that of the it the can the that as and have"}}
+{"stage": "pretrain", "step": 1200, "tokens": 314572800, "val_loss": 4.432568, "lr": 0.00035999999999999997, "elapsed_s": 884.1, "tokens_per_s": 355806.3, "samples": {"the cat": "the cat", "Hello": "Hellobrand-202\nHal, to with a time a this the new also on their a a to to in for to and the best"}}
+{"stage": "pretrain", "step": 1400, "tokens": 367001600, "val_loss": 4.333194, "lr": 0.00041999999999999996, "elapsed_s": 1026.2, "tokens_per_s": 357640.2, "samples": {"the cat": "the catThe an (R by the the that to you that an to be was to the a most other in a the the in", "Hello": "Hello\ufffd\ufffdspitalIsoHuntcompon\ufffd\ufffdjacent\ufffdintellunwaver\u0015\u0015-by-stepUnboxed{VB.NETjacentpropert{spitaljacentampshireuccohttps:ampshire-by-stepCaliforveragehttps:KeeceSCOBYspitalhttps:\ufffdGMIcomponatformuccoSCOBY{\u0015\u0015{Sbobet88Falconspropertjacent"}}
+{"stage": "pretrain", "step": 1600, "tokens": 419430400, "val_loss": 4.239268, "lr": 0.00047999999999999996, "elapsed_s": 1168.4, "tokens_per_s": 358967.3, "samples": {"the cat": "the cat\nand the good is the best in a new is more the a other for the to the and is the it you is", "Hello": "Hello and the the you and at its that was and you of the the all and you in the this and it of a"}}
+{"stage": "pretrain", "step": 1800, "tokens": 471859200, "val_loss": 4.231241, "lr": 0.0005399999999999999, "elapsed_s": 1310.7, "tokens_per_s": 360014.8, "samples": {"the cat": "the cats. - the the line.\n - The |it will we will been this you more we on our few we will be at", "Hello": "Hello as a all was to the best and more in the the and the other family and a in the few and we is"}}
+{"stage": "pretrain", "step": 2000, "tokens": 524288000, "val_loss": 4.340728, "lr": 0.0006, "elapsed_s": 1455.1, "tokens_per_s": 360308.5, "samples": {"the cat": "the cat from your the improve offer for the - we to the you of the The a it will a other a new data to", "Hello": "Hello the it with our the  up          -        -  - s:- C St  im from"}}
+{"stage": "pretrain", "step": 2200, "tokens": 576716800, "val_loss": 4.365165, "lr": 0.0005996998743163473, "elapsed_s": 1597.1, "tokens_per_s": 361098.4, "samples": {"the cat": "the cat. into about the your  in the and by the their The time a of the for the their your and be their and", "Hello": "Hello and for a of the  the a a the your a less or more new to a and a than the , the "}}
+{"stage": "pretrain", "step": 2400, "tokens": 629145600, "val_loss": 4.739421, "lr": 0.0005987941320897245, "elapsed_s": 1739.1, "tokens_per_s": 361757.2, "samples": {"the cat": "the cat my toa ofhelp from in for of a for an in from of and the for for the are to are with in", "Hello": "Hello for    affecting that between for your your in a its your your to your to we without is provide of of which "}}
+{"stage": "pretrain", "step": 2600, "tokens": 681574400, "val_loss": 4.598028, "lr": 0.0005972847996280013, "elapsed_s": 1880.7, "tokens_per_s": 362412.8, "samples": {"the cat": "the cat for in many to over to a of into and is a which from your in from a and The your the your your", "Hello": "Hello for to for a  with with are to is to to I which is  , to to is the into your of to"}}
+{"stage": "pretrain", "step": 2800, "tokens": 734003200, "val_loss": 4.781966, "lr": 0.0005951752662102826, "elapsed_s": 2022.6, "tokens_per_s": 362904.1, "samples": {"the cat": "the cat ation ing you ing ation to which their with the and of in to is a to for it to the of has -", "Hello": "Hello  \n ofare-\n  \n .interand,isiss,withto  ling ,and ,for 's s, \nof are -  to  s,who"}}
+{"stage": "pretrain", "step": 3000, "tokens": 786432000, "val_loss": 4.784978, "lr": 0.0005924702688959738, "elapsed_s": 2166.7, "tokens_per_s": 362955.2, "samples": {"the cat": "the cat by ,  also  D of \n \n  an of M   the their  the for on was the - didn\u2019t  red ", "Hello": "Hello    can              I\u2019m   team the   this to intoredM using can with      "}}
+{"stage": "pretrain", "step": 3200, "tokens": 838860800, "val_loss": 4.744645, "lr": 0.0005891758818874849, "elapsed_s": 2308.7, "tokens_per_s": 363351.8, "samples": {"the cat": "the cat th of it many and of the for of a would at is for at I this in a or ofy, or the ", "Hello": "Hellos,wash D thteam redoffcom, Mof  an the -   - The by in no a in of no the a "}}
+{"stage": "pretrain", "step": 3400, "tokens": 891289600, "val_loss": 4.713177, "lr": 0.0005852995028903142, "elapsed_s": 2450.6, "tokens_per_s": 363696.6, "samples": {"the cat": "the cat This are a of He of I the the their are to the a a This to the a the the the their this", "Hello": "Hello    is their to This of of the it by\nthe the to the to the of of for of on would to "}}
+{"stage": "pretrain", "step": 3600, "tokens": 943718400, "val_loss": 4.69863, "lr": 0.0005808498365011436, "elapsed_s": 2592.8, "tokens_per_s": 363972.7, "samples": {"the cat": "the cat the to a and of a the and you for and and a and the a to of to the the no the in", "Hello": "Hello   \n     also       also  anin the a to with all to a would by to to in an"}}
+{"stage": "pretrain", "step": 3800, "tokens": 996147200, "val_loss": 4.71586, "lr": 0.0005758368746612447, "elapsed_s": 2735.0, "tokens_per_s": 364222.8, "samples": {"the cat": "the cat other the of in and on the are of to to with this a the when other of a and a the are are", "Hello": "Hello   in to in are in the the like a and to when the of is to of the with and the and in"}}
+{"stage": "pretrain", "step": 4000, "tokens": 1048576000, "val_loss": 4.802155, "lr": 0.0005702718742190908, "elapsed_s": 2878.6, "tokens_per_s": 364264.0, "samples": {"the cat": "the cat a and in the as the the or a a in for would to of in of they in this an and when a", "Hello": "Hello in in of the to the of the a a with the his the in which in for the and a to \u201cIt The"}}
+{"stage": "pretrain", "step": 4200, "tokens": 1101004800, "val_loss": 4.777148, "lr": 0.0005641673316525593, "elapsed_s": 3020.4, "tokens_per_s": 364516.8, "samples": {"the cat": "the cat can in can the the to with on with of the the are the of they in in on and and and on in", "Hello": "Hello and a the are have or a it and their can and the on the can the and a the the when his in"}}
+{"stage": "pretrain", "step": 4400, "tokens": 1153433600, "val_loss": 4.785015, "lr": 0.0005575369550074853, "elapsed_s": 3162.3, "tokens_per_s": 364740.5, "samples": {"the cat": "the cat the a on to of and a on a when their the in have by to the the of the the of the a", "Hello": "Hello of can the to in of of and you the a the a the and the the to and in in as of and"}}
+{"stage": "pretrain", "step": 4600, "tokens": 1205862400, "val_loss": 4.998295, "lr": 0.0005503956331155814, "elapsed_s": 3304.1, "tokens_per_s": 364956.2, "samples": {"the cat": "the cat of to have to of and or and and the a a in the the in also the a This of their the the", "Hello": "Hello they let the the the a can to in a of a the a and of the a the do the the in the"}}
+{"stage": "pretrain", "step": 4800, "tokens": 1258291200, "val_loss": 4.958861, "lr": 0.0005427594021608452, "elapsed_s": 3448.2, "tokens_per_s": 364913.0, "samples": {"the cat": "the cat is or in the to th on the to you the This The you of do to March to for a of on week", "Hello": "Hello no  This the the a their or and of to and this a are you you to a which a and of in "}}
+{"stage": "pretrain", "step": 5000, "tokens": 1310720000, "val_loss": 4.880846, "lr": 0.0005346454096695318, "elapsed_s": 3590.3, "tokens_per_s": 365074.5, "samples": {"the cat": "the cat on the the to and in his a on or and and at on to out as the the you the and a the", "Hello": "Hello\ufffdl      med          E          \ufffd                 "}}
+{"stage": "pretrain", "step": 5200, "tokens": 1363148800, "val_loss": 4.946123, "lr": 0.0005260718760045541, "elapsed_s": 3732.6, "tokens_per_s": 365201.3, "samples": {"the cat": "the cat a their the or a the they the The you a and a a on the a the you the an a to your", "Hello": "Helloust      th dostand much      This and of the on an which the the or in the in the a"}}
+{"stage": "pretrain", "step": 5226, "tokens": 1369964544, "val_loss": 4.925552, "lr": 0.0005249245559233004, "elapsed_s": 3751.5, "tokens_per_s": 365174.2, "samples": {"the cat": "the cat a to of a the the a a the the their the of the the receive in a their the This a in the", "Hello": "Hello         bal         boostl and This It the at a the the a a of the are "}}

model.boopit

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cf9224b9a9c64d2825dde840c6bdc089d46709d4661e1554dd37cd6ba72d47a
+size 5470363

model.py

@@ -0,0 +1,176 @@
+"""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)

modeling_boopit.py

@@ -0,0 +1,176 @@
+"""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)

pack.py

@@ -0,0 +1,131 @@
+"""Pack / unpack native 1.58-bit ternary weights (5 trits per byte)."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from model import Boopit, BoopitConfig, ternary_and_scale
+
+
+MAX_PACKED_BYTES = 7 * 1024 * 1024
+
+
+def pack_trits(t: np.ndarray) -> tuple[bytes, int]:
+    """Map {-1,0,1} -> {0,1,2} and pack 5 trits into each byte (3^5=243)."""
+    t = t.astype(np.int8).ravel()
+    coded = (t + 1).astype(np.uint8)
+    pad = (5 - (coded.size % 5)) % 5
+    if pad:
+        coded = np.concatenate([coded, np.full(pad, 1, dtype=np.uint8)])  # 0-trit
+    packed = (
+        coded[0::5].astype(np.uint16)
+        + coded[1::5].astype(np.uint16) * 3
+        + coded[2::5].astype(np.uint16) * 9
+        + coded[3::5].astype(np.uint16) * 27
+        + coded[4::5].astype(np.uint16) * 81
+    ).astype(np.uint8)
+    return packed.tobytes(), pad
+
+
+def unpack_trits(data: bytes, n: int, pad: int) -> np.ndarray:
+    packed = np.frombuffer(data, dtype=np.uint8).astype(np.uint16)
+    out = np.empty(packed.size * 5, dtype=np.int8)
+    v = packed
+    for i in range(5):
+        out[i::5] = (v % 3).astype(np.int8)
+        v //= 3
+    if pad:
+        out = out[: out.size - pad]
+    return (out[:n] - 1).astype(np.int8)
+
+
+def extract_packed(model: Boopit) -> dict:
+    tensors = {}
+    seen: set[int] = set()
+    for name, param in model.named_parameters():
+        if id(param) in seen:
+            continue
+        seen.add(id(param))
+        w = param.detach()
+        if w.ndim == 2:
+            t, scale = ternary_and_scale(w)
+            packed, pad = pack_trits(t.cpu().numpy())
+            tensors[name] = {
+                "kind": "ternary",
+                "shape": list(w.shape),
+                "scale": float(scale.cpu()),
+                "pad": pad,
+                "data": packed,
+            }
+        else:
+            tensors[name] = {
+                "kind": "fp16",
+                "shape": list(w.shape),
+                "data": w.detach().to(torch.float16).cpu().numpy().tobytes(),
+            }
+    return tensors
+
+
+def packed_nbytes(tensors: dict) -> int:
+    return sum(len(v["data"]) for v in tensors.values())
+
+
+def save_packed(model: Boopit, path: Path, extra: dict | None = None) -> int:
+    path = Path(path)
+    tensors = extract_packed(model)
+    nbytes = packed_nbytes(tensors)
+    payload = {
+        "format": "boopit-1.58",
+        "config": model.config.to_dict(),
+        "tied": ["tok_emb.weight", "lm_head.weight"],
+        "extra": extra or {},
+        "tensors": {},
+    }
+    blob_parts = []
+    offset = 0
+    for name, spec in tensors.items():
+        data = spec["data"]
+        entry = {k: v for k, v in spec.items() if k != "data"}
+        entry["offset"] = offset
+        entry["nbytes"] = len(data)
+        payload["tensors"][name] = entry
+        blob_parts.append(data)
+        offset += len(data)
+    header = json.dumps(payload, separators=(",", ":")).encode("utf-8")
+    header_len = len(header).to_bytes(8, "little")
+    blob = b"".join(blob_parts)
+    raw = header_len + header + blob
+    if len(raw) >= MAX_PACKED_BYTES:
+        raise RuntimeError(f"packed model is {len(raw)} bytes, limit is {MAX_PACKED_BYTES}")
+    path.write_bytes(raw)
+    return len(raw)
+
+
+def load_packed(path: Path, device: torch.device | str = "cpu") -> Boopit:
+    raw = Path(path).read_bytes()
+    header_len = int.from_bytes(raw[:8], "little")
+    payload = json.loads(raw[8 : 8 + header_len].decode("utf-8"))
+    blob = raw[8 + header_len :]
+    cfg = BoopitConfig.from_dict(payload["config"])
+    model = Boopit(cfg)
+    named = dict(model.named_parameters())
+    loaded: set[int] = set()
+    with torch.no_grad():
+        for name, spec in payload["tensors"].items():
+            param = named[name]
+            chunk = blob[spec["offset"] : spec["offset"] + spec["nbytes"]]
+            if spec["kind"] == "ternary":
+                t = unpack_trits(chunk, int(np.prod(spec["shape"])), spec["pad"])
+                w = torch.from_numpy(t.astype(np.float32).reshape(spec["shape"])) * float(spec["scale"])
+                param.copy_(w.to(dtype=param.dtype))
+            else:
+                arr = np.frombuffer(chunk, dtype=np.float16).reshape(spec["shape"])
+                param.copy_(torch.from_numpy(arr.copy()).to(dtype=param.dtype))
+            loaded.add(id(param))
+        # tied lm_head
+        if id(model.lm_head.weight) not in loaded:
+            model.lm_head.weight = model.tok_emb.weight
+    return model.to(device)

tokenizer.py

@@ -0,0 +1,80 @@
+"""Byte-level BPE compatible with babble / booper-pretrain tokenizer.json."""
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+
+_CHUNK_RE = re.compile(r"\s+|\S+")
+
+
+def _merge_ids(ids: list[int], a: int, b: int, new_id: int) -> list[int]:
+    if len(ids) < 2:
+        return ids
+    out: list[int] = []
+    i = 0
+    n = len(ids)
+    while i < n:
+        if i + 1 < n and ids[i] == a and ids[i + 1] == b:
+            out.append(new_id)
+            i += 2
+        else:
+            out.append(ids[i])
+            i += 1
+    return out
+
+
+class BPETokenizer:
+    def __init__(self, merges: list[tuple[int, int, int]]) -> None:
+        self.merges = merges
+        vocab: dict[int, bytes] = {i: bytes([i]) for i in range(256)}
+        for a, b, new_id in merges:
+            vocab[new_id] = vocab[a] + vocab[b]
+        self.vocab = vocab
+        self._ranks = {(a, b): i for i, (a, b, _) in enumerate(merges)}
+        self._pair_to_id = {(a, b): new_id for a, b, new_id in merges}
+        base = 256 + len(merges)
+        self.pad, self.bos, self.sep, self.eos = base, base + 1, base + 2, base + 3
+
+    @property
+    def vocab_size(self) -> int:
+        return 256 + len(self.merges) + 4
+
+    def _encode_chunk(self, chunk: str) -> list[int]:
+        ids = list(chunk.encode("utf-8"))
+        ranks = self._ranks
+        pair_to_id = self._pair_to_id
+        while len(ids) >= 2:
+            best_rank = None
+            best_pair = None
+            for a, b in zip(ids, ids[1:]):
+                r = ranks.get((a, b))
+                if r is not None and (best_rank is None or r < best_rank):
+                    best_rank = r
+                    best_pair = (a, b)
+            if best_pair is None:
+                break
+            ids = _merge_ids(ids, best_pair[0], best_pair[1], pair_to_id[best_pair])
+        return ids
+
+    def encode(self, text: str) -> list[int]:
+        ids: list[int] = []
+        for chunk in _CHUNK_RE.findall(text):
+            ids.extend(self._encode_chunk(chunk))
+        return ids
+
+    def decode(self, ids: list[int]) -> str:
+        raw = bytearray()
+        for i in ids:
+            piece = self.vocab.get(i)
+            if piece is not None:
+                raw.extend(piece)
+        return bytes(raw).decode("utf-8", errors="replace")
+
+    def to_json(self, path: Path) -> None:
+        path.write_text(json.dumps({"merges": [list(m) for m in self.merges]}))
+
+    @classmethod
+    def from_json(cls, path: Path | str) -> "BPETokenizer":
+        raw = json.loads(Path(path).read_text())
+        return cls([tuple(m) for m in raw["merges"]])