OpenTransformer/binary-transformers

1#!/usr/bin/env python32"""3BINARY TRANSFORMER - Raw network bytes → neural network4No tokenizer. No preprocessing. Just bytes.5 6Vocab = 256 (one token per byte value 0x00-0xFF)7Input: Raw bytes from network stream via stdin8"""9 10import sys11import math12import time13import torch14import torch.nn as nn15import torch.nn.functional as F16from collections import deque17 18DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")19torch.backends.cuda.matmul.allow_tf32 = True20 21# Binary model config - TINY for speed22CONFIG = {23    "d": 128,       # smaller embedding24    "layers": 3,    # fewer layers25    "heads": 4,     26    "vocab": 256,   # ONE TOKEN PER BYTE27    "ctx": 1024,    # longer context (bytes are fine-grained)28}29 30LR = 3e-431UPDATE_EVERY = 64   # bytes between updates32PRINT_EVERY = 50000 # bytes between stats33 34class ByteAttention(nn.Module):35    def __init__(self, d, h):36        super().__init__()37        self.h, self.dk = h, d // h38        self.qkv = nn.Linear(d, 3 * d, bias=False)39        self.proj = nn.Linear(d, d, bias=False)40        41    def forward(self, x, mask=None):42        B, N, D = x.shape43        qkv = self.qkv(x).view(B, N, 3, self.h, self.dk).permute(2, 0, 3, 1, 4)44        q, k, v = qkv[0], qkv[1], qkv[2]45        att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk)46        if mask is not None:47            att = att + mask48        return self.proj((F.softmax(att, -1) @ v).transpose(1, 2).reshape(B, N, D))49 50class ByteBlock(nn.Module):51    def __init__(self, d, h):52        super().__init__()53        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)54        self.attn = ByteAttention(d, h)55        self.ff = nn.Sequential(nn.Linear(d, 4*d), nn.GELU(), nn.Linear(4*d, d))56        57    def forward(self, x, mask):58        x = x + self.attn(self.ln1(x), mask)59        return x + self.ff(self.ln2(x))60 61class BinaryTransformer(nn.Module):62    def __init__(self, cfg):63        super().__init__()64        d, L, h, V = cfg["d"], cfg["layers"], cfg["heads"], cfg["vocab"]65        self.emb = nn.Embedding(V, d)  # 256 embeddings, one per byte66        self.blocks = nn.ModuleList([ByteBlock(d, h) for _ in range(L)])67        self.ln = nn.LayerNorm(d)68        self.head = nn.Linear(d, V, bias=False)69        self.head.weight = self.emb.weight  # tie weights70        71    def forward(self, x):72        B, N = x.shape73        mask = torch.triu(torch.ones(N, N, device=x.device), 1) * -1e974        h = self.emb(x)75        for block in self.blocks:76            h = block(h, mask)77        return self.head(self.ln(h))78    79    def count_params(self):80        return sum(p.numel() for p in self.parameters())81 82class BinaryTrainer:83    def __init__(self, model, lr=LR):84        self.model = model.to(DEVICE)85        self.opt = torch.optim.AdamW(model.parameters(), lr=lr)86        self.ctx_size = CONFIG["ctx"]87        self.buffer = deque(maxlen=self.ctx_size + 1)88        89        self.bytes_seen = 090        self.total_loss = 0.091        self.updates = 092        self.start_time = time.time()93        94    def ingest_byte(self, byte_val):95        """Absorb a single byte (0-255)"""96        self.buffer.append(byte_val)97        self.bytes_seen += 198        99        if len(self.buffer) >= UPDATE_EVERY + 1 and self.bytes_seen % UPDATE_EVERY == 0:100            self._update()101            102        if self.bytes_seen % PRINT_EVERY == 0:103            self._print_stats()104            105        # Save checkpoint every 500k bytes106        if self.bytes_seen % 500000 == 0 and self.bytes_seen > 0:107            self._save()108    109    def _update(self):110        tokens = list(self.buffer)111        x = torch.tensor(tokens[:-1], device=DEVICE, dtype=torch.long).unsqueeze(0)112        y = torch.tensor(tokens[1:], device=DEVICE, dtype=torch.long).unsqueeze(0)113        114        self.model.train()115        logits = self.model(x)116        loss = F.cross_entropy(117            logits[:, -UPDATE_EVERY:].reshape(-1, 256),118            y[:, -UPDATE_EVERY:].reshape(-1)119        )120        121        self.opt.zero_grad()122        loss.backward()123        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)124        self.opt.step()125        126        self.total_loss += loss.item()127        self.updates += 1128    129    def _print_stats(self):130        elapsed = time.time() - self.start_time131        rate = self.bytes_seen / elapsed if elapsed > 0 else 0132        avg_loss = self.total_loss / max(1, self.updates)133        mb = self.bytes_seen / 1_000_000134        135        # Bits per byte (compression metric) - log2(256)=8 is random, lower is learning136        bpb = avg_loss / math.log(2)137        138        print(f"[{elapsed:.0f}s] {mb:.2f}MB | {rate/1000:.1f} KB/s | "139              f"loss={avg_loss:.3f} | bpb={bpb:.2f} | updates={self.updates}", flush=True)140    141    def _save(self):142        avg_loss = self.total_loss / max(1, self.updates)143        mb = self.bytes_seen // 1_000_000144        ckpt = {145            "model": self.model.state_dict(),146            "bytes": self.bytes_seen,147            "loss": avg_loss,148        }149        torch.save(ckpt, f"byte_ckpt_{mb}mb.pt")150        print(f"[SAVED] {mb}MB checkpoint", flush=True)151 152def main():153    print(f"BINARY TRANSFORMER - Raw bytes learning", flush=True)154    print(f"Config: {CONFIG}", flush=True)155    print(f"Device: {DEVICE}", flush=True)156    157    model = BinaryTransformer(CONFIG)158    params = model.count_params()159    print(f"Parameters: {params:,} ({params/1e6:.1f}M)", flush=True)160    print(f"Vocab: 256 (one per byte)", flush=True)161    162    trainer = BinaryTrainer(model)163    164    print(f"Listening for raw bytes on stdin...", flush=True)165    166    # Read raw bytes from stdin167    while True:168        byte = sys.stdin.buffer.read(1)169        if not byte:170            break171        trainer.ingest_byte(byte[0])172    173    print(f"Stream ended. Total bytes: {trainer.bytes_seen:,}", flush=True)174 175if __name__ == "__main__":176    main()177