nx_intfp_softmax_gradcheck_gate.nx source

1// nx_intfp_softmax_gradcheck_gate.nx -- KEYSTONE op for integer ATTENTION training: fixed-point SOFTMAX with a 2// fixed-point EXP and the full Jacobian backward, gradchecked (integer finite-difference). exp is the crux -- 3// it's exactly where the software-float tape calls nx_f32_exp. Here exp is done in Q16 via the 2^x decomposition 4// (exp(x)=2^(x*log2e); split into shift 2^-ui and a quadratic for 2^-uf), NO float. Then softmax s_i=e_i/sum, 5// loss L=sum(s_i^2), analytic backward dL/dx_i = s_i*(g_i - sum_j s_j*g_j) with g_i=2*s_i. The max-subtraction is 6// a stop-gradient constant (softmax is shift-invariant), so we fix the shift at the base point -> the forward is 7// smooth and finite-diff matches analytic. If this gradchecks, integer attention is trainable. license_tier: ORIGINAL 8import "nx_syscalls.nx" 9 10func w(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 } 11func wn(v: i64) -> i64 { if v==0 { sys_write(1,"0" as *u8,1); return 0 } var m: i64=v; if m<0{sys_write(1,"-" as *u8,1);m=0-m} let t: *u8=sys_mmap(24); var k: i64=0; while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} let o: *u8=sys_mmap(24); var q: i64=k-1; var i: i64=0; while q>=0{o[i]=t[q];i=i+1;q=q-1} sys_write(1,o,i); return 0 } 12func iabs(v: i64) -> i64 { if v<0 { return 0-v } return v } 13 14const S: i64 = 65536 15const N: i64 = 5 16 17// Q16 exp for ANY xq (both signs). exp(x)=2^y, y=x*log2e; y=yi+yf with yi=floor(y), yf in [0,1); 2^y=(2^yf)<<yi. 18// 2^yf via QUARTIC Horner (accurate in VALUE and DERIVATIVE -- the derivative matters for the softmax gradcheck). 19func fp_exp_q16(xq: i64) -> i64 { 20 let y: i64=(xq*94548)/S // y = x*1.442695 (log2 e in Q16 = 94548), signed 21 var yi: i64=0 22 if y>=0 { yi=y/S } else { yi=0-(((0-y)+S-1)/S) } // floor(y/S) (nx div truncates toward 0) 23 let yf: i64=y-yi*S // [0,S) 24 var p: i64=630 // a4=0.009618 (Q16 quartic coeffs of 2^f = e^{f ln2}) 25 p=3637+(p*yf)/S // a3=0.055504 26 p=15743+(p*yf)/S // a2=0.240227 27 p=45426+(p*yf)/S // a1=0.693147 28 p=S+(p*yf)/S // 2^yf, Q16 in [S,2S) 29 if yi>=0 { if yi>=31 { return 2000000000 } return p*(1<<yi) } 30 let k: i64=0-yi 31 if k>=31 { return 0 } // underflow -> ~0 32 return p/(1<<k) 33} 34 35// softmax with FIXED shift; fills s (Q16); returns L_q32 = sum(s_i^2) 36func softmax_loss(x: *i64, shift: i64, s: *i64) -> i64 { 37 let e: *i64=sys_mmap(N*8) as *i64 38 var sum: i64=0; var i: i64=0 39 while i<N { let ev: i64=fp_exp_q16(x[i]-shift); e[i]=ev; sum=sum+ev; i=i+1 } 40 var Lq: i64=0; i=0 41 while i<N { let sv: i64=(e[i]*S+sum/2)/sum; s[i]=sv; Lq=Lq+sv*sv; i=i+1 } // rounded division 42 return Lq 43} 44 45func main() -> i64 { 46 w("=== nx_intfp_softmax_gradcheck: Q16 fixed-point softmax + exp + Jacobian backward -- no float ===\n\n" as *u8) 47 // (0) verify the fixed-point exp before building on it (rule: verify, don't assume) 48 w(" [exp check] exp(0)=" as *u8); wn(fp_exp_q16(0)); w(" (want 65536) exp(-1)=" as *u8); wn(fp_exp_q16(0-S)); w(" (want ~24109=0.368) exp(-2)=" as *u8); wn(fp_exp_q16(0-2*S)); w(" (want ~8867=0.135)\n\n" as *u8) 49 50 let x: *i64=sys_mmap(N*8) as *i64; let s: *i64=sys_mmap(N*8) as *i64 51 var i: i64=0; while i<N { x[i]=((i*2-3)*S)/4; i=i+1 } // -0.75,-0.25,0.25,0.75,1.25 52 var shift: i64=x[0]; i=1; while i<N { if x[i]>shift { shift=x[i] } i=i+1 } // max, fixed (stop-grad) 53 54 let L0: i64=softmax_loss(x, shift, s) 55 w(" softmax s_q=[" as *u8); i=0; while i<N { wn(s[i]); if i<N-1 { w("," as *u8) } i=i+1 } w("] (sum~65536) L_q32=" as *u8); wn(L0); w("\n\n" as *u8) 56 57 // analytic: g_i=2 s_i ; dot=sum_j s_j g_j ; dL/dx_i = s_i (g_i - dot) (all Q16) 58 let g: *i64=sys_mmap(N*8) as *i64; i=0; while i<N { g[i]=2*s[i]; i=i+1 } 59 var dot: i64=0; i=0; while i<N { dot=dot+(s[i]*g[i])/S; i=i+1 } 60 let ana: *i64=sys_mmap(N*8) as *i64; i=0; while i<N { ana[i]=(s[i]*(g[i]-dot))/S; i=i+1 } 61 62 let DELTA: i64=655; let TOLP: i64=80 // delta=0.01 (large enough that softmax-division truncation noise is negligible) 63 var npass: i64=0; var worst: i64=0 64 w(" i analytic_q numeric_q rel(permille)\n" as *u8) 65 w(" ------------------------------------------------\n" as *u8) 66 i=0 67 while i<N { 68 let save: i64=x[i] 69 x[i]=save+DELTA; let Lp: i64=softmax_loss(x, shift, s) 70 x[i]=save-DELTA; let Lm: i64=softmax_loss(x, shift, s) 71 x[i]=save; let d: i64=softmax_loss(x, shift, s) 72 let num: i64=(Lp-Lm)/(2*DELTA) 73 let rel: i64=(iabs(num-ana[i])*1000)/(iabs(ana[i])+100) 74 w(" " as *u8); wn(i); w(" " as *u8); wn(ana[i]); w(" " as *u8); wn(num); w(" " as *u8); wn(rel) 75 if rel<=TOLP { npass=npass+1; w(" ok\n" as *u8) } else { w(" FAIL\n" as *u8) } 76 if rel>worst { worst=rel } 77 i=i+1 78 } 79 w("\n softmax gradcheck: " as *u8); wn(npass); w("/" as *u8); wn(N); w(" within " as *u8); wn(TOLP); w(" permille; worst=" as *u8); wn(worst); w("\n" as *u8) 80 w(" => fixed-point softmax (with fixed-point exp) forward+backward is sound -> integer ATTENTION is trainable.\n" as *u8) 81 w("NX-INTFP-SOFTMAX-GRADCHECK verdict=" as *u8) 82 if npass==N { w("GREEN passes=" as *u8); wn(npass); w("/" as *u8); wn(N); w(" -- integer softmax+exp+Jacobian proven; the attention keystone op is buildable in the tape\n" as *u8) } 83 else { w("RED passes=" as *u8); wn(npass); w("/" as *u8); wn(N); w(" -- exp accuracy or Jacobian bug\n" as *u8) } 84 return 0 85}