Upload e2e_bench.py with huggingface_hub

e2e_bench.py

@@ -0,0 +1,156 @@
+#!/usr/bin/env python3
+"""End-to-end training benchmark: Dense vs PyLoop vs Triton sparse backward."""
+
+import math, os, time, urllib.request
+import torch, torch.nn as nn, torch.nn.functional as F
+import tiktoken
+
+# Import our Triton kernels from the module
+from triton_sparse import (
+    TritonChunkedSparseLinear, PythonLoopSparseLinear,
+    sparse_bwd_dW, sparse_bwd_dX, sparse_bwd_dbias
+)
+
+device = 'cuda'
+BS, BLK = 8, 256
+
+# Data
+if not os.path.exists('input.txt'):
+    urllib.request.urlretrieve('https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt', 'input.txt')
+enc = tiktoken.get_encoding('gpt2')
+tokens = torch.tensor(enc.encode(open('input.txt').read()), dtype=torch.long)
+train_data = tokens[:int(0.9*len(tokens))]
+val_data = tokens[int(0.9*len(tokens)):]
+V = enc.n_vocab
+
+def get_batch(data, gen=None):
+    ix = torch.randint(len(data)-BLK-1, (BS,), generator=gen)
+    return (torch.stack([data[i:i+BLK] for i in ix]).to(device),
+            torch.stack([data[i+1:i+BLK+1] for i in ix]).to(device))
+
+# Model
+class SparseFFN(nn.Module):
+    def __init__(self, d, cs=64):
+        super().__init__()
+        self.fc = nn.Linear(d, 4*d)
+        self.proj = nn.Linear(4*d, d)
+        self.do = nn.Dropout(0.1)
+        self.cs = cs
+        self.mode = 'dense'
+        self.active_chunks = None
+
+    def forward(self, x):
+        h = F.gelu(self.fc(x))
+        if self.mode == 'dense' or self.active_chunks is None:
+            return self.do(self.proj(h))
+        elif self.mode == 'pyloop':
+            return self.do(PythonLoopSparseLinear.apply(
+                h, self.proj.weight, self.proj.bias, self.active_chunks, self.cs, False))
+        else:  # triton
+            return self.do(TritonChunkedSparseLinear.apply(
+                h, self.proj.weight, self.proj.bias, self.active_chunks, self.cs, False))
+
+class Attn(nn.Module):
+    def __init__(self, d, nh, bs):
+        super().__init__()
+        self.nh, self.hd = nh, d//nh
+        self.qkv = nn.Linear(d, 3*d)
+        self.proj = nn.Linear(d, d)
+        self.do = nn.Dropout(0.1)
+        self.register_buffer('mask', torch.tril(torch.ones(bs,bs)).view(1,1,bs,bs))
+
+    def forward(self, x):
+        B,T,C = x.shape
+        q,k,v = self.qkv(x).split(C,2)
+        q = q.view(B,T,self.nh,self.hd).transpose(1,2)
+        k = k.view(B,T,self.nh,self.hd).transpose(1,2)
+        v = v.view(B,T,self.nh,self.hd).transpose(1,2)
+        att = (q @ k.transpose(-2,-1)) / math.sqrt(self.hd)
+        att = att.masked_fill(self.mask[:,:,:T,:T]==0, float('-inf'))
+        att = self.do(F.softmax(att, dim=-1))
+        return self.proj((att @ v).transpose(1,2).contiguous().view(B,T,C))
+
+class Block(nn.Module):
+    def __init__(self, d, nh, bs):
+        super().__init__()
+        self.ln1=nn.LayerNorm(d); self.attn=Attn(d,nh,bs)
+        self.ln2=nn.LayerNorm(d); self.mlp=SparseFFN(d)
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        return x + self.mlp(self.ln2(x))
+
+class GPT(nn.Module):
+    def __init__(self, d, nl, nh, bs):
+        super().__init__()
+        self.te=nn.Embedding(V,d); self.pe=nn.Embedding(bs,d)
+        self.blocks=nn.ModuleList([Block(d,nh,bs) for _ in range(nl)])
+        self.ln=nn.LayerNorm(d); self.head=nn.Linear(d,V)
+
+    def forward(self, idx, tgt=None):
+        x = self.te(idx)+self.pe(torch.arange(idx.shape[1],device=idx.device))[None]
+        for b in self.blocks: x = b(x)
+        lo = self.head(self.ln(x))
+        loss = F.cross_entropy(lo.view(-1,lo.size(-1)), tgt.view(-1)) if tgt is not None else None
+        return lo, loss
+
+    def get_ffns(self):
+        return [b.mlp for b in self.blocks]
+
+# Run
+STEPS = 100
+af = 0.10
+cs = 64
+
+print(f"End-to-end training: {STEPS} steps, B={BS}, T={BLK}, active_frac={af}")
+print(f"{'d_model':>7} | {'Mode':>8} | {'ms/step':>10} | {'vs Dense':>10} | {'val_loss':>10}")
+print("-"*60)
+
+for d in [512, 1024, 2048]:
+    nh = 8; nl = 6
+    results = {}
+
+    for mode in ['dense', 'pyloop', 'triton']:
+        torch.manual_seed(42)
+        model = GPT(d, nl, nh, BLK).to(device)
+        opt = torch.optim.AdamW(model.parameters(), lr=5e-4)
+        ffns = model.get_ffns()
+
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+
+        for step in range(STEPS):
+            if mode != 'dense':
+                for ffn in ffns:
+                    ffn.mode = mode
+                    # proj: Linear(4d, d) -> weight shape (d, 4d), out_features=d
+                    nc = ffn.proj.out_features // cs
+                    k = max(1, int(af * nc))
+                    ffn.active_chunks = torch.randperm(nc, device=device)[:k].sort().values
+            else:
+                for ffn in ffns:
+                    ffn.mode = 'dense'; ffn.active_chunks = None
+
+            x, y = get_batch(train_data, torch.Generator().manual_seed(step))
+            opt.zero_grad()
+            _, loss = model(x, y)
+            loss.backward()
+            opt.step()
+
+        torch.cuda.synchronize()
+        ms = 1000 * (time.perf_counter() - t0) / STEPS
+
+        # Eval
+        model.eval()
+        for ffn in ffns: ffn.mode = 'dense'; ffn.active_chunks = None
+        with torch.no_grad():
+            vl = sum(model(*get_batch(val_data, torch.Generator().manual_seed(9999+i)))[1].item() for i in range(20))/20
+
+        results[mode] = (ms, vl)
+        del model; torch.cuda.empty_cache()
+
+    d_ms = results['dense'][0]
+    for mode in ['dense', 'pyloop', 'triton']:
+        ms, vl = results[mode]
+        sp = d_ms / ms
+        print(f"{d:>7} | {mode:>8} | {ms:>9.1f}ms | {sp:>9.2f}x | {vl:>9.4f}")
+    print()