exp4_relaxation.py

#!/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()