nx_nofloat_adamw32_gate.nx source

1// nx_nofloat_adamw32_gate.nx -- the REAL win from the R4 scale-up work: a numerically STABLE fixed-point Adam. 2// DIAGNOSIS: the Q16 AdamW (lib nfa_adamw) DIVERGES/STALLS on small gradients because g^2 UNDERFLOWS in Q16 -- 3// e.g. g=0.06 (Q16 3932) -> qmul(g,g)=235 (0.0036) and the (1-b2)*g^2 increment qmul(66,235)=0, so the 2nd 4// moment v never accumulates, the denominator collapses to eps, and the update mhat/eps explodes (or stalls). 5// FIX: accumulate v in Q32 and compute g^2 = g*g at FULL precision (no >>16) so small gradients survive. 6// T1: Q32-moment AdamW CONVERGES on a small-gradient fit (loss -> ~0). 7// T2 (measured exceed): Q32 final loss << Q16 final loss on the SAME problem (the underflow fix is real). 8// Self-contained controlled experiment (n=8 fit, hand-coded grad 2*(w-t), no transformer). Pure integer. 9// expect_exit: 0 Sovereign: nx_nofloat_autograd (for nfa_adamw + nfa_isqrt) + nx_syscalls. 10import "nx_nofloat_autograd.nx" 11import "nx_syscalls.nx" 12import "nx_gate_emit_lib.nx" 13const Q16: i64 = 65536 14 15 16 17// Q32-moment AdamW (the fix): v in Q32, g^2 = g*g full precision; sqrt(vhat_Q32) via isqrt = Q16 sqrt directly. 18func adamw32_one(w: *i64, g: *i64, m: *i64, v: *i64, n: i64, lr: i64, b1: i64, b2: i64, eps: i64, t: i64) -> i64 { 19 var c1: i64=65536; var c2: i64=65536; var k: i64=0 20 while k<t { if c2==0 { k=t } else { c1=(c1*b1)>>16; c2=(c2*b2)>>16; k=k+1 } } 21 let bc1: i64=65536-c1; let bc2: i64=65536-c2 22 var i: i64=0 23 while i<n { 24 m[i] = ((b1*m[i])>>16) + (((65536-b1)*g[i])>>16) // m Q16 25 let g2: i64 = g[i]*g[i] // Q32 -- FULL precision, no underflow 26 v[i] = ((b2*v[i])>>16) + (((65536-b2)*g2)>>16) // v Q32 27 var mh: i64=m[i]; if bc1>0 { mh=(m[i]<<16)/bc1 } 28 var vh: i64=v[i]; if bc2>0 { vh=(v[i]<<16)/bc2 } 29 if vh<0 { vh=0 } 30 let sq: i64 = nfa_isqrt(vh) // sqrt of Q32 = Q16 sqrt directly 31 let denom: i64 = sq + eps 32 var upd: i64=0; if denom>0 { upd=(mh<<16)/denom } 33 w[i] = w[i] - ((lr*upd)>>16) 34 i=i+1 35 } 36 return 0 37} 38// loss = sum (w-t)^2 in Q32 (raw d*d, so tiny residuals don't underflow the metric itself) 39func loss_q32(w: *i64, t: *i64, n: i64) -> i64 { var s: i64=0; var i: i64=0; while i<n { let d: i64=w[i]-t[i]; s=s+d*d; i=i+1 } return s } 40// grad of (w-t)^2 = 2(w-t), exact in Q16 (linear) 41func fill_grad(w: *i64, t: *i64, g: *i64, n: i64) -> i64 { var i: i64=0; while i<n { g[i]=2*(w[i]-t[i]); i=i+1 } return 0 } 42 43func main() -> i64 { 44 g_puts("nx_nofloat_adamw32 gate (Q32 moments FIX the Q16 AdamW g^2-underflow divergence)\n" as *u8) 45 let n: i64=8; let steps: i64=4000 46 let t: *i64 = sys_mmap(n*8) as *i64 47 t[0]=1966; t[1]=3277; t[2]=1311; t[3]=2621; t[4]=3932; t[5]=655; t[6]=2294; t[7]=2949 // small targets (~0.01-0.06) 48 let g: *i64 = sys_mmap(n*8) as *i64 49 50 // ---- Q16 AdamW (lib nfa_adamw): small gradients underflow g^2 ---- 51 let w16: *i64 = sys_mmap(n*8) as *i64; let m16: *i64 = sys_mmap(n*8) as *i64; let v16: *i64 = sys_mmap(n*8) as *i64 52 var z: i64=0; while z<n { w16[z]=0; m16[z]=0; v16[z]=0; z=z+1 } 53 var s1: i64=0 54 while s1<steps { fill_grad(w16,t,g,n); nfa_adamw(w16,g,m16,v16,n,655,58982,65470,66,0,s1+1); s1=s1+1 } 55 let l16: i64 = loss_q32(w16,t,n) 56 57 // ---- Q32 AdamW (the fix) ---- 58 let w32: *i64 = sys_mmap(n*8) as *i64; let m32: *i64 = sys_mmap(n*8) as *i64; let v32: *i64 = sys_mmap(n*8) as *i64 59 z=0; while z<n { w32[z]=0; m32[z]=0; v32[z]=0; z=z+1 } 60 var s2: i64=0 61 while s2<steps { fill_grad(w32,t,g,n); adamw32_one(w32,g,m32,v32,n,655,58982,65470,66,s2+1); s2=s2+1 } 62 let l32: i64 = loss_q32(w32,t,n) 63 64 g_puts(" [measure] final fit loss (sum (w-t)^2, Q32 raw; lower=better): Q16-AdamW=" as *u8); g_pn(l16); g_puts(" Q32-AdamW=" as *u8); g_pn(l32); g_puts("\n" as *u8) 65 66 var pass: i64=0; var total: i64=0 67 var t1: i64=0; if l32 < 30000000 { t1=1 } // Q32 converged: sum d^2 < 3e7 (Q32) ~= avg |w-t| < ~0.001 68 pass=pass+g_check("T1: Q32-moment AdamW CONVERGES on small-gradient fit (loss -> ~0)" as *u8, t1); total=total+1 69 var t2: i64=0; if l32*10 < l16 { t2=1 } // Q32 at least 10x better than Q16 = the underflow fix is real 70 pass=pass+g_check("T2: MEASURED EXCEED -- Q32 final loss << Q16 final loss (g^2-underflow fixed)" as *u8, t2); total=total+1 71 72 var okall: i64=0; if pass==total { okall=1 } 73 let logf: i64 = sys_openat_append("knowledge/status/nofloat_adamw32.log" as *u8, 420) 74 if logf >= 0 { let x0: i64=sys_write(logf,"NOFLOATADAMW32 q32-moments fix q16 adam underflow measured\n" as *u8,58); sys_close(logf) } 75 g_puts("---- adamw32 gate: passed " as *u8); g_pn(pass); g_puts(" / " as *u8); g_pn(total); g_puts(" ----\n" as *u8) 76 if okall==1 { g_puts("verdict=GREEN (Q32-moment AdamW: a numerically STABLE fixed-point optimizer; fixes the Q16 underflow)\n" as *u8); sys_exit(0); return 0 } 77 g_puts("verdict=RED\n" as *u8); sys_exit(1); return 1 78}