code wiki / _hdl_build / nx_f32_lm_train_gate.nx
nx_f32_lm_train_gate.nx source
↩ module page · 132 lines · 9034 B
1import "nx_gate_gn.nx"
2import "nx_gate_base.nx"
3// nx_f32_lm_train_gate.nx -- THE CAPSTONE: a real LM trained END-TO-END with the sovereign stack. Composes every rung
4// -- softmax (R2a) + cross-entropy (R3c) + the softmax-onehot gradient + Adam (R3d) -- into a working train loop on a
5// real (tiny) next-token task. A bigram model W[V][V] (logits for the next token given the current one) is trained from
6// ZERO on a cyclic sequence 0->1->2->3->0; it must LEARN the transitions (loss drops, argmax becomes correct). This is
7// R5 in miniature: it proves the whole machinery LEARNS, not just that each piece gradchecks. Sovereign throughout.
8// T0 INITIAL: untrained loss = log(V) per token (uniform softmax) = log(4) ~ 1.386.
9// T1 LOSS DROPS: training loss falls sharply over epochs.
10// T2 LOSS -> ~0: final loss << initial -- the model fit the data.
11// T3 LEARNED: for every token i, argmax(softmax(W[i])) == (i+1) mod V -- it predicts the next token correctly.
12// T4 CONFIDENT: p(correct next) > 0.9 -- not just argmax, but high probability.
13// T5 = end-to-end sovereign LM training WORKS (forward->loss->backward->Adam->weights->it learns).
14// license_tier: ORIGINAL
15import "nx_f32_hw.nx"
16import "nx_syscalls.nx"
17
18func grow(name: *u8, ok: i64) -> i64 { if ok==1 { gw(" PASS " as *u8) } else { gw(" FAIL " as *u8) } gw(name); gw("
19" as *u8); return ok }
20func gm(x: i64) -> i64 { return gn(f32_int(f32_mul(x, f32_of(1000)))) }
21func 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 }
22func f32_max2(a: i64, b: i64) -> i64 { if f32_le(a,b)==1 { return b } return a }
23func f32_sqrt(x: i64) -> i64 { if (x & 0x7FFFFFFF)==0 { return f32_of(0) } var y: i64=x; var i: i64=0; while i<16 { y=f32_div(f32_add(y, f32_div(x,y)), f32_of(2)); i=i+1 } return y }
24// RANGE-REDUCED f32_exp (the production version): exp(x)=2^(x*log2e)=2^n * 2^f, n=round(t), f in [-0.5,0.5] so the
25// Taylor for 2^f=exp(f*ln2) is on a tiny accurate range; 2^n built from the IEEE-754 exponent field. Handles large |x|
26// (training grows logits) -- the Taylor-only version diverged for |x|>~6, which the end-to-end train exposed.
27func f32_exp(x: i64) -> i64 {
28 let log2e: i64=f32_div(f32_of(1442695),f32_of(1000000)); let ln2: i64=f32_div(f32_of(693147),f32_of(1000000)); let half: i64=f32_div(f32_of(1),f32_of(2))
29 let t: i64=f32_mul(x, log2e)
30 var n: i64=0; if f32_le(f32_of(0), t)==1 { n=f32_int(f32_add(t,half)) } else { n=f32_int(f32_sub(t,half)) }
31 let arg: i64=f32_mul(f32_sub(t, f32_of(n)), ln2)
32 var p2f: i64=f32_of(1); var term: i64=f32_of(1); var k: i64=1
33 while k<=8 { term=f32_div(f32_mul(term,arg), f32_of(k)); p2f=f32_add(p2f,term); k=k+1 }
34 var ef: i64=n+127
35 if ef<=0 { return f32_of(0) }
36 if ef>=255 { ef=254 }
37 return f32_mul(p2f, (ef & 0xFF) << 23)
38}
39func f32_log(x: i64) -> i64 { let b: i64=x & 0xFFFFFFFF; let e: i64=((b>>23)&0xFF)-127; let m: i64=(b & 0x7FFFFF)|0x3F800000; let u: i64=f32_div(f32_sub(m,f32_of(1)),f32_add(m,f32_of(1))); let u2: i64=f32_mul(u,u); var t: i64=u; var s: i64=u; var k: i64=1; while k<=7 { t=f32_mul(t,u2); s=f32_add(s,f32_div(t,f32_of((2*k)+1))); k=k+1 } let ln2: i64=f32_div(f32_of(693147),f32_of(1000000)); return f32_add(f32_mul(f32_of(e),ln2),f32_mul(f32_of(2),s)) }
40
41func softmax(z: *i64, n: i64, out: *i64) -> i64 {
42 var mx: i64=z[0]; var i: i64=1; while i<n { mx=f32_max2(mx,z[i]); i=i+1 }
43 var sum: i64=f32_of(0); i=0; while i<n { out[i]=f32_exp(f32_sub(z[i],mx)); sum=f32_add(sum,out[i]); i=i+1 }
44 i=0; while i<n { out[i]=f32_div(out[i],sum); i=i+1 }
45 return 0
46}
47func argmax(p: *i64, n: i64) -> i64 { var bi: i64=0; var i: i64=1; while i<n { if f32_le(p[bi],p[i])==1 { bi=i } i=i+1 } return bi }
48
49func main() -> i64 {
50 gw("=== nx_f32_lm_train_gate: CAPSTONE -- a real bigram LM trained END-TO-END with the sovereign stack ===\n" as *u8)
51 var pass: i64=0; var total: i64=0
52 let V: i64=4
53 let b1: i64=f32_div(f32_of(9),f32_of(10)); let b2: i64=f32_div(f32_of(999),f32_of(1000))
54 let lr: i64=f32_div(f32_of(1),f32_of(10)); let eps: i64=f32_div(f32_of(1),f32_of(100000000)); let one: i64=f32_of(1)
55 let W: *i64=sys_mmap(8*64) as *i64; let m: *i64=sys_mmap(8*64) as *i64; let vv: *i64=sys_mmap(8*64) as *i64
56 var i: i64=0; while i<V*V { W[i]=f32_of(0); m[i]=f32_of(0); vv[i]=f32_of(0); i=i+1 }
57 let p: *i64=sys_mmap(64) as *i64; let logits: *i64=sys_mmap(64) as *i64
58 var b1t: i64=one; var b2t: i64=one
59
60 // loss over the 4 transitions i -> (i+1)%V.
61 var loss_start: i64=f32_of(0); var loss_end: i64=f32_of(0)
62 var epoch: i64=1
63 while epoch<=400 {
64 var eloss: i64=f32_of(0)
65 var inp: i64=0
66 while inp<V {
67 let tgt: i64=(inp+1)%V
68 var j: i64=0; while j<V { logits[j]=W[inp*V+j]; j=j+1 }
69 softmax(logits, V, p)
70 eloss=f32_add(eloss, f32_neg(f32_log(p[tgt]))) // CE for this transition
71 b1t=f32_mul(b1t,b1); b2t=f32_mul(b2t,b2)
72 j=0
73 while j<V {
74 var g: i64=p[j]; if j==tgt { g=f32_sub(p[j],one) } // softmax - onehot
75 let idx: i64=inp*V+j
76 m[idx]=f32_add(f32_mul(b1,m[idx]), f32_mul(f32_sub(one,b1),g))
77 vv[idx]=f32_add(f32_mul(b2,vv[idx]), f32_mul(f32_sub(one,b2),f32_mul(g,g)))
78 let mhat: i64=f32_div(m[idx],f32_sub(one,b1t)); let vhat: i64=f32_div(vv[idx],f32_sub(one,b2t))
79 W[idx]=f32_sub(W[idx], f32_div(f32_mul(lr,mhat), f32_add(f32_sqrt(vhat),eps)))
80 j=j+1
81 }
82 inp=inp+1
83 }
84 if epoch==1 { loss_start=eloss }
85 loss_end=eloss
86 epoch=epoch+1
87 }
88
89 // T0 initial loss ~ V*log(V) = 4*1.386 = 5.545.
90 let lsi: i64=f32_int(f32_mul(loss_start,f32_of(1000)))
91 total=total+1; if lsi>=5400 { if lsi<=5700 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } } else { gw(" [FAIL] " as *u8) }
92 gw("T0 INITIAL: untrained loss=" as *u8); gn(lsi); gw("m ~= 4*log(4)=5545m (uniform softmax -- the model knows nothing yet)\n" as *u8)
93
94 // T1 + T2 loss dropped to ~0.
95 let lei: i64=f32_int(f32_mul(loss_end,f32_of(1000)))
96 total=total+1; if f32_le(loss_end,loss_start)==1 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) }
97 gw("T1 LOSS DROPS: " as *u8); gn(lsi); gw("m -> " as *u8); gn(lei); gw("m over 400 epochs\n" as *u8)
98 total=total+1; if lei>=0 { if lei<=200 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } } else { gw(" [FAIL] " as *u8) }
99 gw("T2 LOSS -> ~0: final loss=" as *u8); gn(lei); gw("m (<< initial) -- the model fit the transitions\n" as *u8)
100
101 // T3 LEARNED: argmax(softmax(W[i])) == (i+1)%V for all i.
102 var learned: i64=1; i=0
103 while i<V {
104 var j: i64=0; while j<V { logits[j]=W[i*V+j]; j=j+1 }
105 softmax(logits, V, p)
106 let pred: i64=argmax(p, V); let want: i64=(i+1)%V
107 gw(" token " as *u8); gn(i); gw(" -> predicts " as *u8); gn(pred); gw(" (want " as *u8); gn(want); gw("), p=" as *u8); gm(p[want]); gw("m\n" as *u8)
108 if pred!=want { learned=0 }
109 i=i+1
110 }
111 total=total+1; if learned==1 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) }
112 gw("T3 LEARNED: every token predicts its correct next token (argmax) -- the LM learned the pattern from zero\n" as *u8)
113
114 // T4 CONFIDENT: p(correct) > 0.9 for token 0.
115 var j2: i64=0; while j2<V { logits[j2]=W[0*V+j2]; j2=j2+1 }
116 softmax(logits, V, p)
117 let pc: i64=f32_int(f32_mul(p[1],f32_of(1000)))
118 total=total+1; if pc>=900 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) }
119 gw("T4 CONFIDENT: p(token0 -> token1)=" as *u8); gn(pc); gw("m (>900 = >0.9 probability, not just argmax)\n" as *u8)
120
121 // T5.
122 total=total+1; if learned==1 { if lei<=200 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } } else { gw(" [FAIL] " as *u8) }
123 gw("T5 END-TO-END TRAIN WORKS: forward->CE loss->softmax-onehot gradient->Adam->weights updated->the model LEARNED\n" as *u8)
124
125 gw("\n CAPSTONE: a real LM trained from ZERO to correct predictions, entirely on the sovereign stack -- softmax + cross-entropy +\n" as *u8)
126 gw(" Adam composed into a working train loop. This is R5 in miniature: every rung (R1 autograd .. R3d Adam) COMPOSES into a\n" as *u8)
127 gw(" machine that LEARNS. To scale to the 0.5-1B: same loop over a real transformer (GQA-wired blocks) + the crawl-on-burst\n" as *u8)
128 gw(" tokenized corpus, on the RTX 5080 (~1wk) or 1 H100 (~1day). The math + optimizer + loop are PROVEN sovereign.\n" as *u8)
129 gw("LM-TRAIN verdict=" as *u8)
130 if pass==total { gw("GREEN passes=" as *u8); gn(pass); gw("/" as *u8); gn(total); gw(" -- end-to-end sovereign LM training works (a model learned from zero)\n" as *u8); sys_exit(0); return 0 }
131 gw("RED passes=" as *u8); gn(pass); gw("/" as *u8); gn(total); gw("\n" as *u8); sys_exit(1); return 1
132}