1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
|
import torch, pickle, numpy as np
from pathlib import Path
from scipy.sparse.linalg import LinearOperator, eigs
import lt_ep_train as L
from lt_ep_train import EQBlock, relax
L.DD=Path('data/tinystories_bpe'); L.vocab=pickle.load(open(L.DD/'meta.pkl','rb'))['vocab_size']
dev='cuda'; B=2; T=256; eps=0.1
torch.manual_seed(1234); idx,y=L.get_batch('val',B,T); idx=idx.to(dev) if hasattr(idx,'to') else idx
def load_blk(path):
ck=torch.load(path,map_location=dev)
blk=EQBlock(512,16,256,256,s=1.0,c=1.0,attn_mode='thick'); blk.qknorm=True
with torch.no_grad():
for p,w in zip(blk.allp,ck['allp']): p.copy_(w.to(dev))
return blk, ck.get('best','?')
@torch.no_grad()
def lead_eigs(blk, relax_steps=250, k=6, hrel=1e-3):
xin=blk.embed(idx).detach(); z=relax(blk,xin.clone(),xin,relax_steps,eps).detach()
F0=blk.force(z,xin).detach(); zn=z.norm().item(); g=F0.norm().item()/(zn+1e-9); h=hrel*zn; shp=z.shape; N=z.numel()
def Jv(vt):
nv=vt.norm().item()
if nv<1e-20: return torch.zeros_like(vt)
return (blk.force(z+h*(vt/nv),xin).detach()-F0)/h*nv
def matvec(v):
vt=torch.from_numpy(np.ascontiguousarray(v).astype('float32')).reshape(shp).to(dev)
return (vt+eps*Jv(vt)).double().cpu().numpy().reshape(-1)
op=LinearOperator((N,N),matvec=matvec,dtype='float64')
vals=eigs(op,k=k,which='LM',return_eigenvectors=False,maxiter=4000,tol=1e-5)
return g, sorted(vals,key=lambda x:-abs(x))
for tag,path in [("ep_jacreg ~2.75 (ADAPTIVE jacreg)","runs/ep_jacreg.pt"),
("redx s3200 2.74 (FROZEN jacreg, BLEW)","runs/redx_traj/s3200.pt")]:
try:
blk,best=load_blk(path); g,vals=lead_eigs(blk)
print(f"=== {tag} best={best} g_floor={g:.4f} ===")
for lam in vals[:4]:
mu=(lam-1)/eps
print(f" |lam|={abs(lam):.5f} mu={mu.real:+.4f}{mu.imag:+.4f}j [{'UNSTABLE' if abs(lam)>1+1e-4 else 'STABLE'}{' rot' if abs(lam.imag)>1e-3 else ' real'}{' ReMu<0' if mu.real<-1e-4 else ' ReMu>=0'}]")
except Exception as e: print(f"=== {tag}: ERR {e} ===")
print("=== DONE === ep_jacreg ReMu<0 => jacreg pushed it STABLE where frozen-jacreg redx is ReMu>0 (mechanism confirmed)")
|
