Upload ablations_lite.py with huggingface_hub

ablations_lite.py

@@ -0,0 +1,324 @@
+#!/usr/bin/env python3
+"""Minimal ablation suite — PyTorch only, no Triton. Addresses the 3 critique gaps."""
+import argparse,json,math,os,random,sys,time,urllib.request
+from collections import defaultdict
+import torch,torch.nn as nn,torch.nn.functional as F
+import tiktoken
+print("imports ok",flush=True)
+
+# ── Data ──
+class Corpus:
+    _i=None
+    @classmethod
+    def get(cls,bs,dev):
+        if cls._i is None: cls._i=cls(bs,dev)
+        return cls._i
+    def __init__(self,bs,dev):
+        self.block_size,self.device=bs,dev
+        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)
+        enc=tiktoken.get_encoding("gpt2"); t=enc.encode(open(p).read())
+        self.vocab_size=enc.n_vocab; d=torch.tensor(t,dtype=torch.long)
+        si=int(0.9*len(d)); self.train_data,self.val_data=d[:si],d[si:]
+        print(f"Corpus: V={self.vocab_size} train={len(self.train_data):,} val={len(self.val_data):,}",flush=True)
+    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 mg(s):
+    g=torch.Generator(device="cpu"); g.manual_seed(s); return g
+
+# ── Sparse Linear ──
+class SparseBwd(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx,x,w,b,ac,cs,sdx):
+        ctx.save_for_backward(x,w,ac); ctx.hb=b is not None; ctx.sdx=sdx; ctx.cs=cs
+        return F.linear(x,w,b)
+    @staticmethod
+    def backward(ctx,gy):
+        x,w,ac=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.hb else None
+        gx=torch.zeros_like(xf) if ctx.sdx else gf@w
+        for c in ac.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.sdx: gx+=sl@w[s:e]
+        return gx.reshape(x.shape),gw,gb,None,None,None
+
+class SL(nn.Linear):
+    def __init__(self,i,o,bias=True):
+        super().__init__(i,o,bias=bias)
+        self.se=False; self.sdx=False; self.ac=None; self.cs=64
+    def forward(self,x):
+        if not self.se or self.ac is None: return F.linear(x,self.weight,self.bias)
+        return SparseBwd.apply(x,self.weight,self.bias,self.ac,self.cs,self.sdx)
+
+# ── Model ──
+class Attn(nn.Module):
+    def __init__(self,d,nh,bs,do):
+        super().__init__(); self.nh=nh; self.hd=d//nh
+        self.qkv=SL(d,3*d); self.proj=SL(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.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=(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.proj((a@v).transpose(1,2).contiguous().view(B,T,C))
+
+class FFN(nn.Module):
+    def __init__(self,d,do,fm=4):
+        super().__init__(); self.fc=SL(d,fm*d); self.proj=SL(fm*d,d); self.drop=nn.Dropout(do)
+    def forward(self,x): return self.drop(self.proj(F.gelu(self.fc(x))))
+
+class Blk(nn.Module):
+    def __init__(self,d,nh,bs,do,fm=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,fm)
+    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,fm=4):
+        super().__init__(); self.te=nn.Embedding(V,d); self.pe=nn.Embedding(bs,d)
+        self.blocks=nn.Sequential(*[Blk(d,nh,bs,do,fm) 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)))
+        return lo,F.cross_entropy(lo.view(-1,lo.size(-1)),tgt.view(-1)) if tgt is not None else None
+    def np(self): return sum(p.numel() for p in self.parameters())
+
+def gsl(m): return [x for x in m.modules() if isinstance(x,SL)]
+
+# ── Scheduler ──
+class Sched:
+    def __init__(self,model,pol,frac,cs,dev,beta=0.95):
+        self.pol,self.frac,self.cs,self.dev,self.beta=pol,frac,cs,dev,beta
+        self.lins=gsl(model); self.m2i,self.m2l={},{}; off=0
+        for m in self.lins:
+            m.cs=cs; nc=m.out_features//cs; assert m.out_features%cs==0
+            self.m2i[m]=torch.arange(off,off+nc,device=dev)
+            self.m2l[m]=torch.arange(nc,device=dev); off+=nc
+        self.nc=off; self.ema=torch.zeros(self.nc,device=dev)
+        self.act=torch.zeros(self.nc,dtype=torch.bool,device=dev)
+    def gf(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.gf(step,wu,an)
+        if f>=0.999: self.act.fill_(True); self._inst(); return
+        k=max(1,int(f*self.nc)); self.act.fill_(False)
+        if self.pol=="random": idx=torch.randperm(self.nc,device=self.dev)[:k]
+        else: idx=torch.topk(self.ema+1e-9*torch.rand_like(self.ema),k=k).indices
+        self.act[idx]=True; self._inst()
+    def _inst(self):
+        for m,gi in self.m2i.items(): m.ac=self.m2l[m][self.act[gi]]
+    @torch.no_grad()
+    def update(self,step,wu):
+        cur=torch.zeros_like(self.ema)
+        for m,ids in self.m2i.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.act; 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]
+        return cur
+    @torch.no_grad()
+    def oracle_scores(self):
+        sc=torch.zeros(self.nc,device=self.dev)
+        for m,ids in self.m2i.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 overlap(self,k):
+        o=set(torch.topk(self.oracle_scores(),k=k).indices.tolist())
+        p=set(self.act.nonzero(as_tuple=True)[0].tolist())
+        if not o or not p: return 0.,0.
+        i=o&p; return len(i)/len(o|p),len(i)/len(o)
+
+# ── Adam phantom/frozen ──
+class CAdam:
+    def __init__(self,model,lr=3e-4,cs=64,mm="phantom"):
+        self.model,self.lr,self.cs,self.mm=model,lr,cs,mm
+        self.st={}; self.p2m={}
+        for m in gsl(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.st: self.st[p]={"m":torch.zeros_like(p),"v":torch.zeros_like(p)}
+            m,v=self.st[p]["m"],self.st[p]["v"]
+            sm=self.p2m.get(p); ac=getattr(sm,'ac',None) if sm else None
+            if ac is None:
+                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)
+            elif self.mm=="phantom":
+                m.mul_(0.9).add_(p.grad,alpha=0.1); v.mul_(0.999).addcmul_(p.grad,p.grad,value=0.001)
+                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)
+            else:
+                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)
+
+# ── Eval ──
+@torch.no_grad()
+def ev(model,corpus,bs,n=20,seed=9999):
+    model.eval(); ls=[model(*corpus.get_batch("val",bs,mg(seed+i)))[1].item() for i in range(n)]
+    model.train(); a=sum(ls)/len(ls); return a,math.exp(min(a,20))
+
+# ── Single run ──
+def run1(pol,steps,bs,bsz,nl,nh,d,cs,af,wu,an,lr,dev,seed,mm="phantom",fm=4,mo=False,oi=50):
+    torch.manual_seed(seed); random.seed(seed)
+    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
+    corpus=Corpus.get(bsz,dev)
+    model=GPT(corpus.vocab_size,bsz,nl,nh,d,0.1,fm).to(dev)
+    for m in gsl(model): m.cs=cs
+    dense=pol=="dense"; sched=None if dense else Sched(model,pol,af,cs,dev)
+    opt=CAdam(model,lr,cs,mm); np_=model.np(); overlaps=[]
+    if dev=="cuda": torch.cuda.synchronize()
+    t0=time.perf_counter()
+    for step in range(steps):
+        x,y=corpus.get_batch("train",bs,mg(step))
+        if dense:
+            for m in gsl(model): m.se=False; m.ac=None
+        else:
+            sched.choose(step,wu,an)
+            for m in gsl(model): m.se=True; m.sdx=False
+        opt.zero_grad(); _,loss=model(x,y); loss.backward()
+        if sched:
+            sched.update(step,wu)
+            if mo and step%oi==0 and step>=wu+an:
+                saved={p:p.grad.clone() for p in model.parameters() if p.grad is not None}
+                for m in gsl(model): m.se=False
+                for p in model.parameters(): p.grad=None
+                _,lo=model(x,y); lo.backward()
+                k=max(1,int(af*sched.nc)); j,r=sched.overlap(k)
+                overlaps.append((step,j,r))
+                for p in model.parameters():
+                    if p in saved: p.grad=saved[p]
+                for m in gsl(model): m.se=True
+        opt.step()
+        if step%200==0: print(f"    step {step}/{steps} loss={loss.item():.4f}",flush=True)
+    if dev=="cuda": torch.cuda.synchronize()
+    wall=time.perf_counter()-t0
+    for m in gsl(model): m.se=False
+    vl,vp=ev(model,corpus,bs,n=30)
+    del model; torch.cuda.empty_cache() if dev=="cuda" else None
+    return {"vl":vl,"vp":vp,"wall":wall,"ms":1000*wall/steps,"np":np_,"tl":loss.item(),"ov":overlaps}
+
+def runs(cfg,seeds):
+    rs=[]; 
+    for s in seeds: cfg["seed"]=s; rs.append(run1(**cfg))
+    vls=[r["vl"] for r in rs]; ml=sum(vls)/len(vls)
+    sl=(sum((x-ml)**2 for x in vls)/max(1,len(vls)-1))**0.5
+    return {"ml":ml,"sl":sl,"rs":rs,"ms":sum(r["ms"] for r in rs)/len(rs)}
+
+# ── Experiments ──
+def exp1(dev,steps,seeds,d,nl,nh,bs,bsz,cs,af,wu,an,lr):
+    """Phantom momentum ablation"""
+    print("\n"+"="*80+"\nEXP 1: Phantom Momentum\n"+"="*80,flush=True)
+    base=dict(steps=steps,bs=bs,bsz=bsz,nl=nl,nh=nh,d=d,cs=cs,af=af,wu=wu,an=an,lr=lr,dev=dev)
+    cfgs=[("dense","dense","phantom"),("ema+phantom","ema","phantom"),("ema+frozen","ema","frozen"),
+          ("random+phantom","random","phantom"),("random+frozen","random","frozen")]
+    R={}
+    for name,pol,mm in cfgs:
+        print(f"\n--- {name} ---",flush=True)
+        R[name]=runs({**base,"pol":pol,"mm":mm},seeds)
+    print(f"\n{'Method':<20} | {'Val Loss':>18} | {'ms/step':>8}",flush=True)
+    print("-"*52,flush=True)
+    for name,_,_ in cfgs:
+        r=R[name]; print(f"{name:<20} | {r['ml']:.4f} ± {r['sl']:.4f}   | {r['ms']:>7.1f}",flush=True)
+    return R
+
+def exp2(dev,steps,seeds,d,nl,nh,bs,bsz,cs,af,wu,an,lr):
+    """Compute-matched baselines"""
+    print("\n"+"="*80+"\nEXP 2: Compute-Matched\n"+"="*80,flush=True)
+    base=dict(steps=steps,bs=bs,bsz=bsz,nl=nl,nh=nh,d=d,cs=cs,af=af,wu=wu,an=an,lr=lr,dev=dev,mm="phantom")
+    print("\n--- Sparse EMA ---",flush=True)
+    sp=runs({**base,"pol":"ema"},seeds)
+    print("\n--- Dense same steps ---",flush=True)
+    ds=runs({**base,"pol":"dense"},seeds)
+    ms=int(steps*(1+1+af)/3)
+    print(f"\n--- Dense matched {ms} steps ---",flush=True)
+    dm=runs({**base,"pol":"dense","steps":ms},seeds)
+    sfm=max(1,round(4*af))
+    print(f"\n--- Small dense ffn_mult={sfm} ---",flush=True)
+    dd=runs({**base,"pol":"dense","fm":sfm},seeds)
+    R={"sparse_ema":sp,"dense_same":ds,f"dense_{ms}steps":dm,f"dense_ffn{sfm}":dd}
+    print(f"\n{'Method':<25} | {'Params':>7} | {'Val Loss':>18} | {'ms/step':>8}",flush=True)
+    print("-"*65,flush=True)
+    for n,r in R.items():
+        np_=r["rs"][0]["np"]
+        print(f"{n:<25} | {np_/1e6:>6.1f}M | {r['ml']:.4f} ± {r['sl']:.4f}   | {r['ms']:>7.1f}",flush=True)
+    return R
+
+def exp3(dev,steps,seeds,d,nl,nh,bs,bsz,cs,af,wu,an,lr):
+    """Predictor accuracy"""
+    print("\n"+"="*80+"\nEXP 3: Predictor Accuracy\n"+"="*80,flush=True)
+    base=dict(steps=steps,bs=bs,bsz=bsz,nl=nl,nh=nh,d=d,cs=cs,af=af,wu=wu,an=an,lr=lr,dev=dev,mm="phantom",mo=True,oi=25)
+    R={}
+    for pol in ["ema","random"]:
+        print(f"\n--- {pol} ---",flush=True)
+        R[pol]=runs({**base,"pol":pol},seeds)
+    for pol in ["ema","random"]:
+        print(f"\n{pol.upper()} overlap:",flush=True)
+        sd=defaultdict(lambda:{"j":[],"r":[]})
+        for res in R[pol]["rs"]:
+            for s,j,r in res["ov"]: 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"  step {s:>5}: jaccard={mj:.4f} recall={mr:.4f}",flush=True)
+    print(f"\n{'Pol':<8} | {'Val Loss':>18} | {'ms/step':>8}",flush=True)
+    for p in ["ema","random"]:
+        r=R[p]; print(f"{p:<8} | {r['ml']:.4f} ± {r['sl']:.4f}   | {r['ms']:>7.1f}",flush=True)
+    return R
+
+def main():
+    p=argparse.ArgumentParser()
+    p.add_argument("--exp",default="all",choices=["all","exp1","exp2","exp3"])
+    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("--d",type=int,default=1024)
+    p.add_argument("--nl",type=int,default=4); p.add_argument("--nh",type=int,default=8)
+    p.add_argument("--bs",type=int,default=8); p.add_argument("--bsz",type=int,default=256)
+    p.add_argument("--cs",type=int,default=64); p.add_argument("--af",type=float,default=0.10)
+    p.add_argument("--wu",type=int,default=50); p.add_argument("--an",type=int,default=200)
+    p.add_argument("--lr",type=float,default=3e-4)
+    a=p.parse_args(); seeds=[int(s) for s in a.seeds.split(",")]
+    if a.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",flush=True)
+    print(f"d={a.d} nl={a.nl} nh={a.nh} steps={a.steps} seeds={seeds} cs={a.cs} af={a.af}",flush=True)
+    sh=dict(dev=a.device,steps=a.steps,seeds=seeds,d=a.d,nl=a.nl,nh=a.nh,bs=a.bs,bsz=a.bsz,cs=a.cs,af=a.af,wu=a.wu,an=a.an,lr=a.lr)
+    t0=time.time()
+    exps={"exp1":exp1,"exp2":exp2,"exp3":exp3} if a.exp=="all" else {a.exp:{"exp1":exp1,"exp2":exp2,"exp3":exp3}[a.exp]}
+    for n,fn in exps.items():
+        print(f"\n{'#'*60}\n# {n} ({(time.time()-t0)/60:.1f}m)\n{'#'*60}",flush=True)
+        r=fn(**sh)
+        with open(f"{n}.json","w") as f: json.dump(r,f,indent=2,default=str)
+        print(f"✓ {n}.json saved",flush=True)
+    print(f"\nDONE in {(time.time()-t0)/60:.1f}m",flush=True)
+
+if __name__=="__main__": main()