Upload exp4_relaxation.py with huggingface_hub

exp4_relaxation.py

@@ -0,0 +1,417 @@
+#!/usr/bin/env python3
+"""
+Experiment 4: Graph Laplacian Weight Relaxation
+
+After each sparse gradient step, treat updated (active) chunk weights as
+Dirichlet boundary conditions and relax inactive chunk weights via
+diffusion on the chunk similarity graph.
+
+This is NOT gradient imputation (predicting what the gradient would have
+been). This is post-hoc weight smoothing: given that the active chunks
+moved, nudge the inactive chunks toward structural consistency.
+
+The similarity graph comes from the EMA gradient history (same as KNN
+scheduler in v18). Chunks with correlated gradient histories are
+"neighbors" in the graph — their weights should co-vary.
+
+Modes tested:
+  - dense:          standard dense training (reference)
+  - ema_only:       sparse EMA, no relaxation (existing method)
+  - ema+relax_graph: sparse EMA + graph Laplacian relaxation on inactive weights
+  - ema+relax_roll:  sparse EMA + naive spatial relaxation (torch.roll, control)
+
+The graph relaxation should outperform roll relaxation on dense Linear layers
+because roll assumes spatial adjacency that doesn't exist.
+"""
+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 (reuse from ablations_lite) ──
+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
+
+# ── Model (same as ablations_lite) ──
+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)
+
+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 with similarity matrix ──
+class Sched:
+    def __init__(self,model,frac,cs,dev,beta=0.95,sim_hist=128,min_sim=8):
+        self.frac,self.cs,self.dev,self.beta=frac,cs,dev,beta
+        self.sim_hist,self.min_sim=sim_hist,min_sim
+        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)
+        self.mass_history=[]; self.similarity=None
+    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)
+        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]
+        # Build similarity 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:
+                self._build_sim()
+        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)
+        # Only allow similarity within same layer's chunks
+        ok=torch.zeros_like(S,dtype=torch.bool)
+        for _,ids in self.m2i.items(): ok[ids[:,None],ids[None,:]]=True
+        self.similarity=torch.where(ok,S,torch.zeros_like(S))
+
+# ── Graph Laplacian Relaxation ──
+class WeightRelaxer:
+    """
+    After each sparse optimizer step, relax inactive chunk weights via
+    diffusion on the chunk similarity graph.
+
+    For each SparseLinear layer:
+      1. Reshape weight into (n_chunks, chunk_size, d_in)
+      2. For inactive chunks: new_w[c] = (1-alpha)*w[c] + alpha * sum_j(S[c,j]*w[j]) / sum_j(S[c,j])
+         where S is the similarity matrix restricted to the same layer.
+      3. Active chunks are clamped (Dirichlet boundary).
+
+    alpha controls relaxation strength. iterations controls convergence depth.
+    """
+    def __init__(self, sched, alpha=0.1, iterations=3):
+        self.sched = sched
+        self.alpha = alpha
+        self.iterations = iterations
+
+    @torch.no_grad()
+    def relax(self):
+        S = self.sched.similarity
+        if S is None:
+            return  # No similarity built yet (still in warmup)
+
+        act = self.sched.act  # (n_chunks_total,) bool
+
+        for m, ids in self.sched.m2i.items():
+            nc = len(ids)
+            cs = self.sched.cs
+            d_in = m.weight.shape[1]
+
+            # Local similarity matrix for this layer
+            S_local = S[ids][:, ids]  # (nc, nc)
+
+            # Normalize: each row sums to 1 (or 0 for isolated chunks)
+            row_sum = S_local.sum(dim=1, keepdim=True) + 1e-12
+            S_norm = S_local / row_sum  # (nc, nc)
+
+            # Local active mask
+            local_act = act[ids]  # (nc,) bool
+            local_inact = ~local_act
+
+            if local_inact.sum() == 0:
+                continue  # All active, nothing to relax
+
+            # Reshape weight: (O, I) -> (nc, cs, I)
+            W = m.weight.data.view(nc, cs, d_in)
+
+            for _ in range(self.iterations):
+                # Compute neighbor-weighted average for ALL chunks
+                # W_avg[c] = sum_j S_norm[c,j] * W[j]
+                # Shape: (nc, cs, I) = (nc, nc) @ (nc, cs*I) reshaped
+                W_flat = W.reshape(nc, -1)  # (nc, cs*I)
+                W_avg = (S_norm @ W_flat).view(nc, cs, d_in)  # (nc, cs, I)
+
+                # Blend: only for inactive chunks
+                # w_new = (1 - alpha) * w_old + alpha * w_avg
+                W[local_inact] = (1 - self.alpha) * W[local_inact] + self.alpha * W_avg[local_inact]
+
+            # Write back
+            m.weight.data = W.view(m.out_features, d_in)
+
+
+class NaiveRollRelaxer:
+    """Control: spatial relaxation via torch.roll (wrong neighborhood for dense layers)."""
+    def __init__(self, sched, alpha=0.1, iterations=3):
+        self.sched = sched
+        self.alpha = alpha
+        self.iterations = iterations
+
+    @torch.no_grad()
+    def relax(self):
+        act = self.sched.act
+
+        for m, ids in self.sched.m2i.items():
+            nc = len(ids)
+            cs = self.sched.cs
+            d_in = m.weight.shape[1]
+            local_act = act[ids]
+            local_inact = ~local_act
+
+            if local_inact.sum() == 0:
+                continue
+
+            W = m.weight.data.view(nc, cs, d_in)
+
+            for _ in range(self.iterations):
+                # Spatial neighbors: previous and next chunk
+                W_prev = torch.roll(W, 1, dims=0)
+                W_next = torch.roll(W, -1, dims=0)
+                W_avg = (W_prev + W_next) / 2.0
+
+                W[local_inact] = (1 - self.alpha) * W[local_inact] + self.alpha * W_avg[local_inact]
+
+            m.weight.data = W.view(m.out_features, d_in)
+
+
+# ── Adam (phantom mode only for simplicity) ──
+class CAdam:
+    def __init__(self,model,lr=3e-4,cs=64):
+        self.model,self.lr,self.cs=model,lr,cs
+        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)
+            else:
+                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)
+
+# ── 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(relax_mode, steps, bs, bsz, nl, nh, d, cs, af, wu, an, lr, dev, seed,
+         relax_alpha=0.1, relax_iters=3):
+    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).to(dev)
+    for m in gsl(model): m.cs=cs
+    dense=(relax_mode=="dense")
+    sched=None if dense else Sched(model,af,cs,dev)
+    opt=CAdam(model,lr,cs)
+
+    # Set up relaxer
+    relaxer=None
+    if relax_mode=="ema+relax_graph" and sched:
+        relaxer=WeightRelaxer(sched, alpha=relax_alpha, iterations=relax_iters)
+    elif relax_mode=="ema+relax_roll" and sched:
+        relaxer=NaiveRollRelaxer(sched, alpha=relax_alpha, iterations=relax_iters)
+
+    np_=model.np()
+    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)
+        opt.step()
+
+        # Post-optimizer relaxation
+        if relaxer and step >= wu + an:  # Only after annealing completes
+            relaxer.relax()
+
+        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()}
+
+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)}
+
+# ── Main experiment ──
+def main():
+    p=argparse.ArgumentParser()
+    p.add_argument("--device",default="cuda"); p.add_argument("--steps",type=int,default=500)
+    p.add_argument("--seeds",default="42,123"); 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)
+    p.add_argument("--relax_alpha",type=float,default=0.1)
+    p.add_argument("--relax_iters",type=int,default=3)
+    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} steps={a.steps} seeds={seeds} alpha={a.relax_alpha} iters={a.relax_iters}",flush=True)
+
+    base=dict(steps=a.steps,bs=a.bs,bsz=a.bsz,nl=a.nl,nh=a.nh,d=a.d,cs=a.cs,af=a.af,
+              wu=a.wu,an=a.an,lr=a.lr,dev=a.device,relax_alpha=a.relax_alpha,relax_iters=a.relax_iters)
+
+    configs=[
+        ("dense",           "dense"),
+        ("ema_only",        "ema_only"),
+        ("ema+relax_graph", "ema+relax_graph"),
+        ("ema+relax_roll",  "ema+relax_roll"),
+    ]
+
+    print("\n"+"="*80,flush=True)
+    print("EXP 4: Graph Laplacian Weight Relaxation",flush=True)
+    print("="*80,flush=True)
+
+    R={}
+    for name,mode in configs:
+        print(f"\n--- {name} ---",flush=True)
+        R[name]=runs({**base,"relax_mode":mode},seeds)
+
+    print(f"\n{'Method':<20} | {'Val Loss':>18} | {'ms/step':>8} | {'train_loss':>10}",flush=True)
+    print("-"*65,flush=True)
+    for name,_ in configs:
+        r=R[name]; tl=sum(x["tl"] for x in r["rs"])/len(r["rs"])
+        print(f"{name:<20} | {r['ml']:.4f} ± {r['sl']:.4f}   | {r['ms']:>7.1f} | {tl:>9.4f}",flush=True)
+
+    # Also sweep alpha
+    print(f"\n--- Alpha sweep (graph relaxation) ---",flush=True)
+    print(f"{'alpha':>6} | {'iters':>5} | {'Val Loss':>18} | {'ms/step':>8}",flush=True)
+    print("-"*50,flush=True)
+    for alpha in [0.01, 0.05, 0.1, 0.2, 0.5]:
+        r=runs({**base,"relax_mode":"ema+relax_graph","relax_alpha":alpha,"relax_iters":3},seeds)
+        print(f"{alpha:>6.2f} | {3:>5} | {r['ml']:.4f} ± {r['sl']:.4f}   | {r['ms']:>7.1f}",flush=True)
+
+    with open("exp4.json","w") as f: json.dump(R,f,indent=2,default=str)
+    print("\n✓ exp4.json saved",flush=True)
+
+if __name__=="__main__": main()