OpenTransformer/binary-transformers

1#!/usr/bin/env python32"""3PURE BINARY TRANSFORMER - BITS ALL THE WAY DOWN4- Vocab = 2 (0 and 1)5- Weights = binary (-1 or +1, stored as bits)6- Activations = binary where possible7 8Uses Straight-Through Estimator (STE) for gradients.9XNOR + popcount for matmul = insanely fast on hardware.10"""11 12import sys13import math14import time15import torch16import torch.nn as nn17import torch.nn.functional as F18from collections import deque19 20DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")21 22# Config for pure binary transformer23CONFIG = {24    "d": 256,       # must be divisible by heads25    "layers": 6,26    "heads": 8,27    "vocab": 2,     # 0 and 128    "ctx": 2048,29}30 31LR = 1e-332UPDATE_EVERY = 25633PRINT_EVERY = 5000034 35# ============== BINARY LAYERS ==============36 37class BinarySign(torch.autograd.Function):38    """Binarize to -1/+1 with straight-through estimator"""39    @staticmethod40    def forward(ctx, x):41        ctx.save_for_backward(x)42        return x.sign()43    44    @staticmethod45    def backward(ctx, grad_output):46        x, = ctx.saved_tensors47        # STE: pass gradient through if |x| <= 148        grad_input = grad_output.clone()49        grad_input[x.abs() > 1] = 050        return grad_input51 52def binarize(x):53    return BinarySign.apply(x)54 55class BinaryLinear(nn.Module):56    """Linear layer with binary weights (-1/+1)"""57    def __init__(self, in_features, out_features, bias=False):58        super().__init__()59        self.in_features = in_features60        self.out_features = out_features61        62        # Real-valued weights for training, binarized during forward63        self.weight = nn.Parameter(torch.randn(out_features, in_features) * 0.1)64        if bias:65            self.bias = nn.Parameter(torch.zeros(out_features))66        else:67            self.bias = None68    69    def forward(self, x):70        # Binarize weights to -1/+171        binary_weight = binarize(self.weight)72        73        # Scale factor for better gradients (from XNOR-Net paper)74        # alpha = mean(|W|)75        alpha = self.weight.abs().mean()76        77        out = F.linear(x, binary_weight * alpha, self.bias)78        return out79 80class BinaryAttention(nn.Module):81    """Attention with binary QKV projections"""82    def __init__(self, d, h):83        super().__init__()84        self.h, self.dk = h, d // h85        self.q_proj = BinaryLinear(d, d)86        self.k_proj = BinaryLinear(d, d)87        self.v_proj = BinaryLinear(d, d)88        self.out_proj = BinaryLinear(d, d)89        90    def forward(self, x, mask=None):91        B, N, D = x.shape92        93        q = self.q_proj(x).view(B, N, self.h, self.dk).transpose(1, 2)94        k = self.k_proj(x).view(B, N, self.h, self.dk).transpose(1, 2)95        v = self.v_proj(x).view(B, N, self.h, self.dk).transpose(1, 2)96        97        # Standard attention (values stay real for now)98        att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk)99        if mask is not None:100            att = att + mask101        att = F.softmax(att, dim=-1)102        103        out = (att @ v).transpose(1, 2).reshape(B, N, D)104        return self.out_proj(out)105 106class BinaryMLP(nn.Module):107    """MLP with binary weights"""108    def __init__(self, d):109        super().__init__()110        self.fc1 = BinaryLinear(d, d * 4)111        self.fc2 = BinaryLinear(d * 4, d)112    113    def forward(self, x):114        # Binary weights, but ReLU activation (could binarize this too)115        x = F.gelu(self.fc1(x))116        return self.fc2(x)117 118class BinaryBlock(nn.Module):119    def __init__(self, d, h):120        super().__init__()121        self.ln1 = nn.LayerNorm(d)122        self.attn = BinaryAttention(d, h)123        self.ln2 = nn.LayerNorm(d)124        self.mlp = BinaryMLP(d)125        126    def forward(self, x, mask):127        x = x + self.attn(self.ln1(x), mask)128        return x + self.mlp(self.ln2(x))129 130class PureBinaryTransformer(nn.Module):131    """132    Transformer where:133    - Input vocab = 2 (bits)134    - All linear weights are binary (-1/+1)135    """136    def __init__(self, cfg):137        super().__init__()138        d, L, h = cfg["d"], cfg["layers"], cfg["heads"]139        140        # Embeddings stay real (only 2 of them anyway)141        self.emb = nn.Embedding(2, d)142        143        # Binary blocks144        self.blocks = nn.ModuleList([BinaryBlock(d, h) for _ in range(L)])145        146        self.ln = nn.LayerNorm(d)147        self.head = BinaryLinear(d, 2)  # Binary output projection too!148        149    def forward(self, x):150        B, N = x.shape151        mask = torch.triu(torch.ones(N, N, device=x.device), 1) * -1e9152        153        h = self.emb(x)154        for block in self.blocks:155            h = block(h, mask)156        157        return self.head(self.ln(h))158    159    def count_params(self):160        return sum(p.numel() for p in self.parameters())161    162    def count_binary_params(self):163        """Count params that are binarized"""164        count = 0165        for name, module in self.named_modules():166            if isinstance(module, BinaryLinear):167                count += module.weight.numel()168        return count169 170def byte_to_bits(byte_val):171    return [(byte_val >> (7 - i)) & 1 for i in range(8)]172 173class BinaryTrainer:174    def __init__(self, model, lr=LR):175        self.model = model.to(DEVICE)176        self.opt = torch.optim.AdamW(model.parameters(), lr=lr)177        self.ctx_size = CONFIG["ctx"]178        self.buffer = deque(maxlen=self.ctx_size + 1)179        180        self.bits_seen = 0181        self.bytes_seen = 0182        self.total_loss = 0.0183        self.updates = 0184        self.start_time = time.time()185        186    def ingest_byte(self, byte_val):187        bits = byte_to_bits(byte_val)188        for bit in bits:189            self.buffer.append(bit)190            self.bits_seen += 1191            192            if len(self.buffer) >= UPDATE_EVERY + 1 and self.bits_seen % UPDATE_EVERY == 0:193                self._update()194        195        self.bytes_seen += 1196        197        if self.bits_seen % PRINT_EVERY == 0:198            self._print_stats()199            200        if self.bytes_seen % 500000 == 0 and self.bytes_seen > 0:201            self._save()202    203    def _update(self):204        tokens = list(self.buffer)205        x = torch.tensor(tokens[:-1], device=DEVICE, dtype=torch.long).unsqueeze(0)206        y = torch.tensor(tokens[1:], device=DEVICE, dtype=torch.long).unsqueeze(0)207        208        self.model.train()209        logits = self.model(x)210        loss = F.cross_entropy(211            logits[:, -UPDATE_EVERY:].reshape(-1, 2),212            y[:, -UPDATE_EVERY:].reshape(-1)213        )214        215        self.opt.zero_grad()216        loss.backward()217        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)218        self.opt.step()219        220        self.total_loss += loss.item()221        self.updates += 1222    223    def _print_stats(self):224        elapsed = time.time() - self.start_time225        bytes_per_sec = self.bytes_seen / elapsed if elapsed > 0 else 0226        avg_loss = self.total_loss / max(1, self.updates)227        228        entropy = avg_loss / math.log(2)229        compression = (1.0 - entropy) * 100230        231        print(f"[{elapsed:.0f}s] {self.bytes_seen/1000:.1f}KB | {bytes_per_sec/1000:.2f} KB/s | "232              f"loss={avg_loss:.4f} | entropy={entropy:.3f} | compression={compression:.1f}%", flush=True)233    234    def _save(self):235        avg_loss = self.total_loss / max(1, self.updates)236        kb = self.bytes_seen // 1000237        ckpt = {238            "model": self.model.state_dict(),239            "bits": self.bits_seen,240            "bytes": self.bytes_seen,241            "loss": avg_loss,242        }243        torch.save(ckpt, f"/workspace/purebit_ckpt_{kb}kb.pt")244        print(f"[SAVED] purebit_ckpt_{kb}kb.pt", flush=True)245 246def main():247    print(f"PURE BINARY TRANSFORMER - BITS ALL THE WAY DOWN", flush=True)248    print(f"Config: {CONFIG}", flush=True)249    print(f"Device: {DEVICE}", flush=True)250    251    model = PureBinaryTransformer(CONFIG)252    total_params = model.count_params()253    binary_params = model.count_binary_params()254    255    print(f"Total Parameters: {total_params:,} ({total_params/1e6:.2f}M)", flush=True)256    print(f"Binary Parameters: {binary_params:,} ({binary_params/total_params*100:.1f}%)", flush=True)257    print(f"Vocab: 2 (input bits)", flush=True)258    print(f"Weights: BINARY (-1/+1)", flush=True)259    print(f"", flush=True)260    print(f"🔥 BITS IN, BITS WEIGHTS, BITS OUT 🔥", flush=True)261    262    trainer = BinaryTrainer(model)263    264    print(f"Listening for bytes...", flush=True)265    266    while True:267        byte = sys.stdin.buffer.read(1)268        if not byte:269            break270        trainer.ingest_byte(byte[0])271    272    print(f"Done. {trainer.bytes_seen:,} bytes = {trainer.bits_seen:,} bits", flush=True)273 274if __name__ == "__main__":275    main()276