Add experiments/final_showdown.py

experiments/final_showdown.py

@@ -0,0 +1,172 @@
+#!/usr/bin/env python3
+"""
+FINAL SHOWDOWN: Standard depth vs Ultra-heavy mechanisms
+Question: At equal compute budget, does any heavy approach beat just adding layers?
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import time
+import math
+
+DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.backends.cuda.matmul.allow_tf32 = True
+VOCAB = 128256
+
+def alibi_bias(n_heads, n_tokens):
+    def slopes(n):
+        start = 2 ** (-2 ** -(math.log2(n) - 3))
+        return [start * (start ** i) for i in range(n)]
+    s = slopes(n_heads) if math.log2(n_heads).is_integer() else slopes(2 ** math.floor(math.log2(n_heads)))[:n_heads]
+    s = torch.tensor(s, device=DEV).view(1, n_heads, 1, 1)
+    i = torch.arange(n_tokens, device=DEV).view(1, 1, n_tokens, 1)
+    j = torch.arange(n_tokens, device=DEV).view(1, 1, 1, n_tokens)
+    return -s * (j - i).clamp_min(0).float()
+
+def causal_mask(n):
+    return torch.triu(torch.full((1, 1, n, n), float("-inf"), device=DEV), 1)
+
+
+class StandardAttn(nn.Module):
+    def __init__(self, d, h):
+        super().__init__()
+        self.h, self.dk = h, d // h
+        self.qkv = nn.Linear(d, 3*d, bias=False)
+        self.proj = nn.Linear(d, d, bias=False)
+    
+    def forward(self, x, mask=None):
+        B, N, _ = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.h, self.dk).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk) + alibi_bias(self.h, N)
+        if mask is not None: att = att + mask
+        return self.proj((att.softmax(-1) @ v).transpose(1, 2).reshape(B, N, -1))
+
+
+class DoubleAttn(nn.Module):
+    """Simplest heavy: two sequential attention ops"""
+    def __init__(self, d, h):
+        super().__init__()
+        self.attn1 = StandardAttn(d, h)
+        self.attn2 = StandardAttn(d, h)
+        self.gate = nn.Linear(d * 2, d)
+    
+    def forward(self, x, mask=None):
+        o1 = self.attn1(x, mask)
+        o2 = self.attn2(x + o1, mask)
+        return self.gate(torch.cat([o1, o2], dim=-1))
+
+
+class RecurrentAttn(nn.Module):
+    """Same attention applied k times"""
+    def __init__(self, d, h, k=4):
+        super().__init__()
+        self.attn = StandardAttn(d, h)
+        self.depth_emb = nn.Embedding(k, d)
+        self.k = k
+    
+    def forward(self, x, mask=None):
+        for i in range(self.k):
+            x = x + self.attn(x + self.depth_emb.weight[i], mask)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, d, h, mode="standard"):
+        super().__init__()
+        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
+        if mode == "standard":
+            self.attn = StandardAttn(d, h)
+        elif mode == "double":
+            self.attn = DoubleAttn(d, h)
+        elif mode == "recurrent":
+            self.attn = RecurrentAttn(d, h, k=4)
+        self.ff = nn.Sequential(nn.Linear(d, 4*d), nn.GELU(), nn.Linear(4*d, d))
+    
+    def forward(self, x, mask=None):
+        x = x + self.attn(self.ln1(x), mask)
+        return x + self.ff(self.ln2(x))
+
+
+class Model(nn.Module):
+    def __init__(self, d, layers, h, mode="standard"):
+        super().__init__()
+        self.emb = nn.Embedding(VOCAB, d)
+        self.blocks = nn.ModuleList([Block(d, h, mode) for _ in range(layers)])
+        self.ln = nn.LayerNorm(d)
+        self.head = nn.Linear(d, VOCAB, bias=False)
+        self.head.weight = self.emb.weight
+    
+    def forward(self, x, mask=None):
+        x = self.emb(x)
+        for b in self.blocks: x = b(x, mask)
+        return self.head(self.ln(x))
+    
+    def count_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+
+def train(model, steps, batch, seq):
+    opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
+    mask = causal_mask(seq - 1)
+    losses, times = [], []
+    
+    for step in range(steps):
+        ids = torch.randint(0, VOCAB, (batch, seq), device=DEV)
+        start = time.time()
+        opt.zero_grad()
+        loss = F.cross_entropy(model(ids[:, :-1], mask).view(-1, VOCAB), ids[:, 1:].reshape(-1))
+        loss.backward()
+        opt.step()
+        times.append(time.time() - start)
+        losses.append(loss.item())
+        
+        if step % 50 == 0 or step == steps - 1:
+            tok_s = batch * seq / times[-1]
+            print(f"Step {step:3d} | Loss {loss.item():.4f} | {tok_s:.0f} tok/s")
+    
+    return sum(losses[-20:]) / 20, batch * seq / (sum(times[-20:]) / 20)
+
+
+def main():
+    print(f"Device: {DEV}")
+    if torch.cuda.is_available():
+        print(f"GPU: {torch.cuda.get_device_name()}")
+    
+    d, h, batch, seq = 256, 8, 16, 128
+    
+    configs = [
+        # (name, layers, mode, target_steps)
+        ("Standard-4L", 4, "standard", 500),
+        ("Standard-8L", 8, "standard", 250),  # ~2x slower, so half steps
+        ("Standard-16L", 16, "standard", 125),  # ~4x slower
+        ("Double-4L", 4, "double", 250),  # ~2x slower
+        ("Recurrent-4L", 4, "recurrent", 125),  # ~4x slower (k=4 iterations)
+    ]
+    
+    results = []
+    for name, layers, mode, steps in configs:
+        print(f"\n{'='*60}")
+        print(f"{name}")
+        print(f"{'='*60}")
+        
+        model = Model(d, layers, h, mode).to(DEV)
+        params = model.count_params()
+        print(f"Parameters: {params:,}")
+        
+        avg_loss, avg_toks = train(model, steps, batch, seq)
+        results.append((name, avg_loss, avg_toks, params, steps))
+        del model
+        torch.cuda.empty_cache()
+    
+    print(f"\n{'='*60}")
+    print("FINAL RESULTS (roughly compute-matched)")
+    print(f"{'='*60}")
+    for name, loss, toks, params, steps in results:
+        total_tok = steps * batch * seq
+        print(f"{name:15s} | Loss {loss:.4f} | {toks:.0f} tok/s | {params/1e6:.1f}M | {total_tok/1e6:.1f}M tok trained")
+
+
+if __name__ == "__main__":
+    main()