OpenTransformer/binary-transformers

1#!/usr/bin/env python32"""3BIT-LEVEL TRANSFORMER - The Ultimate Zero-Overhead Model4Vocab = 2 (just 0 and 1)5No tokenization. No bytes. Pure binary.6 7Each byte becomes 8 tokens (bits).8Model learns ALL structure from raw bits.9"""10 11import sys12import math13import time14import torch15import torch.nn as nn16import torch.nn.functional as F17from collections import deque18 19DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")20torch.backends.cuda.matmul.allow_tf32 = True21 22# BIT-LEVEL CONFIG - ABSOLUTE UNIT23CONFIG = {24    "d": 768,        # GPT-2 small size25    "layers": 12,    # DEEP for bit pattern learning26    "heads": 12,27    "vocab": 2,      # JUST 0 AND 1!28    "ctx": 4096,     # 512 bytes of context29}30 31LR = 3e-4           # learning rate32UPDATE_EVERY = 2048  # bits between updates (256 bytes worth) - BIGGER BATCHES33PRINT_EVERY = 100000   # bits34 35class BitAttention(nn.Module):36    def __init__(self, d, h):37        super().__init__()38        self.h, self.dk = h, d // h39        self.qkv = nn.Linear(d, 3 * d, bias=False)40        self.proj = nn.Linear(d, d, bias=False)41        42    def forward(self, x, mask=None):43        B, N, D = x.shape44        qkv = self.qkv(x).view(B, N, 3, self.h, self.dk).permute(2, 0, 3, 1, 4)45        q, k, v = qkv[0], qkv[1], qkv[2]46        att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk)47        if mask is not None:48            att = att + mask49        return self.proj((F.softmax(att, -1) @ v).transpose(1, 2).reshape(B, N, D))50 51class BitBlock(nn.Module):52    def __init__(self, d, h):53        super().__init__()54        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)55        self.attn = BitAttention(d, h)56        self.ff = nn.Sequential(nn.Linear(d, 4*d), nn.GELU(), nn.Linear(4*d, d))57        58    def forward(self, x, mask):59        x = x + self.attn(self.ln1(x), mask)60        return x + self.ff(self.ln2(x))61 62class BitTransformer(nn.Module):63    """Transformer with vocab=2 (just 0 and 1)"""64    def __init__(self, cfg):65        super().__init__()66        d, L, h = cfg["d"], cfg["layers"], cfg["heads"]67        self.emb = nn.Embedding(2, d)  # ONLY 2 EMBEDDINGS!68        self.blocks = nn.ModuleList([BitBlock(d, h) for _ in range(L)])69        self.ln = nn.LayerNorm(d)70        self.head = nn.Linear(d, 2, bias=False)  # predict 0 or 171        72    def forward(self, x):73        B, N = x.shape74        mask = torch.triu(torch.ones(N, N, device=x.device), 1) * -1e975        h = self.emb(x)76        for block in self.blocks:77            h = block(h, mask)78        return self.head(self.ln(h))79    80    def count_params(self):81        return sum(p.numel() for p in self.parameters())82 83def byte_to_bits(byte_val):84    """Convert byte to 8 bits (MSB first)"""85    return [(byte_val >> (7 - i)) & 1 for i in range(8)]86 87def bits_to_byte(bits):88    """Convert 8 bits back to byte"""89    val = 090    for i, b in enumerate(bits[:8]):91        val |= (b << (7 - i))92    return val93 94class BitTrainer:95    def __init__(self, model, lr=LR):96        self.model = model.to(DEVICE)97        self.opt = torch.optim.AdamW(model.parameters(), lr=lr)98        self.ctx_size = CONFIG["ctx"]99        self.buffer = deque(maxlen=self.ctx_size + 1)100        101        self.bits_seen = 0102        self.bytes_seen = 0103        self.total_loss = 0.0104        self.updates = 0105        self.start_time = time.time()106        107    def ingest_byte(self, byte_val):108        """Convert byte to 8 bits and absorb"""109        bits = byte_to_bits(byte_val)110        for bit in bits:111            self.buffer.append(bit)112            self.bits_seen += 1113            114            if len(self.buffer) >= UPDATE_EVERY + 1 and self.bits_seen % UPDATE_EVERY == 0:115                self._update()116        117        self.bytes_seen += 1118        119        if self.bits_seen % PRINT_EVERY == 0:120            self._print_stats()121            122        if self.bytes_seen % 500000 == 0 and self.bytes_seen > 0:123            self._save()124    125    def _update(self):126        bits = list(self.buffer)127        x = torch.tensor(bits[:-1], device=DEVICE, dtype=torch.long).unsqueeze(0)128        y = torch.tensor(bits[1:], device=DEVICE, dtype=torch.long).unsqueeze(0)129        130        self.model.train()131        logits = self.model(x)132        loss = F.cross_entropy(133            logits[:, -UPDATE_EVERY:].reshape(-1, 2),134            y[:, -UPDATE_EVERY:].reshape(-1)135        )136        137        self.opt.zero_grad()138        loss.backward()139        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)140        self.opt.step()141        142        self.total_loss += loss.item()143        self.updates += 1144    145    def _print_stats(self):146        elapsed = time.time() - self.start_time147        bits_per_sec = self.bits_seen / elapsed if elapsed > 0 else 0148        bytes_per_sec = self.bytes_seen / elapsed if elapsed > 0 else 0149        avg_loss = self.total_loss / max(1, self.updates)150        151        # For bits: random is 1.0 (coin flip), lower = learning152        # Entropy in bits per bit153        entropy = avg_loss / math.log(2)154        compression = (1.0 - entropy) * 100  # % compression vs random155        156        print(f"[{elapsed:.0f}s] {self.bytes_seen/1000:.1f}KB | {bytes_per_sec/1000:.1f} KB/s | "157              f"loss={avg_loss:.4f} | entropy={entropy:.3f} bit/bit | "158              f"compression={compression:.1f}%", flush=True)159    160    def _save(self):161        avg_loss = self.total_loss / max(1, self.updates)162        kb = self.bytes_seen // 1000163        ckpt = {164            "model": self.model.state_dict(),165            "bits": self.bits_seen,166            "bytes": self.bytes_seen,167            "loss": avg_loss,168        }169        torch.save(ckpt, f"/workspace/bit_ckpt_{kb}kb.pt")170        print(f"[SAVED] bit_ckpt_{kb}kb.pt", flush=True)171 172def main():173    print(f"BIT-LEVEL TRANSFORMER - Vocab = 2 (just 0 and 1)", flush=True)174    print(f"Config: {CONFIG}", flush=True)175    print(f"Device: {DEVICE}", flush=True)176    177    model = BitTransformer(CONFIG)178    params = model.count_params()179    print(f"Parameters: {params:,} ({params/1e6:.2f}M)", flush=True)180    print(f"Vocab: 2 (literally just 0 and 1)", flush=True)181    print(f"Each byte = 8 bit tokens", flush=True)182    183    trainer = BitTrainer(model)184    185    print(f"Listening for bytes (FAST batch mode)...", flush=True)186    187    # Read in large chunks for speed188    CHUNK_SIZE = 8192  # 8KB chunks = 65536 bits189    while True:190        chunk = sys.stdin.buffer.read(CHUNK_SIZE)191        if not chunk:192            break193        for byte in chunk:194            trainer.ingest_byte(byte)195    196    print(f"Stream ended. Total: {trainer.bytes_seen:,} bytes = {trainer.bits_seen:,} bits", flush=True)197 198if __name__ == "__main__":199    main()200