code wiki / _hdl_build / nx_fnet_model_gate.nx

nx_fnet_model_gate.nx source

↩ module page · 293 lines · 13669 B

1// nx_fnet_model_gate.nx -- GATE for MODEL-001: the FIRST end-to-end sovereign SUB-QUADRATIC MODEL. A real 2// sequence classifier with no attention anywhere: 3// tokens -> EMBED (trained) -> FNET token-mix (sub-quadratic) -> relu FFN -> linear -> softmax-CE 4// Trained by AdamW JOINTLY over the embedding table + FFN/head weights. The embedding has no attention matrix 5// and no quadratic cost -- token mixing is the parameter-free Fourier transform. Embeddings are trained without 6// a dedicated op: the tape yields dL/dx at the input leaf, which is SCATTERED back into the embedding rows. 7// 8// Task: classify whether token[0] > token[3] over 8 length-4 sequences (vocab 4). This genuinely needs the FNet 9// to MIX positions 0 and 3 and the relu to compare -- a bag-of-words linear model cannot do it. 10// 11// G_train the model learns: accuracy >= 7/8 AND final loss < first loss. 12// G_repro bit-exact: train twice -> identical accuracy + final loss bits. 13// 14// Evidence -> knowledge/status/fnet_model.log (FNETMODELGATE authored=organ ... verdict=GREEN). license_tier: ORIGINAL 15import "nx_autograd_tensor.nx" // ta_* incl ta_fnet + transitively the f32 tower + fft_f32 16import "nx_syscalls.nx" 17import "nx_gate_verdict.nx" 18 19const MV: i64 = 4 // vocab 20const MN: i64 = 4 // sequence length 21const MD: i64 = 4 // d_model 22const MH: i64 = 8 // FFN hidden 23const MC: i64 = 2 // classes 24const MND: i64 = 16 // MN*MD (flattened FNet output) 25 26const MO_LOG: *u8 = "knowledge/status/fnet_model.log" 27 28func mo_w(fd: i64, s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(fd, s, n); return 0 } 29func mo_wn(fd: i64, v: i64) -> i64 { 30 let bb: *u8 = sys_mmap(28); var m: i64 = v 31 if m < 0 { m = 0 - m; sys_write(fd, "-" as *u8, 1) } 32 let t: *u8 = sys_mmap(28); var k: i64 = 0 33 if m == 0 { t[0] = 48; k = 1 } 34 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 } 35 var i: i64 = 0 36 while i < k { bb[i] = t[k - 1 - i]; i = i + 1 } 37 sys_write(fd, bb, k); return 0 38} 39 40// build x[MN*MD] by gathering embedding rows for the token sequence seq[MN] 41func mo_embed(E: *i64, seq: *i64, soff: i64, xout: *i64) -> i64 { 42 var i: i64 = 0 43 while i < MN { 44 let tok: i64 = seq[soff + i] 45 var j: i64 = 0 46 while j < MD { xout[i * MD + j] = E[tok * MD + j]; j = j + 1 } 47 i = i + 1 48 } 49 return 0 50} 51 52// one forward example -> logits node. caller has already placed embeddings into x. 53func mo_fwd(tape: *i64, vals: *i64, st: *i64, x: *i64, nW1: i64, nb1: i64, nW2: i64, nb2: i64) -> i64 { 54 let xl: i64 = ta_leaf(tape, vals, st, MN, MD, x, 0) 55 let m: i64 = ta_fnet(tape, vals, st, xl) 56 let h: i64 = ta_relu(tape, vals, st, ta_vadd(tape, vals, st, ta_matvec(tape, vals, st, nW1, m), nb1)) 57 let lo: i64 = ta_vadd(tape, vals, st, ta_matvec(tape, vals, st, nW2, h), nb2) 58 return lo 59} 60 61// full-batch forward over 8 examples; writes shared-leaf indices wb[0..3] and per-example input-leaf idx xl8[0..7]. 62func mo_build(tape: *i64, vals: *i64, st: *i64, E: *i64, W1: *i64, b1: *i64, W2: *i64, b2: *i64, 63 seqs: *i64, labels: *i64, c8: *i64, wb: *i64, xl8: *i64) -> i64 { 64 st[0] = 0; st[1] = 0 65 let nW1: i64 = ta_leaf(tape, vals, st, MH, MND, W1, 0) 66 let nb1: i64 = ta_leaf(tape, vals, st, MH, 1, b1, 0) 67 let nW2: i64 = ta_leaf(tape, vals, st, MC, MH, W2, 0) 68 let nb2: i64 = ta_leaf(tape, vals, st, MC, 1, b2, 0) 69 wb[0] = nW1; wb[1] = nb1; wb[2] = nW2; wb[3] = nb2 70 let x: *i64 = (sys_mmap(MND * 8)) as *i64 71 let th: *i64 = (sys_mmap(MC * 8)) as *i64 72 var sumn: i64 = 0 - 1 73 var s: i64 = 0 74 while s < 8 { 75 mo_embed(E, seqs, s * MN, x) 76 let xl: i64 = ta_leaf(tape, vals, st, MN, MD, x, 0) 77 xl8[s] = xl 78 let m: i64 = ta_fnet(tape, vals, st, xl) 79 let h: i64 = ta_relu(tape, vals, st, ta_vadd(tape, vals, st, ta_matvec(tape, vals, st, nW1, m), nb1)) 80 let lo: i64 = ta_vadd(tape, vals, st, ta_matvec(tape, vals, st, nW2, h), nb2) 81 var j: i64 = 0 82 while j < MC { th[j] = TA_F32_ZERO; j = j + 1 } 83 th[labels[s]] = TA_F32_ONE 84 let tgt: i64 = ta_leaf(tape, vals, st, MC, 1, th, 0) 85 let ls: i64 = ta_softce(tape, vals, st, lo, tgt) 86 if s == 0 { sumn = ls } else { sumn = ta_vadd(tape, vals, st, sumn, ls) } 87 s = s + 1 88 } 89 let inv8: i64 = ta_leaf(tape, vals, st, 1, 1, c8, 0) 90 return ta_matvec(tape, vals, st, inv8, sumn) 91} 92 93func mo_predict(tape: *i64, vals: *i64, st: *i64, E: *i64, W1: *i64, b1: *i64, W2: *i64, b2: *i64, seqs: *i64, s: i64) -> i64 { 94 st[0] = 0; st[1] = 0 95 let nW1: i64 = ta_leaf(tape, vals, st, MH, MND, W1, 0) 96 let nb1: i64 = ta_leaf(tape, vals, st, MH, 1, b1, 0) 97 let nW2: i64 = ta_leaf(tape, vals, st, MC, MH, W2, 0) 98 let nb2: i64 = ta_leaf(tape, vals, st, MC, 1, b2, 0) 99 let x: *i64 = (sys_mmap(MND * 8)) as *i64 100 mo_embed(E, seqs, s * MN, x) 101 let lo: i64 = mo_fwd(tape, vals, st, x, nW1, nb1, nW2, nb2) 102 if nx_f32_gt(ta_val(tape, vals, lo, 1), ta_val(tape, vals, lo, 0)) == 1 { return 1 } 103 return 0 104} 105 106// AdamW update of n params in place 107func mo_adamw(p: *i64, m: *i64, v: *i64, g: *i64, n: i64, lr: i64, beta1: i64, beta2: i64, om1: i64, om2: i64, eps: i64, c1: i64, c2: i64) -> i64 { 108 var i: i64 = 0 109 while i < n { 110 let gi: i64 = g[i] 111 m[i] = nx_f32_add(nx_f32_mul(beta1, m[i]), nx_f32_mul(om1, gi)) 112 v[i] = nx_f32_add(nx_f32_mul(beta2, v[i]), nx_f32_mul(om2, nx_f32_mul(gi, gi))) 113 let mhat: i64 = nx_f32_div(m[i], c1) 114 let vhat: i64 = nx_f32_div(v[i], c2) 115 p[i] = nx_f32_sub(p[i], nx_f32_mul(lr, nx_f32_div(mhat, nx_f32_add(nx_f32_sqrt(vhat), eps)))) 116 i = i + 1 117 } 118 return 0 119} 120 121// train the whole model. writes final E/W1/b1/W2/b2 in place; returns nothing (params are the outputs). 122func mo_train(tape: *i64, vals: *i64, grads: *i64, st: *i64, E: *i64, W1: *i64, b1: *i64, W2: *i64, b2: *i64, 123 seqs: *i64, labels: *i64, epochs: i64, lf: *i64, ll: *i64) -> i64 { 124 // init params (deterministic symmetry-breaking) 125 ta_det_init(E, MV * MD, 5) 126 ta_det_init(W1, MH * MND, 3) 127 ta_det_init(W2, MC * MH, 7) 128 var z: i64 = 0 129 while z < MH { b1[z] = TA_F32_ZERO; z = z + 1 } 130 z = 0 131 while z < MC { b2[z] = TA_F32_ZERO; z = z + 1 } 132 // adamw state 133 let mE: *i64 = (sys_mmap(MV * MD * 8)) as *i64; let vE: *i64 = (sys_mmap(MV * MD * 8)) as *i64 134 let mW1: *i64 = (sys_mmap(MH * MND * 8)) as *i64; let vW1: *i64 = (sys_mmap(MH * MND * 8)) as *i64 135 let mb1: *i64 = (sys_mmap(MH * 8)) as *i64; let vb1: *i64 = (sys_mmap(MH * 8)) as *i64 136 let mW2: *i64 = (sys_mmap(MC * MH * 8)) as *i64; let vW2: *i64 = (sys_mmap(MC * MH * 8)) as *i64 137 let mb2: *i64 = (sys_mmap(MC * 8)) as *i64; let vb2: *i64 = (sys_mmap(MC * 8)) as *i64 138 z = 0 139 while z < MV * MD { mE[z] = TA_F32_ZERO; vE[z] = TA_F32_ZERO; z = z + 1 } 140 z = 0 141 while z < MH * MND { mW1[z] = TA_F32_ZERO; vW1[z] = TA_F32_ZERO; z = z + 1 } 142 z = 0 143 while z < MH { mb1[z] = TA_F32_ZERO; vb1[z] = TA_F32_ZERO; z = z + 1 } 144 z = 0 145 while z < MC * MH { mW2[z] = TA_F32_ZERO; vW2[z] = TA_F32_ZERO; z = z + 1 } 146 z = 0 147 while z < MC { mb2[z] = TA_F32_ZERO; vb2[z] = TA_F32_ZERO; z = z + 1 } 148 149 let beta1: i64 = ta_constf(9, 10); let beta2: i64 = ta_constf(999, 1000) 150 let om1: i64 = ta_constf(1, 10); let om2: i64 = ta_constf(1, 1000) 151 let lr: i64 = ta_constf(1, 50); let eps: i64 = ta_constf(1, 100000000) 152 var b1t: i64 = TA_F32_ONE; var b2t: i64 = TA_F32_ONE 153 154 let c8: *i64 = (sys_mmap(8)) as *i64; c8[0] = ta_constf(1, 8) 155 let wb: *i64 = (sys_mmap(4 * 8)) as *i64 156 let xl8: *i64 = (sys_mmap(8 * 8)) as *i64 157 let gW1: *i64 = (sys_mmap(MH * MND * 8)) as *i64 158 let gb1: *i64 = (sys_mmap(MH * 8)) as *i64 159 let gW2: *i64 = (sys_mmap(MC * MH * 8)) as *i64 160 let gb2: *i64 = (sys_mmap(MC * 8)) as *i64 161 let dE: *i64 = (sys_mmap(MV * MD * 8)) as *i64 162 163 var ep: i64 = 0 164 while ep < epochs { 165 let loss: i64 = mo_build(tape, vals, st, E, W1, b1, W2, b2, seqs, labels, c8, wb, xl8) 166 ta_backward(tape, vals, grads, st[0], loss) 167 if ep == 0 { *lf = ta_val(tape, vals, loss, 0) } 168 *ll = ta_val(tape, vals, loss, 0) 169 // read shared-leaf grads 170 var i: i64 = 0 171 while i < MH * MND { gW1[i] = ta_grad(tape, grads, wb[0], i); i = i + 1 } 172 i = 0 173 while i < MH { gb1[i] = ta_grad(tape, grads, wb[1], i); i = i + 1 } 174 i = 0 175 while i < MC * MH { gW2[i] = ta_grad(tape, grads, wb[2], i); i = i + 1 } 176 i = 0 177 while i < MC { gb2[i] = ta_grad(tape, grads, wb[3], i); i = i + 1 } 178 // scatter input-leaf grads into the embedding table 179 i = 0 180 while i < MV * MD { dE[i] = TA_F32_ZERO; i = i + 1 } 181 var s: i64 = 0 182 while s < 8 { 183 var pos: i64 = 0 184 while pos < MN { 185 let tok: i64 = seqs[s * MN + pos] 186 var j: i64 = 0 187 while j < MD { 188 dE[tok * MD + j] = nx_f32_add(dE[tok * MD + j], ta_grad(tape, grads, xl8[s], pos * MD + j)) 189 j = j + 1 190 } 191 pos = pos + 1 192 } 193 s = s + 1 194 } 195 // adamw step 196 b1t = nx_f32_mul(b1t, beta1); b2t = nx_f32_mul(b2t, beta2) 197 let c1: i64 = nx_f32_sub(TA_F32_ONE, b1t); let c2: i64 = nx_f32_sub(TA_F32_ONE, b2t) 198 mo_adamw(E, mE, vE, dE, MV * MD, lr, beta1, beta2, om1, om2, eps, c1, c2) 199 mo_adamw(W1, mW1, vW1, gW1, MH * MND, lr, beta1, beta2, om1, om2, eps, c1, c2) 200 mo_adamw(b1, mb1, vb1, gb1, MH, lr, beta1, beta2, om1, om2, eps, c1, c2) 201 mo_adamw(W2, mW2, vW2, gW2, MC * MH, lr, beta1, beta2, om1, om2, eps, c1, c2) 202 mo_adamw(b2, mb2, vb2, gb2, MC, lr, beta1, beta2, om1, om2, eps, c1, c2) 203 ep = ep + 1 204 } 205 return 0 206} 207 208func main() -> i64 { 209 var ok: i64 = 1 210 let tape: *i64 = (sys_mmap(2048 * 7 * 8)) as *i64 211 let vals: *i64 = (sys_mmap(8192 * 8)) as *i64 212 let grads: *i64 = (sys_mmap(8192 * 8)) as *i64 213 let st: *i64 = (sys_mmap(2 * 8)) as *i64 214 215 // dataset: 8 sequences, label = (token[0] > token[3]) ? 1 : 0 216 let seqs: *i64 = (sys_mmap(32 * 8)) as *i64 217 seqs[0]=0; seqs[1]=1; seqs[2]=2; seqs[3]=3 218 seqs[4]=3; seqs[5]=2; seqs[6]=1; seqs[7]=0 219 seqs[8]=1; seqs[9]=1; seqs[10]=2; seqs[11]=2 220 seqs[12]=2; seqs[13]=2; seqs[14]=1; seqs[15]=1 221 seqs[16]=0; seqs[17]=0; seqs[18]=3; seqs[19]=3 222 seqs[20]=3; seqs[21]=3; seqs[22]=0; seqs[23]=0 223 seqs[24]=1; seqs[25]=2; seqs[26]=3; seqs[27]=0 224 seqs[28]=2; seqs[29]=1; seqs[30]=0; seqs[31]=3 225 let labels: *i64 = (sys_mmap(8 * 8)) as *i64 226 labels[0]=0; labels[1]=1; labels[2]=0; labels[3]=1; labels[4]=0; labels[5]=1; labels[6]=1; labels[7]=0 227 228 let E: *i64 = (sys_mmap(MV * MD * 8)) as *i64 229 let W1: *i64 = (sys_mmap(MH * MND * 8)) as *i64 230 let b1: *i64 = (sys_mmap(MH * 8)) as *i64 231 let W2: *i64 = (sys_mmap(MC * MH * 8)) as *i64 232 let b2: *i64 = (sys_mmap(MC * 8)) as *i64 233 let lf: *i64 = (sys_mmap(8)) as *i64 234 let ll: *i64 = (sys_mmap(8)) as *i64 235 mo_train(tape, vals, grads, st, E, W1, b1, W2, b2, seqs, labels, 2500, lf, ll) 236 237 var acc: i64 = 0 238 var s: i64 = 0 239 while s < 8 { 240 if mo_predict(tape, vals, st, E, W1, b1, W2, b2, seqs, s) == labels[s] { acc = acc + 1 } 241 s = s + 1 242 } 243 var trainPass: i64 = 1 244 if acc < 7 { trainPass = 0 } 245 if nx_f32_lt(*ll, *lf) != 1 { trainPass = 0 } 246 if trainPass != 1 { ok = 0 } 247 248 // bit-exact reproducibility 249 let E2: *i64 = (sys_mmap(MV * MD * 8)) as *i64 250 let W1b: *i64 = (sys_mmap(MH * MND * 8)) as *i64 251 let b1b: *i64 = (sys_mmap(MH * 8)) as *i64 252 let W2b: *i64 = (sys_mmap(MC * MH * 8)) as *i64 253 let b2b: *i64 = (sys_mmap(MC * 8)) as *i64 254 let lf2: *i64 = (sys_mmap(8)) as *i64 255 let ll2: *i64 = (sys_mmap(8)) as *i64 256 mo_train(tape, vals, grads, st, E2, W1b, b1b, W2b, b2b, seqs, labels, 2500, lf2, ll2) 257 var acc2: i64 = 0 258 s = 0 259 while s < 8 { 260 if mo_predict(tape, vals, st, E2, W1b, b1b, W2b, b2b, seqs, s) == labels[s] { acc2 = acc2 + 1 } 261 s = s + 1 262 } 263 var reproPass: i64 = 1 264 if acc2 != acc { reproPass = 0 } 265 if *ll2 != *ll { reproPass = 0 } 266 if reproPass != 1 { ok = 0 } 267 268 var fdi: i64 = 1 269 while fdi >= 0 { 270 var out: i64 = 1 271 if fdi == 0 { out = sys_openat_append(MO_LOG, 420) } 272 if out >= 0 { 273 mo_w(out, "FNETMODELGATE authored=organ model=embed+fnet-mix+relu-ffn+softmaxCE task=t0>t3 no-attention" as *u8) 274 mo_w(out, " | accuracy=" as *u8); mo_wn(out, acc); mo_w(out, "/8" as *u8) 275 mo_w(out, " loss_first_milli=" as *u8); mo_wn(out, ta_f32_to_milli(*lf)) 276 mo_w(out, " loss_last_milli=" as *u8); mo_wn(out, ta_f32_to_milli(*ll)) 277 mo_w(out, " | bitexact_repro=" as *u8); mo_wn(out, reproPass) 278 if ok == 1 { mo_w(out, " verdict=GREEN\n" as *u8) } else { mo_w(out, " verdict=RED\n" as *u8) } 279 if fdi == 0 { sys_close(out) } 280 } 281 fdi = fdi - 1 282 } 283 284 // MIGRATED onto nx_gate_verdict by nx_gate_dry_apply (D001, minimal form): every check 285 // row above is untouched, so the PASS/FAIL vector cannot change; only the hand-rolled 286 // verdict emission is replaced by the ONE shared base class. Proven by nx_gate_migrate verify. 287 let ctr__dry: *i64 = gv_ctr() 288 ctr__dry[0] = ok 289 ctr__dry[1] = 1 290 let rc__dry: i64 = gv_verdict("FNET-MODEL-GATE" as *u8, ctr__dry, "teeth unchanged; verdict emission migrated onto the shared base class" as *u8) 291 sys_exit(rc__dry) 292 return rc__dry 293}