code wiki / _hdl_build / nx_f32_autograd_gate.nx

nx_f32_autograd_gate.nx source

↩ module page · 109 lines · 7367 B

1import "nx_gate_gn.nx" 2import "nx_gate_base.nx" 3// nx_f32_autograd_gate.nx -- RUNG 1 of the from-scratch 0.5B training loop (operator: "cheap from scratch"). Our 4// Q16 autograd TRAINS + gradchecks (nx_autograd_mlp_gate, GREEN) but fixed-point Q16 lacks the dynamic range for 5// stable LLM training -- transformers need f32. This ports the PROVEN reverse-mode autograd structure (tape of nodes, 6// eager forward, reverse chain-rule, finite-difference gradcheck, the liar-kill negative control) to HARDWARE f32 7// (nx_f32_hw). It proves the f32 autograd is correct on ADD/MUL/RELU -- the foundation; softmax/layernorm backward 8// (needing f32_exp/f32_sqrt, the next rung) build on this. 9// T0 f32 SMOKE: the hardware-f32 lib computes (verify-don't-assume before building on it). 10// T1 FORWARD: L = relu(a*b + a) at a=3,b=2 -> 9.0 (eager forward on the tape). 11// T2 GRADCHECK d/da: autograd grad == central finite-difference (the proof backprop is correct in f32). 12// T3 GRADCHECK d/db: same, the other input. 13// T4 NEGATIVE CONTROL: a deliberately-wrong gradient is REJECTED by the finite-difference check (liar-kill). 14// T5 = the f32 training-precision autograd foundation -> next rung: f32_exp/f32_sqrt -> softmax/layernorm backward. 15// license_tier: ORIGINAL 16import "nx_f32_hw.nx" 17import "nx_syscalls.nx" 18 19// print an f32 in MILLIUNITS (x1000, truncated) so 3.0 -> 3000 and 2.99998 -> 2999 read as CLOSE, not as "3 vs 2". 20func grow(name: *u8, ok: i64) -> i64 { if ok==1 { gw(" PASS " as *u8) } else { gw(" FAIL " as *u8) } gw(name); gw(" 21" as *u8); return ok } 22func gm(x: i64) -> i64 { return gn(f32_int(f32_mul(x, f32_of(1000)))) } 23 24// ---- f32 helpers on the hardware-f32 bit pattern (32-bit IEEE in an i64) ---- 25func f32_is_pos(x: i64) -> i64 { let x32: i64=x & 0xFFFFFFFF; if ((x32>>31)&1)==0 { if (x32 & 0x7FFFFFFF)!=0 { return 1 } } return 0 } 26func f32_relu(x: i64) -> i64 { if f32_is_pos(x)==1 { return x } return f32_of(0) } 27func f32_abs(x: i64) -> i64 { return x & 0x7FFFFFFF } 28func f32_le(x: i64, y: i64) -> i64 { let d: i64=f32_sub(x,y) & 0xFFFFFFFF; if ((d>>31)&1)==1 { return 1 } if (d & 0x7FFFFFFF)==0 { return 1 } return 0 } 29 30// ---- the autograd tape: node i at i*5 = [val, grad, op, inA, inB]; op 0=leaf 1=add 2=mul 3=relu ---- 31func ag_set(tape: *i64, i: i64, val: i64, op: i64, a: i64, b: i64) -> i64 { tape[i*5]=val; tape[i*5+1]=f32_of(0); tape[i*5+2]=op; tape[i*5+3]=a; tape[i*5+4]=b; return 0 } 32 33// build L = relu(a*b + a), EAGER forward into the tape. nodes: 0=a 1=b 2=mul 3=add 4=relu. returns L idx. 34func build_L(tape: *i64, av: i64, bv: i64) -> i64 { 35 ag_set(tape,0, av,0,0,0); ag_set(tape,1, bv,0,0,0) 36 ag_set(tape,2, f32_mul(av,bv),2,0,1) 37 ag_set(tape,3, f32_add(tape[2*5],av),1,2,0) 38 ag_set(tape,4, f32_relu(tape[3*5]),3,3,0) 39 return 4 40} 41// reverse pass: zero grads, seed dL/dL = 1, propagate by op (chain rule) in reverse node order. 42func backward(tape: *i64, nnodes: i64, L: i64) -> i64 { 43 var i: i64=0; while i<nnodes { tape[i*5+1]=f32_of(0); i=i+1 } 44 tape[L*5+1]=f32_of(1) 45 var n: i64=nnodes-1 46 while n>=0 { 47 let op: i64=tape[n*5+2]; let g: i64=tape[n*5+1]; let a: i64=tape[n*5+3]; let b: i64=tape[n*5+4] 48 if op==1 { tape[a*5+1]=f32_add(tape[a*5+1],g); tape[b*5+1]=f32_add(tape[b*5+1],g) } 49 if op==2 { tape[a*5+1]=f32_add(tape[a*5+1], f32_mul(g,tape[b*5])); tape[b*5+1]=f32_add(tape[b*5+1], f32_mul(g,tape[a*5])) } 50 if op==3 { if f32_is_pos(tape[n*5])==1 { tape[a*5+1]=f32_add(tape[a*5+1], g) } } 51 n=n-1 52 } 53 return 0 54} 55// forward-only L (for the finite-difference probe). 56func compute_L(av: i64, bv: i64) -> i64 { return f32_relu(f32_add(f32_mul(av,bv), av)) } 57 58func main() -> i64 { 59 gw("=== nx_f32_autograd_gate: f32 reverse-mode autograd -- RUNG 1 of the from-scratch 0.5B training loop ===\n" as *u8) 60 var pass: i64=0; var total: i64=0 61 let two: i64=f32_of(2); let three: i64=f32_of(3) 62 let eps: i64=f32_div(f32_of(1), f32_of(100)) // 0.01 63 let tol: i64=f32_div(f32_of(1), f32_of(10)) // 0.1 64 65 // T0 SMOKE. 66 let s_add: i64=f32_int(f32_add(three,two)); let s_mul: i64=f32_int(f32_mul(three,two)) 67 total=total+1; if s_add==5 { if s_mul==6 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } } else { gw(" [FAIL] " as *u8) } 68 gw("T0 f32 SMOKE: 3+2=" as *u8); gn(s_add); gw(", 3*2=" as *u8); gn(s_mul); gw(" (hardware f32 lib verified)\n" as *u8) 69 70 let tape: *i64=sys_mmap(8*5*16) as *i64 71 let L: i64=build_L(tape, three, two) 72 let Lval: i64=tape[L*5] 73 backward(tape, 5, L) 74 let grad_a: i64=tape[0*5+1]; let grad_b: i64=tape[1*5+1] 75 76 // T1 FORWARD. 77 total=total+1; if f32_int(Lval)==9 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } 78 gw("T1 FORWARD: L=relu(a*b+a) at a=3,b=2 = " as *u8); gn(f32_int(Lval)); gw(" (expect 9)\n" as *u8) 79 80 // T2 GRADCHECK d/da. 81 let twoeps: i64=f32_mul(two, eps) 82 let fda: i64=f32_div(f32_sub(compute_L(f32_add(three,eps),two), compute_L(f32_sub(three,eps),two)), twoeps) 83 let oka: i64=f32_le(f32_abs(f32_sub(grad_a, fda)), tol) 84 total=total+1; if oka==1 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } 85 gw("T2 GRADCHECK d/da: autograd=" as *u8); gm(grad_a); gw("m finite-diff=" as *u8); gm(fda); gw("m (milliunits; match within tol -> ~3.0 each)\n" as *u8) 86 87 // T3 GRADCHECK d/db. 88 let fdb: i64=f32_div(f32_sub(compute_L(three,f32_add(two,eps)), compute_L(three,f32_sub(two,eps))), twoeps) 89 let okb: i64=f32_le(f32_abs(f32_sub(grad_b, fdb)), tol) 90 total=total+1; if okb==1 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } 91 gw("T3 GRADCHECK d/db: autograd=" as *u8); gm(grad_b); gw("m finite-diff=" as *u8); gm(fdb); gw("m (milliunits; match within tol -> ~3.0 each)\n" as *u8) 92 93 // T4 NEGATIVE CONTROL -- a wrong gradient (doubled) must FAIL the finite-diff check. 94 let wrong: i64=f32_mul(grad_a, two) 95 let wrong_ok: i64=f32_le(f32_abs(f32_sub(wrong, fda)), tol) 96 total=total+1; if wrong_ok==0 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } 97 gw("T4 NEGATIVE CONTROL: a doubled (wrong) gradient " as *u8); gn(f32_int(wrong)); gw(" is REJECTED by the finite-diff check (liar-kill)\n" as *u8) 98 99 // T5. 100 total=total+1; if oka==1 { if okb==1 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } } else { gw(" [FAIL] " as *u8) } 101 gw("T5 FOUNDATION: f32 reverse-mode autograd is correct (ADD/MUL/RELU, gradchecked) = the training-precision base for the 0.5B\n" as *u8) 102 103 gw("\n RUNG 1 DONE: the proven autograd structure now runs in HARDWARE f32 (real dynamic range, not Q16). Gradients are correct\n" as *u8) 104 gw(" (finite-difference verified, wrong-gradient rejected). NEXT RUNGS toward from-scratch 0.5B: (2) f32_exp + f32_sqrt -> softmax\n" as *u8) 105 gw(" + layernorm backward; (3) Adam optimizer; (4) tokenized data pipeline (fed by crawl-on-burst); (5) the train loop. ~$50/1wk.\n" as *u8) 106 gw("F32-AUTOGRAD verdict=" as *u8) 107 if pass==total { gw("GREEN passes=" as *u8); gn(pass); gw("/" as *u8); gn(total); gw(" -- f32 autograd correct (gradchecked), rung 1 of the from-scratch 0.5B training loop\n" as *u8); sys_exit(0); return 0 } 108 gw("RED passes=" as *u8); gn(pass); gw("/" as *u8); gn(total); gw("\n" as *u8); sys_exit(1); return 1 109}