Upload e2e_full.py with huggingface_hub

e2e_full.py

@@ -0,0 +1,262 @@
+#!/usr/bin/env python3
+"""
+Self-contained E2E training benchmark: Dense vs PyLoop-sparse vs Triton-sparse.
+Includes all Triton kernels inline. Runs d_model ∈ {512, 1024, 2048}.
+"""
+
+import math, os, time, urllib.request
+import torch, torch.nn as nn, torch.nn.functional as F
+import triton, triton.language as tl
+import tiktoken
+
+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))
+
+# ═══════════ TRITON KERNELS ═══════════
+
+@triton.autotune(
+    configs=[
+        triton.Config({'BN': 32, 'BK': 64,  'BM': 32}, num_stages=3, num_warps=4),
+        triton.Config({'BN': 64, 'BK': 64,  'BM': 32}, num_stages=3, num_warps=4),
+        triton.Config({'BN': 64, 'BK': 128, 'BM': 32}, num_stages=3, num_warps=4),
+        triton.Config({'BN': 32, 'BK': 128, 'BM': 64}, num_stages=3, num_warps=4),
+        triton.Config({'BN': 64, 'BK': 64,  'BM': 64}, num_stages=4, num_warps=4),
+    ],
+    key=['M', 'd_in', 'CS'],
+)
+@triton.jit
+def _sparse_bwd_dW_kernel(
+    X_ptr, dY_ptr, dW_ptr, chunk_ids_ptr,
+    M, d_in, d_out, num_active,
+    stride_xm, stride_xk, stride_dym, stride_dyn, stride_dwn, stride_dwk,
+    CS: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, BM: tl.constexpr,
+):
+    pid0 = tl.program_id(0); pid1 = tl.program_id(1)
+    N_BLOCKS = tl.cdiv(CS, BN)
+    cli = pid0 // N_BLOCKS; nbi = pid0 % N_BLOCKS; kbi = pid1
+    if cli >= num_active: return
+    cidx = tl.load(chunk_ids_ptr + cli); cs0 = cidx * CS
+    rn = nbi * BN + tl.arange(0, BN); rk = kbi * BK + tl.arange(0, BK)
+    na = cs0 + rn; nm = rn < CS; km = rk < d_in
+    acc = tl.zeros((BN, BK), dtype=tl.float32)
+    for ms in range(0, M, BM):
+        rm = ms + tl.arange(0, BM); mm = rm < M
+        x = tl.load(X_ptr + rm[:, None]*stride_xm + rk[None, :]*stride_xk, mask=mm[:, None] & km[None, :], other=0.0)
+        dy = tl.load(dY_ptr + rm[:, None]*stride_dym + na[None, :]*stride_dyn, mask=mm[:, None] & nm[None, :], other=0.0)
+        acc = tl.dot(tl.trans(dy), x, acc=acc)
+    tl.store(dW_ptr + na[:, None]*stride_dwn + rk[None, :]*stride_dwk, acc.to(dW_ptr.dtype.element_ty), mask=nm[:, None] & km[None, :])
+
+def sparse_bwd_dW(X, dY, active, cs, d_out):
+    M, d_in = X.shape; na = active.shape[0]
+    dW = torch.zeros(d_out, d_in, device=X.device, dtype=X.dtype)
+    if na == 0: return dW
+    cids = active.to(torch.int32).contiguous()
+    grid = lambda META: (na * triton.cdiv(cs, META['BN']), triton.cdiv(d_in, META['BK']))
+    _sparse_bwd_dW_kernel[grid](X, dY, dW, cids, M, d_in, d_out, na,
+        X.stride(0), X.stride(1), dY.stride(0), dY.stride(1), dW.stride(0), dW.stride(1), CS=cs)
+    return dW
+
+@triton.jit
+def _sparse_bwd_dbias_kernel(
+    dY_ptr, dB_ptr, chunk_ids_ptr, M, d_out, num_active,
+    stride_dym, stride_dyn, CS: tl.constexpr, BM: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    cl = pid // CS; ci = pid % CS
+    if cl >= num_active: return
+    cidx = tl.load(chunk_ids_ptr + cl); ca = cidx * CS + ci
+    acc = 0.0
+    for ms in range(0, M, BM):
+        rm = ms + tl.arange(0, BM); mm = rm < M
+        acc += tl.sum(tl.load(dY_ptr + rm*stride_dym + ca*stride_dyn, mask=mm, other=0.0))
+    tl.store(dB_ptr + ca, acc.to(dB_ptr.dtype.element_ty))
+
+def sparse_bwd_dbias(dY, active, cs, d_out):
+    M = dY.shape[0]; na = active.shape[0]
+    dB = torch.zeros(d_out, device=dY.device, dtype=dY.dtype)
+    if na == 0: return dB
+    cids = active.to(torch.int32).contiguous()
+    _sparse_bwd_dbias_kernel[(na * cs,)](dY, dB, cids, M, d_out, na, dY.stride(0), dY.stride(1), CS=cs, BM=128)
+    return dB
+
+# ═══════════ AUTOGRAD ═══════════
+
+class TritonSparse(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, w, b, active, cs, sdx):
+        ctx.save_for_backward(x, w, active); ctx.has_bias = b is not None; ctx.sdx = sdx; ctx.cs = cs
+        return F.linear(x, w, b)
+    @staticmethod
+    def backward(ctx, gy):
+        x, w, active = ctx.saved_tensors; cs = ctx.cs; do, di = w.shape
+        xf = x.reshape(-1, di); gf = gy.reshape(-1, do)
+        gw = sparse_bwd_dW(xf, gf, active, cs, do)
+        gb = sparse_bwd_dbias(gf, active, cs, do) if ctx.has_bias else None
+        gx = gf @ w  # dense dX
+        return gx.reshape(x.shape), gw, gb, None, None, None
+
+class PyLoopSparse(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, w, b, active, cs, sdx):
+        ctx.save_for_backward(x, w, active); ctx.has_bias = b is not None; ctx.sdx = sdx; ctx.cs = cs
+        return F.linear(x, w, b)
+    @staticmethod
+    def backward(ctx, gy):
+        x, w, active = ctx.saved_tensors; cs = ctx.cs
+        xf = x.reshape(-1, x.shape[-1]); gf = gy.reshape(-1, gy.shape[-1])
+        gw = torch.zeros_like(w)
+        gb = torch.zeros(w.shape[0], device=w.device, dtype=w.dtype) if ctx.has_bias else None
+        gx = gf @ w
+        for c in active.tolist():
+            s, e = c*cs, (c+1)*cs
+            gw[s:e] = gf[:, s:e].t() @ xf
+            if ctx.has_bias: gb[s:e] = gf[:, s:e].sum(0)
+        return gx.reshape(x.shape), gw, gb, None, None, None
+
+# ═══════════ 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(PyLoopSparse.apply(h, self.proj.weight, self.proj.bias, self.active_chunks, self.cs, False))
+        else:
+            return self.do(TritonSparse.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)
+        a = self.do(F.softmax((q@k.transpose(-2,-1))/math.sqrt(self.hd)+self.mask[:,:,:T,:T].log(), dim=-1))
+        return self.proj((a@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))
+        return lo, F.cross_entropy(lo.view(-1,lo.size(-1)), tgt.view(-1)) if tgt is not None else None
+    def get_ffns(self): return [b.mlp for b in self.blocks]
+    def nparams(self): return sum(p.numel() for p in self.parameters())
+
+# ═══════════ RUN ═══════════
+
+STEPS = 500
+af = 0.10
+cs = 64
+
+if torch.cuda.is_available():
+    print(f"GPU: {torch.cuda.get_device_name()} | VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB")
+
+print(f"E2E training: {STEPS} steps, B={BS}, T={BLK}, active_frac={af}, chunk_size={cs}")
+print(f"{'d_model':>7} | {'Mode':>8} | {'Params':>8} | {'ms/step':>10} | {'vs Dense':>10} | {'val_loss':>10} | {'train_loss':>10}")
+print("-"*80)
+
+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)
+        npar = model.nparams()
+        opt = torch.optim.AdamW(model.parameters(), lr=5e-4)
+        ffns = model.get_ffns()
+
+        # Triton warmup (compile kernels before timing)
+        if mode == 'triton':
+            for ffn in ffns:
+                ffn.mode = mode
+                nc = ffn.proj.out_features // cs
+                k = max(1, int(af * nc))
+                ffn.active_chunks = torch.randperm(nc, device=device)[:k].sort().values
+            x, y = get_batch(train_data, torch.Generator().manual_seed(99999))
+            opt.zero_grad(); _, loss = model(x, y); loss.backward(); opt.step()
+            # Reset model
+            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()
+        last_loss = 0.0
+
+        for step in range(STEPS):
+            if mode != 'dense':
+                for ffn in ffns:
+                    ffn.mode = mode
+                    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()
+            last_loss = loss.item()
+
+            if step % 100 == 0:
+                print(f"  [{mode}] d={d} step {step}/{STEPS} loss={last_loss:.4f}")
+
+        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, last_loss, npar)
+        del model, opt; torch.cuda.empty_cache()
+
+    d_ms = results['dense'][0]
+    for mode in ['dense', 'pyloop', 'triton']:
+        ms, vl, tl_, np_ = results[mode]
+        sp = d_ms / ms
+        print(f"{d:>7} | {mode:>8} | {np_/1e6:>7.1f}M | {ms:>9.1f}ms | {sp:>9.2f}x | {vl:>9.4f} | {tl_:>9.4f}")
+    print()