Fix: backward() inside @torch .no_grad() — use torch.enable_grad() for dense gradient computation

exp5_mechanism.py

@@ -29,9 +29,6 @@ import torch,torch.nn as nn,torch.nn.functional as F
 import tiktoken
 print("imports ok",flush=True)
 
-# ═══════════════════════════════════════════════════════════════
-# DATA
-# ═══════════════════════════════════════════════════════════════
 class Corpus:
     _i=None
     @classmethod
@@ -56,9 +53,6 @@ class Corpus:
 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):
@@ -86,9 +80,6 @@ class SL(nn.Linear):
         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
@@ -127,9 +118,6 @@ class GPT(nn.Module):
     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 (builds similarity matrix during warmup)
-# ═══════════════════════════════════════════════════════════════
 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
@@ -184,17 +172,12 @@ class Sched:
         for _,ids in self.m2i.items(): ok[ids[:,None],ids[None,:]]=True
         self.similarity=torch.where(ok,S,torch.zeros_like(S))
     def mask_jaccard(self):
-        """Jaccard between current and previous active set."""
         if self.prev_act.sum()==0: return 0.0
         i=(self.act&self.prev_act).sum().item()
         u=(self.act|self.prev_act).sum().item()
         return i/max(u,1)
 
-# ═══════════════════════════════════════════════════════════════
-# RELAXERS
-# ═══════════════════════════════════════════════════════════════
 class GraphRelaxer:
-    """Graph Laplacian relaxation using the real similarity matrix."""
     def __init__(self, sched, alpha=0.1, iters=3):
         self.sched,self.alpha,self.iters=sched,alpha,iters
     @torch.no_grad()
@@ -213,11 +196,10 @@ class GraphRelaxer:
                 Wf=W.reshape(nc,-1); Wa=(S_n@Wf).view(nc,cs,di)
                 W[li]=(1-self.alpha)*W[li]+self.alpha*Wa[li]
             m.weight.data=W.view(m.out_features,di)
-            deltas[m]=W[li]-W_before  # (n_inactive, cs, di)
+            deltas[m]=W[li]-W_before
         return deltas
 
 class RollRelaxer:
-    """Spatial neighbor relaxation via torch.roll."""
     def __init__(self, sched, alpha=0.1, iters=3):
         self.sched,self.alpha,self.iters=sched,alpha,iters
     @torch.no_grad()
@@ -237,7 +219,6 @@ class RollRelaxer:
         return deltas
 
 class RandomRelaxer:
-    """Control: random similarity matrix (same sparsity pattern, random values)."""
     def __init__(self, sched, alpha=0.1, iters=3):
         self.sched,self.alpha,self.iters=sched,alpha,iters
         self._rand_sim=None
@@ -245,7 +226,6 @@ class RandomRelaxer:
         if self._rand_sim is not None: return self._rand_sim
         S=self.sched.similarity
         if S is None: return None
-        # Random positive values with same mask structure
         R=torch.rand_like(S)*S.abs().mean()
         R.fill_diagonal_(0)
         ok=torch.zeros_like(R,dtype=torch.bool)
@@ -272,7 +252,6 @@ class RandomRelaxer:
         return deltas
 
 class ShuffledGraphRelaxer:
-    """Control: real similarity stats, shuffled structure within each layer."""
     def __init__(self, sched, alpha=0.1, iters=3):
         self.sched,self.alpha,self.iters=sched,alpha,iters
         self._shuf_sim=None
@@ -281,11 +260,9 @@ class ShuffledGraphRelaxer:
         S=self.sched.similarity
         if S is None: return None
         Ss=S.clone()
-        # Shuffle within each layer block
         for _,ids in self.sched.m2i.items():
             n=len(ids)
-            block=Ss[ids][:,ids].clone()  # (n,n)
-            # Shuffle rows and columns with same permutation
+            block=Ss[ids][:,ids].clone()
             perm=torch.randperm(n,device=S.device)
             block=block[perm][:,perm]
             block.fill_diagonal_(0)
@@ -312,12 +289,8 @@ class ShuffledGraphRelaxer:
         return deltas
 
 class NullRelaxer:
-    """No-op relaxer."""
     def relax(self): return {}
 
-# ═══════════════════════════════════════════════════════════════
-# OPTIMIZER
-# ═══════════════════════════════════════════════════════════════
 class CAdam:
     def __init__(self,model,lr=3e-4,cs=64):
         self.model,self.lr,self.cs=model,lr,cs
@@ -343,18 +316,11 @@ class CAdam:
                     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))
 
-# ═══════════════════════════════════════════════════════════════
-# ORACLE GRADIENT DIAGNOSTIC
-# ═══════════════════════════════════════════════════════════════
-@torch.no_grad()
 def compute_relaxer_diagnostics(model, sched, relaxer_deltas, x, y, corpus, bs, cs):
     """
     Compare relaxer delta on inactive chunks to what dense gradient would have been.
@@ -362,38 +328,33 @@ def compute_relaxer_diagnostics(model, sched, relaxer_deltas, x, y, corpus, bs,
     """
     if not relaxer_deltas: return None, None
 
-    # Compute dense gradients
+    # Need gradients enabled for the dense forward/backward
     for m in gsl(model): m.se=False
     for p in model.parameters(): p.grad=None
-    _,lo=model(x,y); lo.backward()
+    with torch.enable_grad():
+        _,lo=model(x,y)
+        lo.backward()
 
     cos_sims=[]; mag_ratios=[]
-    for m,delta in relaxer_deltas.items():
-        if m not in sched.m2i: continue
-        ids=sched.m2i[m]; nc=len(ids); di=m.weight.shape[1]
-        la=sched.act[ids]; li=~la
-        if li.sum()==0 or m.weight.grad is None: continue
-
-        # Dense gradient for inactive chunks, reshaped
-        dense_g=m.weight.grad.view(nc,cs,di)[li]  # (n_inact, cs, di)
-
-        # Flatten for cosine/magnitude
-        d_flat=delta.reshape(-1); g_flat=dense_g.reshape(-1)
-        dn=d_flat.norm(); gn=g_flat.norm()
-        if dn>1e-12 and gn>1e-12:
-            cos_sims.append(F.cosine_similarity(d_flat.unsqueeze(0),g_flat.unsqueeze(0)).item())
-            mag_ratios.append((dn/gn).item())
-
-    # Restore sparse mode
+    with torch.no_grad():
+        for m,delta in relaxer_deltas.items():
+            if m not in sched.m2i: continue
+            ids=sched.m2i[m]; nc=len(ids); di=m.weight.shape[1]
+            la=sched.act[ids]; li=~la
+            if li.sum()==0 or m.weight.grad is None: continue
+            dense_g=m.weight.grad.view(nc,cs,di)[li]
+            d_flat=delta.reshape(-1); g_flat=dense_g.reshape(-1)
+            dn=d_flat.norm(); gn=g_flat.norm()
+            if dn>1e-12 and gn>1e-12:
+                cos_sims.append(F.cosine_similarity(d_flat.unsqueeze(0),g_flat.unsqueeze(0)).item())
+                mag_ratios.append((dn/gn).item())
+
     for m in gsl(model): m.se=True
     for p in model.parameters(): p.grad=None
 
     if not cos_sims: return None, None
     return sum(cos_sims)/len(cos_sims), sum(mag_ratios)/len(mag_ratios)
 
-# ═══════════════════════════════════════════════════════════════
-# SINGLE RUN
-# ══════════════════════════════════��════════════════════════════
 def run1(mode, steps, bs, bsz, nl, nh, d, cs, af, wu, an, lr, dev, seed,
          alpha=0.1, iters=3, diag_interval=100):
     torch.manual_seed(seed); random.seed(seed)
@@ -410,12 +371,10 @@ def run1(mode, steps, bs, bsz, nl, nh, d, cs, af, wu, an, lr, dev, seed,
     if is_sparse:
         sched=Sched(model,af,cs,dev)
     elif needs_relax:
-        # Dense + relax: need scheduler for similarity matrix but run dense forward/backward
         sched=Sched(model,af,cs,dev)
 
     opt=CAdam(model,lr,cs)
 
-    # Create relaxer
     if not needs_relax:
         relaxer=NullRelaxer()
     elif "random" in mode:
@@ -443,7 +402,6 @@ def run1(mode, steps, bs, bsz, nl, nh, d, cs, af, wu, an, lr, dev, seed,
             for m in gsl(model): m.se=True; m.sdx=False
         else:
             for m in gsl(model): m.se=False; m.ac=None
-            # For dense+relax: still run scheduler to build similarity & set active mask
             if sched:
                 sched.choose(step,wu,an)
 
@@ -455,10 +413,8 @@ def run1(mode, steps, bs, bsz, nl, nh, d, cs, af, wu, an, lr, dev, seed,
 
         opt.step()
 
-        # Relaxation (only after annealing completes)
         relax_deltas={}
         if needs_relax and step>=wu+an:
-            # For dense+relax: temporarily set active mask so relaxer knows what's "active"
             if is_dense and sched:
                 for m,ids in sched.m2i.items():
                     m.ac=sched.m2l[m][sched.act[ids]]
@@ -466,7 +422,6 @@ def run1(mode, steps, bs, bsz, nl, nh, d, cs, af, wu, an, lr, dev, seed,
             if is_dense and sched:
                 for m in gsl(model): m.ac=None
 
-        # Diagnostics
         if step%50==0:
             vl,_=ev(model,corpus,bs,n=10,seed=7777)
             val_curve.append((step,vl))
@@ -500,9 +455,6 @@ def runs(cfg,seeds):
     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
-# ═══════════════════════════════════════════════════════════════
 def main():
     p=argparse.ArgumentParser()
     p.add_argument("--device",default="cuda"); p.add_argument("--steps",type=int,default=1000)
@@ -522,7 +474,6 @@ def main():
     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,alpha=0.1,iters=3)
 
-    # ── Part 1: Main configs ──
     configs=[
         ("dense",                "dense"),
         ("dense+relax_graph",    "dense+relax_graph"),
@@ -549,7 +500,6 @@ def main():
         r=R[name]
         print(f"{name:<25} | {r['ml']:.4f} ± {r['sl']:.4f}      | {r['ms']:>7.1f}",flush=True)
 
-    # ── Part 2: Alpha sweep ──
     print(f"\n--- Alpha sweep (ema+relax_graph, 5 seeds) ---",flush=True)
     print(f"{'alpha':>6} | {'Val Loss':>20} | {'ms/step':>8}",flush=True)
     print("-"*42,flush=True)
@@ -559,7 +509,6 @@ def main():
         alpha_results[alpha]=r
         print(f"{alpha:>6.2f} | {r['ml']:.4f} ± {r['sl']:.4f}      | {r['ms']:>7.1f}",flush=True)
 
-    # ── Part 3: Diagnostics summary ──
     print(f"\n--- Diagnostics (grad_cos, mag_ratio) ---",flush=True)
     for name in ["ema+relax_graph","ema+relax_roll","ema+relax_random","ema+relax_shuffled"]:
         if name not in R: continue
@@ -571,7 +520,6 @@ def main():
             gc_m=sum(gc_all)/len(gc_all); mr_m=sum(mr_all)/len(mr_all)
             print(f"  {name:<25}: grad_cos={gc_m:.4f}  mag_ratio={mr_m:.4f}",flush=True)
 
-    # Save
     all_results={"configs":R,"alpha_sweep":alpha_results}
     with open("exp5.json","w") as f:
         json.dump(all_results,f,indent=2,default=str)