Upload ablations.py with huggingface_hub

ablations.py

@@ -1 +1,824 @@
-See file content
+#!/usr/bin/env python3
+"""
+Sparse Transformer: Definitive Ablation Suite
+
+Builds on v18_fast_knn_triton.py. Addresses all three structural gaps
+identified in the critique:
+
+  1. PHANTOM MOMENTUM ABLATION
+     - "phantom": standard Adam — inactive chunks' moments decay on zero grad (default)
+     - "frozen": inactive chunks' Adam state (m, v) is completely frozen
+     Compare across all schedulers to isolate whether convergence is driven
+     by the chunking algorithm or by phantom momentum acting as regularization.
+
+  2. COMPUTE-MATCHED BASELINES
+     - Dense at same steps (standard comparison)
+     - Dense at fewer steps matching sparse FLOPs
+     - Natively smaller dense model matching sparse active capacity
+
+  3. UNIFIED HARDWARE
+     Everything on CUDA (A10G). Single hardware stack.
+
+Plus: KNN vs EMA vs Random vs Oracle predictor comparison with proper
+oracle overlap measurement.
+
+Run:
+  python ablations.py --device cuda --steps 1000 --n_embd 1024 --experiment all
+  python ablations.py --device cuda --experiment phantom_momentum
+  python ablations.py --device cuda --experiment compute_matched
+  python ablations.py --device cuda --experiment predictor_accuracy
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import random
+import sys
+import time
+from collections import defaultdict
+from typing import Dict, List, Literal, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    import triton
+    import triton.language as tl
+    HAS_TRITON = True
+except ImportError:
+    HAS_TRITON = False
+
+try:
+    import tiktoken
+    HAS_TIKTOKEN = True
+except ImportError:
+    HAS_TIKTOKEN = False
+
+# ═══════════════════════════════════════════════════════════════
+# TRITON KERNELS (from v18_triton, no autotune, block_ptr)
+# ═══════════════════════════════════════════════════════════════
+
+if HAS_TRITON:
+    @triton.jit
+    def _sparse_bwd_dW_db_kernel(
+        X_ptr, dY_ptr, dW_ptr, dB_ptr, chunk_ids_ptr,
+        M: tl.constexpr, d_in: tl.constexpr, d_out: tl.constexpr,
+        num_active: tl.constexpr,
+        stride_xm: tl.constexpr, stride_xk: tl.constexpr,
+        stride_dym: tl.constexpr, stride_dyn: tl.constexpr,
+        stride_dwn: tl.constexpr, stride_dwk: tl.constexpr,
+        HAS_BIAS: tl.constexpr,
+        CS: tl.constexpr, BK: tl.constexpr, BM: tl.constexpr,
+    ):
+        cli = tl.program_id(0)
+        kbi = tl.program_id(1)
+        cidx = tl.load(chunk_ids_ptr + cli)
+        cs0 = cidx * CS
+        ko = kbi * BK
+
+        dy_bp = tl.make_block_ptr(dY_ptr, (d_out, M), (stride_dyn, stride_dym),
+                                   (cs0, 0), (CS, BM), (1, 0))
+        x_bp = tl.make_block_ptr(X_ptr, (M, d_in), (stride_xm, stride_xk),
+                                  (0, ko), (BM, BK), (1, 0))
+
+        acc = tl.zeros((CS, BK), dtype=tl.float32)
+        do_bias = HAS_BIAS and (kbi == 0)
+        acc_b = tl.zeros((CS,), dtype=tl.float32)
+
+        for _ in range(0, M, BM):
+            dy_t = tl.load(dy_bp, boundary_check=(0, 1))
+            x = tl.load(x_bp, boundary_check=(0, 1))
+            acc = tl.dot(dy_t, x, acc=acc)
+            if do_bias:
+                acc_b += tl.sum(dy_t, axis=1)
+            dy_bp = tl.advance(dy_bp, (0, BM))
+            x_bp = tl.advance(x_bp, (BM, 0))
+
+        dw_bp = tl.make_block_ptr(dW_ptr, (d_out, d_in), (stride_dwn, stride_dwk),
+                                   (cs0, ko), (CS, BK), (1, 0))
+        tl.store(dw_bp, acc.to(dW_ptr.dtype.element_ty), boundary_check=(0, 1))
+
+        if do_bias:
+            rn = cs0 + tl.arange(0, CS)
+            tl.store(dB_ptr + rn, acc_b.to(dB_ptr.dtype.element_ty), mask=rn < d_out)
+
+    @triton.jit
+    def _sparse_bwd_dX_kernel(
+        dY_ptr, W_ptr, dX_ptr, chunk_ids_ptr,
+        M: tl.constexpr, d_in: tl.constexpr, d_out: tl.constexpr,
+        num_active: tl.constexpr,
+        stride_dym: tl.constexpr, stride_dyn: tl.constexpr,
+        stride_wn: tl.constexpr, stride_wk: tl.constexpr,
+        stride_dxm: tl.constexpr, stride_dxk: tl.constexpr,
+        CS: tl.constexpr, BM: tl.constexpr, BK: tl.constexpr,
+    ):
+        pm = tl.program_id(0)
+        pk = tl.program_id(1)
+        mo = pm * BM
+        ko = pk * BK
+        acc = tl.zeros((BM, BK), dtype=tl.float32)
+        for i in range(0, num_active):
+            cidx = tl.load(chunk_ids_ptr + i)
+            cs0 = cidx * CS
+            dy_bp = tl.make_block_ptr(dY_ptr, (M, d_out), (stride_dym, stride_dyn),
+                                       (mo, cs0), (BM, CS), (1, 0))
+            w_bp = tl.make_block_ptr(W_ptr, (d_out, d_in), (stride_wn, stride_wk),
+                                      (cs0, ko), (CS, BK), (1, 0))
+            dy = tl.load(dy_bp, boundary_check=(0, 1))
+            w = tl.load(w_bp, boundary_check=(0, 1))
+            acc = tl.dot(dy, w, acc=acc)
+        dx_bp = tl.make_block_ptr(dX_ptr, (M, d_in), (stride_dxm, stride_dxk),
+                                   (mo, ko), (BM, BK), (1, 0))
+        tl.store(dx_bp, acc.to(dX_ptr.dtype.element_ty), boundary_check=(0, 1))
+
+
+def triton_bwd_dW_db(xf, gyf, active, cs, d_out, has_bias):
+    M, d_in = xf.shape
+    na = active.numel()
+    dW = torch.zeros(d_out, d_in, device=xf.device, dtype=xf.dtype)
+    dB = torch.zeros(d_out, device=xf.device, dtype=xf.dtype) if has_bias else None
+    if na == 0: return dW, dB
+    cids = active.to(torch.int32).contiguous()
+    BK, BM = 64, 64
+    _sparse_bwd_dW_db_kernel[(na, triton.cdiv(d_in, BK))](
+        xf, gyf, dW, dB if has_bias else dW, cids,
+        M, d_in, d_out, na,
+        xf.stride(0), xf.stride(1), gyf.stride(0), gyf.stride(1),
+        dW.stride(0), dW.stride(1),
+        HAS_BIAS=has_bias, CS=cs, BK=BK, BM=BM, num_warps=4)
+    return dW, dB
+
+def triton_bwd_dX(gyf, w, active, cs, M, d_in):
+    na = active.numel()
+    d_out = gyf.shape[1]
+    dX = torch.zeros(M, d_in, device=gyf.device, dtype=gyf.dtype)
+    if na == 0: return dX
+    cids = active.to(torch.int32).contiguous()
+    BM, BK = 64, 64
+    _sparse_bwd_dX_kernel[(triton.cdiv(M, BM), triton.cdiv(d_in, BK))](
+        gyf, w, dX, cids,
+        M, d_in, d_out, na,
+        gyf.stride(0), gyf.stride(1), w.stride(0), w.stride(1),
+        dX.stride(0), dX.stride(1),
+        CS=cs, BM=BM, BK=BK, num_warps=4)
+    return dX
+
+# ═══════════════════════════════════════════════════════════════
+# AUTOGRAD
+# ═══════════════════════════════════════════════════════════════
+
+class TritonSparseLinearFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, w, b, active, cs, sparse_dx):
+        ctx.save_for_backward(x, w, active)
+        ctx.has_bias = b is not None
+        ctx.sparse_dx = sparse_dx
+        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).contiguous()
+        gf = gy.reshape(-1, do).contiguous()
+        M = xf.shape[0]
+        gw, gb = triton_bwd_dW_db(xf, gf, active, cs, do, ctx.has_bias)
+        gx = triton_bwd_dX(gf, w.contiguous(), active, cs, M, di) if ctx.sparse_dx else gf @ w
+        return gx.reshape(x.shape), gw, gb, None, None, None
+
+class PyLoopSparseLinearFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, w, b, active, cs, sparse_dx):
+        ctx.save_for_backward(x, w, active)
+        ctx.has_bias = b is not None
+        ctx.sparse_dx = sparse_dx
+        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 = torch.zeros_like(xf) if ctx.sparse_dx else gf @ w
+        for c in active.tolist():
+            s, e = c * cs, (c+1) * cs
+            sl = gf[:, s:e]
+            gw[s:e] = sl.t() @ xf
+            if gb is not None: gb[s:e] = sl.sum(0)
+            if ctx.sparse_dx: gx += sl @ w[s:e]
+        return gx.reshape(x.shape), gw, gb, None, None, None
+
+# ═══════════════════════════════════════════════════════════════
+# MODEL
+# ═══════════════════════════════════════════════════════════════
+
+class SparseLinear(nn.Linear):
+    def __init__(self, inf, outf, bias=True):
+        super().__init__(inf, outf, bias=bias)
+        self.sparse_enabled = False
+        self.sparse_dx = False
+        self.active_chunks = None
+        self.chunk_size = 64
+        self.backend = "triton"  # "triton" or "torch"
+
+    def forward(self, x):
+        if not self.sparse_enabled or self.active_chunks is None:
+            return F.linear(x, self.weight, self.bias)
+        fn = TritonSparseLinearFn if (self.backend == "triton" and HAS_TRITON) else PyLoopSparseLinearFn
+        return fn.apply(x, self.weight, self.bias, self.active_chunks, self.chunk_size, self.sparse_dx)
+
+class Attn(nn.Module):
+    def __init__(self, d, nh, bs, do):
+        super().__init__()
+        self.nh, self.hd = nh, d // nh
+        self.c_attn = SparseLinear(d, 3*d)
+        self.c_proj = SparseLinear(d, d)
+        self.drop = nn.Dropout(do)
+        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.c_attn(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 = (q @ k.transpose(-2,-1)) / math.sqrt(self.hd)
+        a = a.masked_fill(self.mask[:,:,:T,:T]==0, float("-inf"))
+        a = self.drop(F.softmax(a, dim=-1))
+        return self.c_proj((a @ v).transpose(1,2).contiguous().view(B,T,C))
+
+class FFN(nn.Module):
+    def __init__(self, d, do, ffn_mult=4):
+        super().__init__()
+        self.c_fc = SparseLinear(d, ffn_mult * d)
+        self.c_proj = SparseLinear(ffn_mult * d, d)
+        self.drop = nn.Dropout(do)
+    def forward(self, x):
+        return self.drop(self.c_proj(F.gelu(self.c_fc(x))))
+
+class Block(nn.Module):
+    def __init__(self, d, nh, bs, do, ffn_mult=4):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(d); self.attn = Attn(d, nh, bs, do)
+        self.ln2 = nn.LayerNorm(d); self.mlp = FFN(d, do, ffn_mult)
+    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, V, bs, nl, nh, d, do, ffn_mult=4):
+        super().__init__()
+        self.te = nn.Embedding(V, d); self.pe = nn.Embedding(bs, d)
+        self.blocks = nn.Sequential(*[Block(d, nh, bs, do, ffn_mult) for _ in range(nl)])
+        self.ln = nn.LayerNorm(d); self.head = nn.Linear(d, V)
+    def forward(self, idx, tgt=None):
+        B,T = idx.shape
+        x = self.te(idx) + self.pe(torch.arange(T, device=idx.device))[None]
+        lo = self.head(self.ln(self.blocks(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 nparams(self): return sum(p.numel() for p in self.parameters())
+
+def get_sparse_linears(m): return [x for x in m.modules() if isinstance(x, SparseLinear)]
+
+# ═══════════════════════════════════════════════════════════════
+# DATA
+# ═══════════════════════════════════════════════════════════════
+
+class Corpus:
+    """Uses tiktoken GPT-2 BPE on Tiny Shakespeare if available, else char-level synthetic."""
+    _inst = None
+    @classmethod
+    def get(cls, bs, dev):
+        if cls._inst is None or cls._inst.block_size != bs:
+            cls._inst = cls(bs, dev)
+        return cls._inst
+
+    def __init__(self, block_size, device):
+        self.block_size, self.device = block_size, device
+        import urllib.request
+        p = "input.txt"
+        if not os.path.exists(p):
+            urllib.request.urlretrieve("https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt", p)
+        text = open(p).read()
+        if HAS_TIKTOKEN:
+            enc = tiktoken.get_encoding("gpt2")
+            tokens = enc.encode(text)
+            self.vocab_size = enc.n_vocab
+        else:
+            chars = sorted(set(text))
+            stoi = {c:i for i,c in enumerate(chars)}
+            tokens = [stoi[c] for c in text]
+            self.vocab_size = len(chars)
+        data = torch.tensor(tokens, dtype=torch.long)
+        si = int(0.9 * len(data))
+        self.train_data, self.val_data = data[:si], data[si:]
+        print(f"Corpus: V={self.vocab_size}, train={len(self.train_data):,}, val={len(self.val_data):,}")
+
+    def get_batch(self, split, bs, gen=None):
+        d = self.train_data if split == "train" else self.val_data
+        ix = torch.randint(len(d)-self.block_size-1, (bs,), generator=gen)
+        x = torch.stack([d[i:i+self.block_size] for i in ix])
+        y = torch.stack([d[i+1:i+self.block_size+1] for i in ix])
+        return x.to(self.device), y.to(self.device)
+
+def make_gen(s):
+    g = torch.Generator(device="cpu"); g.manual_seed(s); return g
+
+# ═══════════════════════════════════════════════════════════════
+# SCHEDULER (from v18, with KNN)
+# ═══════════════════════════════════════════════════════════════
+
+class ChunkScheduler:
+    def __init__(self, model, policy, frac, cs, dev, beta=0.95, knn_k=3,
+                 sim_hist=128, min_sim_hist=8):
+        self.policy, self.frac, self.cs, self.dev = policy, frac, cs, dev
+        self.beta, self.knn_k = beta, knn_k
+        self.sim_hist, self.min_sim_hist = sim_hist, min_sim_hist
+        self.linears = get_sparse_linears(model)
+        self.m2ids, self.m2loc = {}, {}
+        off = 0
+        for m in self.linears:
+            m.chunk_size = cs
+            nc = m.out_features // cs
+            assert m.out_features % cs == 0
+            self.m2ids[m] = torch.arange(off, off+nc, device=dev)
+            self.m2loc[m] = torch.arange(nc, device=dev)
+            off += nc
+        self.nc = off
+        self.ema = torch.zeros(self.nc, device=dev)
+        self.active = torch.zeros(self.nc, dtype=torch.bool, device=dev)
+        self.mass_history = []
+        self.similarity = None
+        self.scores = torch.zeros(self.nc, device=dev)
+
+    def get_frac(self, step, wu, an):
+        if step < wu: return 1.0
+        if an > 0 and step < wu + an:
+            p = (step - wu) / an
+            return self.frac + (1-self.frac) * 0.5 * (1 + math.cos(math.pi * p))
+        return self.frac
+
+    def choose(self, step, wu, an):
+        f = self.get_frac(step, wu, an)
+        if f >= 0.999:
+            self.active.fill_(True)
+            self._install(); return
+        k = max(1, int(f * self.nc))
+        self.active.fill_(False)
+        if self.policy == "random":
+            idx = torch.randperm(self.nc, device=self.dev)[:k]
+        elif self.policy == "ema":
+            idx = torch.topk(self.ema + 1e-9*torch.rand_like(self.ema), k=k).indices
+        elif self.policy == "knn":
+            base = self.scores if self.scores.sum() > 1e-12 else self.ema
+            idx = torch.topk(base + 1e-9*torch.rand_like(base), k=k).indices
+        else:
+            raise ValueError(self.policy)
+        self.active[idx] = True
+        self._install()
+
+    def _install(self):
+        for m, gids in self.m2ids.items():
+            m.active_chunks = self.m2loc[m][self.active[gids]]
+
+    @torch.no_grad()
+    def update(self, step, wu):
+        cur = torch.zeros_like(self.ema)
+        for m, ids in self.m2ids.items():
+            if m.weight.grad is None: continue
+            s = m.weight.grad.square().view(len(ids), self.cs, -1).sum((1,2))
+            if m.bias is not None and m.bias.grad is not None:
+                s += m.bias.grad.square().view(len(ids), self.cs).sum(1)
+            cur[ids] = torch.sqrt(s + 1e-30)
+        obs = self.active
+        new = obs & (self.ema == 0)
+        old = obs & ~new
+        self.ema[new] = cur[new]
+        self.ema[old] = self.beta*self.ema[old] + (1-self.beta)*cur[old]
+        # KNN similarity building during warmup
+        if step < wu:
+            self.mass_history.append(cur.clone())
+            if len(self.mass_history) > self.sim_hist:
+                self.mass_history = self.mass_history[-self.sim_hist:]
+            if len(self.mass_history) >= self.min_sim_hist:
+                self.similarity = self._build_sim()
+        if self.policy == "knn":
+            self.scores = self._knn_scores(self.active, cur)
+        else:
+            self.scores = self.ema.clone()
+        return cur
+
+    def _build_sim(self):
+        H = torch.stack(self.mass_history)
+        H = (H - H.mean(0, keepdim=True)) / (H.std(0, keepdim=True) + 1e-6)
+        S = torch.clamp((H.T @ H) / max(1, H.shape[0]-1), min=0)
+        S.fill_diagonal_(0)
+        ok = torch.zeros_like(S, dtype=torch.bool)
+        for _, ids in self.m2ids.items():
+            ok[ids[:,None], ids[None,:]] = True
+        return torch.where(ok, S, torch.zeros_like(S))
+
+    def _knn_scores(self, active_mask, cur):
+        if self.similarity is None: return self.ema.clone()
+        sc = self.ema.clone()
+        sc[active_mask] = cur[active_mask]
+        aidx = active_mask.nonzero(as_tuple=False).flatten()
+        iidx = (~active_mask).nonzero(as_tuple=False).flatten()
+        if aidx.numel() == 0: return sc
+        S = self.similarity
+        for i in iidx.tolist():
+            w = S[i, aidx]
+            if w.sum() <= 1e-12: continue
+            kk = min(self.knn_k, w.numel())
+            top = torch.topk(w, k=kk)
+            sc[i] = (top.values * cur[aidx[top.indices]]).sum() / (top.values.sum() + 1e-12)
+        return sc
+
+    @torch.no_grad()
+    def oracle_scores(self):
+        """Compute dense gradient magnitudes per chunk (requires dense grads already computed)."""
+        sc = torch.zeros(self.nc, device=self.dev)
+        for m, ids in self.m2ids.items():
+            if m.weight.grad is None: continue
+            s = m.weight.grad.square().view(len(ids), self.cs, -1).sum((1,2))
+            if m.bias is not None and m.bias.grad is not None:
+                s += m.bias.grad.square().view(len(ids), self.cs).sum(1)
+            sc[ids] = torch.sqrt(s + 1e-30)
+        return sc
+
+    def measure_overlap(self, k):
+        """Jaccard and recall of current active vs oracle top-k."""
+        oracle = set(torch.topk(self.oracle_scores(), k=k).indices.tolist())
+        pred = set(self.active.nonzero(as_tuple=True)[0].tolist())
+        if not oracle or not pred: return 0., 0.
+        inter = oracle & pred
+        return len(inter)/len(oracle|pred), len(inter)/len(oracle)
+
+# ══════════════════════��════════════════════════════════════════
+# CHUNKED ADAM WITH PHANTOM/FROZEN MODES
+# ═══════════════════════════════════════════════════════════════
+
+class ChunkedAdam:
+    """
+    Adam with two modes for inactive chunks:
+      phantom: standard — m,v decay even on zero grad (default, original behavior)
+      frozen:  m,v state completely frozen for inactive chunks
+    """
+    def __init__(self, model, lr=3e-4, cs=64, momentum_mode="phantom"):
+        self.model, self.lr, self.cs = model, lr, cs
+        self.momentum_mode = momentum_mode  # "phantom" or "frozen"
+        self.state = {}
+        self.p2m = {}
+        for m in get_sparse_linears(model):
+            if m.weight is not None: self.p2m[m.weight] = m
+            if m.bias is not None: self.p2m[m.bias] = m
+
+    def zero_grad(self):
+        for p in self.model.parameters(): p.grad = None
+
+    @torch.no_grad()
+    def step(self):
+        for p in self.model.parameters():
+            if p.grad is None: continue
+            if p not in self.state:
+                self.state[p] = {"m": torch.zeros_like(p), "v": torch.zeros_like(p)}
+            m, v = self.state[p]["m"], self.state[p]["v"]
+            sm = self.p2m.get(p)
+            ac = getattr(sm, 'active_chunks', None) if sm else None
+
+            if ac is None:
+                # Dense parameter (LN, embeddings, lm_head) — always full update
+                m.mul_(0.9).add_(p.grad, alpha=0.1)
+                v.mul_(0.999).addcmul_(p.grad, p.grad, value=0.001)
+                p.sub_(m / (torch.sqrt(v) + 1e-8), alpha=self.lr)
+            else:
+                if self.momentum_mode == "phantom":
+                    # PHANTOM: update ALL chunks' moments, but only active get real gradients.
+                    # Inactive chunks see grad=0, so m decays and v decays.
+                    # This is the original behavior.
+                    m.mul_(0.9).add_(p.grad, alpha=0.1)
+                    v.mul_(0.999).addcmul_(p.grad, p.grad, value=0.001)
+                    # But only update weights for active chunks
+                    for c in ac.tolist():
+                        s, e = c*self.cs, (c+1)*self.cs
+                        p.data[s:e].sub_(m[s:e] / (torch.sqrt(v[s:e]) + 1e-8), alpha=self.lr)
+                elif self.momentum_mode == "frozen":
+                    # FROZEN: only touch m,v,p for active chunks. Inactive state is untouched.
+                    for c in ac.tolist():
+                        s, e = c*self.cs, (c+1)*self.cs
+                        g = p.grad[s:e]
+                        m[s:e].mul_(0.9).add_(g, alpha=0.1)
+                        v[s:e].mul_(0.999).addcmul_(g, g, value=0.001)
+                        p.data[s:e].sub_(m[s:e] / (torch.sqrt(v[s:e]) + 1e-8), alpha=self.lr)
+
+# ═══════════════════════════════════════════════════════════════
+# EVALUATION
+# ═══════════════════════════════════════════════════════════════
+
+@torch.no_grad()
+def evaluate(model, corpus, bs, n=20, seed=9999):
+    model.eval()
+    losses = []
+    for i in range(n):
+        _, l = model(*corpus.get_batch("val", bs, make_gen(seed+i)))
+        losses.append(l.item())
+    model.train()
+    avg = sum(losses)/len(losses)
+    return avg, math.exp(min(avg, 20))
+
+# ═══════════════════════════════════════════════════════════════
+# SINGLE TRAINING RUN
+# ═══════════════════════════════════════════════════════════════
+
+def run(policy, bwd_mode, steps, bs, block_size, nl, nh, d, cs,
+        active_frac, wu, an, lr, device, seed, backend="triton",
+        momentum_mode="phantom", ffn_mult=4,
+        measure_oracle=False, oracle_interval=50):
+    """Run one training config. Returns dict of results."""
+    torch.manual_seed(seed)
+    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
+    random.seed(seed)
+
+    corpus = Corpus.get(block_size, device)
+    model = GPT(corpus.vocab_size, block_size, nl, nh, d, 0.1, ffn_mult).to(device)
+    for m in get_sparse_linears(model):
+        m.chunk_size = cs
+        m.backend = backend
+
+    is_dense = (policy == "dense")
+    sched = None if is_dense else ChunkScheduler(model, policy, active_frac, cs, device)
+    opt = ChunkedAdam(model, lr=lr, cs=cs, momentum_mode=momentum_mode)
+
+    np_ = model.nparams()
+    overlaps = []
+
+    torch.cuda.synchronize() if device == "cuda" else None
+    t0 = time.perf_counter()
+
+    for step in range(steps):
+        x, y = corpus.get_batch("train", bs, make_gen(step))
+
+        if is_dense:
+            for m in get_sparse_linears(model):
+                m.sparse_enabled = False; m.active_chunks = None
+        else:
+            sched.choose(step, wu, an)
+            for m in get_sparse_linears(model):
+                m.sparse_enabled = True
+                m.sparse_dx = (bwd_mode == "sparse_dX")
+
+        opt.zero_grad()
+        _, loss = model(x, y)
+        loss.backward()
+
+        if sched:
+            sched.update(step, wu)
+
+            # Oracle overlap measurement
+            if measure_oracle and step % oracle_interval == 0 and step >= wu + an:
+                saved = {p: p.grad.clone() for p in model.parameters() if p.grad is not None}
+                for m in get_sparse_linears(model): m.sparse_enabled = False
+                for p in model.parameters(): p.grad = None
+                _, lo = model(x, y); lo.backward()
+                k = max(1, int(active_frac * sched.nc))
+                j, r = sched.measure_overlap(k)
+                overlaps.append((step, j, r))
+                for p in model.parameters():
+                    if p in saved: p.grad = saved[p]
+                for m in get_sparse_linears(model): m.sparse_enabled = True
+
+        opt.step()
+
+        if step % 200 == 0:
+            print(f"    step {step}/{steps} loss={loss.item():.4f}")
+
+    torch.cuda.synchronize() if device == "cuda" else None
+    wall = time.perf_counter() - t0
+
+    for m in get_sparse_linears(model): m.sparse_enabled = False
+    vl, vp = evaluate(model, corpus, bs, n=30)
+
+    del model; torch.cuda.empty_cache() if device == "cuda" else None
+
+    return {
+        "val_loss": vl, "val_ppl": vp, "wall_time": wall,
+        "ms_per_step": 1000*wall/steps, "n_params": np_,
+        "train_loss_final": loss.item(), "overlaps": overlaps,
+    }
+
+def run_seeds(cfg, seeds):
+    results = []
+    for s in seeds:
+        cfg["seed"] = s
+        results.append(run(**cfg))
+    vls = [r["val_loss"] for r in results]
+    ml = sum(vls)/len(vls)
+    sl = (sum((x-ml)**2 for x in vls)/max(1,len(vls)-1))**0.5
+    return {"mean_loss": ml, "std_loss": sl, "results": results,
+            "mean_ms": sum(r["ms_per_step"] for r in results)/len(results)}
+
+# ═══════════════════════════════════════════════════════════════
+# EXPERIMENT 1: PHANTOM MOMENTUM ABLATION
+# ═══════════════════════════════════════════════════════════════
+
+def exp_phantom_momentum(device, steps, seeds, d, nl, nh, bs, block_size, cs, af, wu, an, lr, backend):
+    print("\n" + "="*80)
+    print("EXPERIMENT 1: Phantom Momentum Ablation")
+    print("="*80)
+
+    base = dict(bwd_mode="full_dX", steps=steps, bs=bs, block_size=block_size,
+                nl=nl, nh=nh, d=d, cs=cs, active_frac=af, wu=wu, an=an,
+                lr=lr, device=device, backend=backend)
+
+    configs = [
+        ("dense",               "dense",  "phantom"),
+        ("ema+phantom",         "ema",    "phantom"),
+        ("ema+frozen",          "ema",    "frozen"),
+        ("knn+phantom",         "knn",    "phantom"),
+        ("knn+frozen",          "knn",    "frozen"),
+        ("random+phantom",      "random", "phantom"),
+        ("random+frozen",       "random", "frozen"),
+    ]
+
+    results = {}
+    for name, policy, mm in configs:
+        print(f"\n--- {name} ---")
+        cfg = {**base, "policy": policy, "momentum_mode": mm}
+        results[name] = run_seeds(cfg, seeds)
+
+    print(f"\n{'Method':<22} | {'Val Loss':>18} | {'ms/step':>10}")
+    print("-"*55)
+    for name, _, _ in configs:
+        r = results[name]
+        print(f"{name:<22} | {r['mean_loss']:.4f} ± {r['std_loss']:.4f}   | {r['mean_ms']:>9.1f}")
+
+    return results
+
+# ═══════════════════════════════════════════════════════════════
+# EXPERIMENT 2: COMPUTE-MATCHED BASELINES
+# ═══════════════════════════════════════════════════════════════
+
+def exp_compute_matched(device, steps, seeds, d, nl, nh, bs, block_size, cs, af, wu, an, lr, backend):
+    print("\n" + "="*80)
+    print("EXPERIMENT 2: Compute-Matched Baselines")
+    print("="*80)
+
+    base = dict(bwd_mode="full_dX", steps=steps, bs=bs, block_size=block_size,
+                nl=nl, nh=nh, d=d, cs=cs, active_frac=af, wu=wu, an=an,
+                lr=lr, device=device, backend=backend, momentum_mode="phantom")
+
+    # 1. Sparse reference
+    print("\n--- Sparse (EMA, reference) ---")
+    sparse_r = run_seeds({**base, "policy": "ema"}, seeds)
+
+    # 2. Dense at same steps
+    print("\n--- Dense (same steps) ---")
+    dense_same = run_seeds({**base, "policy": "dense"}, seeds)
+
+    # 3. Dense at compute-matched steps
+    # Sparse does ~70% of dense FLOPs (fwd dense + dX dense + dW at 10%)
+    ratio = (1.0 + 1.0 + af) / 3.0
+    matched_steps = int(steps * ratio)
+    print(f"\n--- Dense (compute-matched, {matched_steps} steps) ---")
+    dense_matched = run_seeds({**base, "policy": "dense", "steps": matched_steps}, seeds)
+
+    # 4. Natively smaller dense model: FFN multiplier = 4 * af = 0.4 (rounded)
+    # This gives a model with ~10% of the FFN capacity
+    small_ffn_mult = max(1, round(4 * af))  # 4*0.1 = 0.4, round to 1
+    print(f"\n--- Small dense (ffn_mult={small_ffn_mult}, capacity-matched) ---")
+    dense_small = run_seeds({**base, "policy": "dense", "ffn_mult": small_ffn_mult}, seeds)
+
+    results = {
+        "sparse_ema": sparse_r,
+        "dense_same_steps": dense_same,
+        f"dense_matched_{matched_steps}steps": dense_matched,
+        f"dense_small_ffn{small_ffn_mult}": dense_small,
+    }
+
+    print(f"\n{'Method':<35} | {'Steps':>6} | {'Params':>8} | {'Val Loss':>18} | {'ms/step':>10}")
+    print("-"*90)
+    for name, r in results.items():
+        np_ = r["results"][0]["n_params"]
+        st = r["results"][0].get("steps", steps) if "steps" in name else steps
+        # read actual steps from config — approximate
+        print(f"{name:<35} | {st if 'matched' not in name else matched_steps:>6} | {np_/1e6:>7.1f}M | {r['mean_loss']:.4f} ± {r['std_loss']:.4f}   | {r['mean_ms']:>9.1f}")
+
+    return results
+
+# ═══════════════════════════════════════════════════════════════
+# EXPERIMENT 3: PREDICTOR ACCURACY (EMA vs KNN vs Oracle)
+# ═══════════════════════════════════════════════════════════════
+
+def exp_predictor_accuracy(device, steps, seeds, d, nl, nh, bs, block_size, cs, af, wu, an, lr, backend):
+    print("\n" + "="*80)
+    print("EXPERIMENT 3: Predictor Accuracy (EMA vs KNN vs Oracle)")
+    print("="*80)
+
+    base = dict(bwd_mode="full_dX", steps=steps, bs=bs, block_size=block_size,
+                nl=nl, nh=nh, d=d, cs=cs, active_frac=af, wu=wu, an=an,
+                lr=lr, device=device, backend=backend, momentum_mode="phantom",
+                measure_oracle=True, oracle_interval=25)
+
+    results = {}
+    for policy in ["ema", "knn", "random"]:
+        print(f"\n--- {policy} ---")
+        results[policy] = run_seeds({**base, "policy": policy}, seeds)
+
+    # Aggregate overlaps
+    for policy in ["ema", "knn", "random"]:
+        print(f"\n{policy.upper()} predictor overlap:")
+        print(f"  {'Step':>6} | {'Jaccard':>10} | {'Recall':>10}")
+        sd = defaultdict(lambda: {"j": [], "r": []})
+        for res in results[policy]["results"]:
+            for s, j, r in res["overlaps"]:
+                sd[s]["j"].append(j); sd[s]["r"].append(r)
+        for s in sorted(sd):
+            mj = sum(sd[s]["j"])/len(sd[s]["j"])
+            mr = sum(sd[s]["r"])/len(sd[s]["r"])
+            print(f"  {s:>6} | {mj:>10.4f} | {mr:>10.4f}")
+
+    print(f"\n{'Policy':<10} | {'Val Loss':>18} | {'ms/step':>10}")
+    print("-"*45)
+    for p in ["ema", "knn", "random"]:
+        r = results[p]
+        print(f"{p:<10} | {r['mean_loss']:.4f} ± {r['std_loss']:.4f}   | {r['mean_ms']:>9.1f}")
+
+    return results
+
+# ═══════════════════════════════════════════════════════════════
+# MAIN
+# ═══════════════════════════════════════════════════════════════
+
+ALL_EXPS = {
+    "phantom_momentum": exp_phantom_momentum,
+    "compute_matched": exp_compute_matched,
+    "predictor_accuracy": exp_predictor_accuracy,
+}
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--experiment", default="all", choices=list(ALL_EXPS)+["all"])
+    p.add_argument("--device", default="cuda")
+    p.add_argument("--steps", type=int, default=1000)
+    p.add_argument("--seeds", default="42,123,456")
+    p.add_argument("--n_embd", type=int, default=1024)
+    p.add_argument("--n_layer", type=int, default=4)
+    p.add_argument("--n_head", type=int, default=8)
+    p.add_argument("--batch_size", type=int, default=8)
+    p.add_argument("--block_size", type=int, default=256)
+    p.add_argument("--chunk_size", type=int, default=64)
+    p.add_argument("--active_fraction", type=float, default=0.10)
+    p.add_argument("--warmup_steps", type=int, default=50)
+    p.add_argument("--anneal_steps", type=int, default=200)
+    p.add_argument("--lr", type=float, default=3e-4)
+    p.add_argument("--backend", default="triton", choices=["triton", "torch"])
+    p.add_argument("--output_dir", default="results")
+    args = p.parse_args()
+
+    seeds = [int(s) for s in args.seeds.split(",")]
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    if args.device == "cuda" and 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"Config: d={args.n_embd} nl={args.n_layer} nh={args.n_head} steps={args.steps} seeds={seeds}")
+    print(f"        cs={args.chunk_size} af={args.active_fraction} backend={args.backend}")
+
+    shared = dict(device=args.device, steps=args.steps, seeds=seeds,
+                  d=args.n_embd, nl=args.n_layer, nh=args.n_head,
+                  bs=args.batch_size, block_size=args.block_size,
+                  cs=args.chunk_size, af=args.active_fraction,
+                  wu=args.warmup_steps, an=args.anneal_steps,
+                  lr=args.lr, backend=args.backend)
+
+    exps = ALL_EXPS if args.experiment == "all" else {args.experiment: ALL_EXPS[args.experiment]}
+    t0 = time.time()
+
+    for name, fn in exps.items():
+        print(f"\n{'#'*80}\n# {name} ({(time.time()-t0)/60:.1f}m elapsed)\n{'#'*80}")
+        sys.stdout.flush()
+        result = fn(**shared)
+
+        def ser(o):
+            if isinstance(o, dict): return {str(k): ser(v) for k,v in o.items()}
+            if isinstance(o, list): return [ser(x) for x in o]
+            return o
+
+        with open(os.path.join(args.output_dir, f"{name}.json"), "w") as f:
+            json.dump(ser(result), f, indent=2, default=str)
+        print(f"✓ {name} saved to {args.output_dir}/{name}.json")
+
+    print(f"\n{'='*80}\nALL COMPLETE in {(time.time()-t0)/60:.1f} minutes\n{'='*80}")
+
+if __name__ == "__main__":
+    main()