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
|
import torch, pickle, math
from pathlib import Path
import lt_ep_train as L
from lt_ep_train import EQBlock
L.DD=Path('data/tinystories_bpe'); L.vocab=pickle.load(open(L.DD/'meta.pkl','rb'))['vocab_size']
dev='cuda'; B=8; T=256; Ttot=400.0 # fixed "time" budget eps*N=Ttot so all eps cover same settling
torch.manual_seed(1234); idx,y=L.get_batch('val',B,T); idx=idx.to(dev) if hasattr(idx,'to') else idx
ck=torch.load('runs/redx_traj/s3200.pt',map_location=dev)
def relax_eps(eps):
N=min(int(Ttot/eps), 24000)
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))
xin=blk.embed(idx).detach(); z=xin.clone(); ress=[]
for t in range(N):
z2=z+eps*blk.force(z,xin).detach()
r=(z2-z).norm().item()/(z.norm().item()+1e-9); ress.append(r); z=z2
if not math.isfinite(r) or r>1e3: return ('DIVERGED',t,r,r)
tail=ress[-min(800,N//4):]
return (N, ress[-1], min(tail), max(tail))
print("=== eps-sweep on redx s3200 (FULL attention, the cycling operator): Euler-artifact vs continuous instability ===")
print("eps*N=400 held fixed (same settling time). Cycle dies as eps shrinks => DISCRETE-EULER ARTIFACT (continuous ODE / analog HW is fine).")
for eps in [0.1, 0.05, 0.03, 0.02, 0.01]:
r=relax_eps(eps)
if r[0]=='DIVERGED': print(f" eps={eps}: DIVERGED at t={r[1]} r={r[2]:.2e}")
else:
N,last,tmin,tmax=r; osc=tmax-tmin
print(f" eps={eps}: N={N:5d} res(last)={last:.3e} tail[min={tmin:.2e},max={tmax:.2e}] osc={osc:.2e} {'CYCLE' if (osc>5e-4 and last>2e-3) else 'CONVERGED' if last<2e-3 else 'floored'}")
print("=== DONE ===")
|