code wiki / _hdl_build / nx_f32_qwen2_train_gate.nx
nx_f32_qwen2_train_gate.nx source
↩ module page · 324 lines · 28181 B
1import "nx_gate_gn.nx"
2import "nx_gate_base.nx"
3// nx_f32_qwen2_train_gate.nx -- FULL-FAITHFUL Qwen2 trainer: the mini-Qwen block + the two params it was missing
4// so its output is a COMPLETE Qwen2 model our no-float inference can RUN (closing loop-b): learned RMSNorm GAMMA
5// + QKV BIAS. Single-head (a valid Qwen2 config: n_heads=1, n_kv=1, head_dim=D). Per layer:
6// x -> RMSNorm*gamma -> Q/K/V(+bias) -> NEOX-RoPE(Q,K) -> causal attn -> O -> +res -> RMSNorm*gamma -> SwiGLU -> +res
7// final RMSNorm*gamma -> head -> per-position CE. GAMMA and BIAS enter via the "tile [D]->[T,D] leaf, scatter-reduce
8// its grad to the [D] buffer" pattern (same as the embed gradient) -- so they ride the SAME verified ta_ tape ops
9// (hadamard for gamma, vadd for bias); no autograd-lib edit. GRADCHECKED THROUGH THE WHOLE NET incl the new params
10// (dWq, dGamma, dBias vs central finite-diff). This model has EVERY tensor Qwen2 inference needs -> next: gguf-export
11// it -> run on arch-config inference -> the first ours-on-our-stack generation + perplexity. ORIGINAL, no gcc/py.
12import "nx_autograd_tensor.nx"
13import "nx_syscalls.nx"
14
15const VMAX: i64 = 40
16const DM: i64 = 16 // perfect square -> trainer 1/sqrt(16)=0.25 EXACTLY equals inference 65536/isqrt(16)
17const HF: i64 = 32
18const NLAYERS: i64 = 2
19const TMAX: i64 = 64
20const EPOCHS: i64 = 1800
21const NODES: i64 = 220
22const CELLS: i64 = 300000
23
24func grow(name: *u8, ok: i64) -> i64 { if ok==1 { gw(" PASS " as *u8) } else { gw(" FAIL " as *u8) } gw(name); gw("
25" as *u8); return ok }
26func gm(x: i64) -> i64 { return gn(ta_f32_to_milli(x)) }
27func f32c(a: i64, b: i64) -> i64 { return nx_f32_div(nx_i32_to_f32(a), nx_i32_to_f32(b)) }
28func f32_absx(x: i64) -> i64 { return x & 0x7FFFFFFF }
29func f32_lex(x: i64, y: i64) -> i64 { if nx_f32_gt(x,y)==1 { return 0 } return 1 }
30func f32_sqrtx(x: i64) -> i64 { if (x & 0x7FFFFFFF)==0 { return nx_i32_to_f32(0) } var y: i64=x; var i: i64=0; while i<16 { y=nx_f32_div(nx_f32_add(y, nx_f32_div(x,y)), nx_i32_to_f32(2)); i=i+1 } return y }
31
32// build NEOX RoPE tables (as in mini-qwen)
33func build_rope(COS: *i64, SIN: *i64, ROT: *i64) -> i64 {
34 let np: i64=DM/2; let lnb: i64=nx_f32_log(nx_i32_to_f32(1000000))
35 var t: i64=0
36 while t<TMAX { var i: i64=0
37 while i<np { let frac: i64=nx_f32_div(nx_f32_mul(nx_i32_to_f32(2*i), lnb), nx_i32_to_f32(DM)); let theta: i64=nx_f32_exp(nx_f32_neg(frac)); let ang: i64=nx_f32_mul(nx_i32_to_f32(t), theta)
38 let c: i64=nx_f32_cos(ang); let s: i64=nx_f32_sin(ang); COS[t*DM+i]=c; COS[t*DM+i+np]=c; SIN[t*DM+i]=s; SIN[t*DM+i+np]=s; i=i+1 }
39 t=t+1 }
40 var j: i64=0; while j<DM*DM { ROT[j]=nx_i32_to_f32(0); j=j+1 }
41 var k: i64=0; while k<np { ROT[(k+np)*DM + k]=nx_f32_neg(nx_i32_to_f32(1)); k=k+1 }
42 k=np; while k<DM { ROT[(k-np)*DM + k]=nx_i32_to_f32(1); k=k+1 }
43 return 0
44}
45func rope_node(tape: *i64, vals: *i64, st: *i64, xn: i64, COSn: i64, SINn: i64, ROTn: i64) -> i64 {
46 let a: i64=ta_hadamard(tape, vals, st, xn, COSn); let rh: i64=ta_matmul(tape, vals, st, xn, ROTn); let b: i64=ta_hadamard(tape, vals, st, rh, SINn); return ta_vadd(tape, vals, st, a, b)
47}
48// tile a [width] param into a [T,width] source buffer (broadcast per row) -> ta_leaf. Returns node id via out.
49func tile_leaf(tape: *i64, vals: *i64, st: *i64, param: *i64, T: i64, width: i64, scratch: *i64) -> i64 {
50 var t: i64=0; while t<T { var d: i64=0; while d<width { scratch[t*width+d]=param[d]; d=d+1 } t=t+1 }
51 return ta_leaf(tape, vals, st, T, width, scratch, 0)
52}
53// scatter-reduce a tiled leaf's grad [T,width] back to a [width] grad buffer (sum over rows).
54func reduce_tiled(tape: *i64, grads: *i64, node: i64, T: i64, width: i64, dst: *i64) -> i64 {
55 var d: i64=0; while d<width { dst[d]=nx_i32_to_f32(0); d=d+1 }
56 var t: i64=0; while t<T { d=0; while d<width { dst[d]=nx_f32_add(dst[d], ta_grad(tape, grads, node, t*width+d)); d=d+1 } t=t+1 }
57 return 0
58}
59
60// gamma-scaled RMSNorm: hadamard(rmsnorm_rows(x), tile(gamma)). returns the node; records the gamma-tile node in gv.
61func rmsnorm_gamma(tape: *i64, vals: *i64, st: *i64, x: i64, gamma: *i64, T: i64, scratch: *i64, gv: *i64, gi: i64) -> i64 {
62 let xn: i64=ta_rmsnorm_rows(tape, vals, st, x)
63 let gnode: i64=tile_leaf(tape, vals, st, gamma, T, DM, scratch); gv[gi]=gnode
64 return ta_hadamard(tape, vals, st, xn, gnode)
65}
66
67// M param layout: [0]=E [1]=Wlm [2]=final_gamma then per layer L (13 each) starting at 3+L*13:
68// +0 Wq +1 Wk +2 Wv +3 Wo +4 Wg +5 Wu +6 Wd +7 attn_gamma +8 ffn_gamma +9 bq +10 bk +11 bv
69func LB(L: i64) -> i64 { return 3+L*13 }
70// leaf-node-id tracking: lv[] mirrors M for the WEIGHT leaves; gv[]/bv[] track the tiled gamma/bias leaves per fwd.
71func build_fwd(tape: *i64, vals: *i64, st: *i64, M: *i64, toks: *i64, T: i64, vv: i64, scale: i64, Xsrc: *i64, COS: *i64, SIN: *i64, ROT: *i64, tsc: *i64, lv: *i64, gv: *i64, bvn: *i64) -> i64 {
72 let E: *i64=M[0] as *i64
73 var tp: i64=0; while tp<T { let tk: i64=toks[tp]; var d: i64=0; while d<DM { Xsrc[tp*DM+d]=E[tk*DM+d]; d=d+1 } tp=tp+1 }
74 var cur: i64=ta_leaf(tape, vals, st, T, DM, Xsrc, 0); lv[0]=cur
75 let COSn: i64=ta_leaf(tape, vals, st, T, DM, COS, 0); let SINn: i64=ta_leaf(tape, vals, st, T, DM, SIN, 0); let ROTn: i64=ta_leaf(tape, vals, st, DM, DM, ROT, 0)
76 var L: i64=0
77 while L<NLAYERS {
78 let b: i64=LB(L)
79 let Wq: i64=ta_leaf(tape, vals, st, DM, DM, M[b+0] as *i64, 0); lv[b+0]=Wq
80 let Wk: i64=ta_leaf(tape, vals, st, DM, DM, M[b+1] as *i64, 0); lv[b+1]=Wk
81 let Wv: i64=ta_leaf(tape, vals, st, DM, DM, M[b+2] as *i64, 0); lv[b+2]=Wv
82 let Wo: i64=ta_leaf(tape, vals, st, DM, DM, M[b+3] as *i64, 0); lv[b+3]=Wo
83 let Wg: i64=ta_leaf(tape, vals, st, DM, HF, M[b+4] as *i64, 0); lv[b+4]=Wg
84 let Wu: i64=ta_leaf(tape, vals, st, DM, HF, M[b+5] as *i64, 0); lv[b+5]=Wu
85 let Wd: i64=ta_leaf(tape, vals, st, HF, DM, M[b+6] as *i64, 0); lv[b+6]=Wd
86 let xn1: i64=rmsnorm_gamma(tape, vals, st, cur, M[b+7] as *i64, T, tsc, gv, 2*L)
87 var Q: i64=ta_matmul(tape, vals, st, xn1, Wq); Q=ta_vadd(tape, vals, st, Q, tile_leaf(tape, vals, st, M[b+9] as *i64, T, DM, tsc)); bvn[3*L+0]=st[0]-1
88 var Kk: i64=ta_matmul(tape, vals, st, xn1, Wk); Kk=ta_vadd(tape, vals, st, Kk, tile_leaf(tape, vals, st, M[b+10] as *i64, T, DM, tsc)); bvn[3*L+1]=st[0]-1
89 var Vv: i64=ta_matmul(tape, vals, st, xn1, Wv); Vv=ta_vadd(tape, vals, st, Vv, tile_leaf(tape, vals, st, M[b+11] as *i64, T, DM, tsc)); bvn[3*L+2]=st[0]-1
90 let Qr: i64=rope_node(tape, vals, st, Q, COSn, SINn, ROTn)
91 let Kr: i64=rope_node(tape, vals, st, Kk, COSn, SINn, ROTn)
92 let sc0: i64=ta_matmul_nt(tape, vals, st, Qr, Kr); let sc: i64=ta_cmul(tape, vals, st, sc0, scale)
93 let A: i64=ta_softmax_rows(tape, vals, st, sc, 1); let ctx: i64=ta_matmul(tape, vals, st, A, Vv); let O: i64=ta_matmul(tape, vals, st, ctx, Wo)
94 let h1: i64=ta_vadd(tape, vals, st, cur, O)
95 let xn2: i64=rmsnorm_gamma(tape, vals, st, h1, M[b+8] as *i64, T, tsc, gv, 2*L+1)
96 let g: i64=ta_matmul(tape, vals, st, xn2, Wg); let gs: i64=ta_silu(tape, vals, st, g); let up: i64=ta_matmul(tape, vals, st, xn2, Wu)
97 let ff: i64=ta_hadamard(tape, vals, st, gs, up); let dn: i64=ta_matmul(tape, vals, st, ff, Wd)
98 cur=ta_vadd(tape, vals, st, h1, dn)
99 L=L+1
100 }
101 let xf: i64=rmsnorm_gamma(tape, vals, st, cur, M[2] as *i64, T, tsc, gv, 2*NLAYERS)
102 let Wlm: i64=ta_leaf(tape, vals, st, DM, vv, M[1] as *i64, 0); lv[1]=Wlm
103 return ta_matmul(tape, vals, st, xf, Wlm)
104}
105func loss_of(tape: *i64, vals: *i64, M: *i64, toks: *i64, tgt: *i64, T: i64, vv: i64, scale: i64, Xsrc: *i64, COS: *i64, SIN: *i64, ROT: *i64, tsc: *i64, lv: *i64, gv: *i64, bvn: *i64) -> i64 {
106 let st: *i64=sys_mmap(16) as *i64; st[0]=0; st[1]=0
107 let lg: i64=build_fwd(tape, vals, st, M, toks, T, vv, scale, Xsrc, COS, SIN, ROT, tsc, lv, gv, bvn)
108 return ta_val(tape, vals, ta_softce_rows(tape, vals, st, lg, tgt), 0)
109}
110func det_fill(dst: *i64, n: i64, seed: i64) -> i64 { var i: i64=0; while i<n { let s: i64=(((i*2654435761)+seed) % 13) - 6; dst[i]=f32c(s, 100); i=i+1 } return 0 }
111func argmax_row(vals: *i64, tape: *i64, lg: i64, row: i64, vv: i64) -> i64 { let off: i64=tape[7*lg+5]+row*vv; var bi: i64=0; var o: i64=1; while o<vv { if nx_f32_gt(vals[off+o], vals[off+bi])==1 { bi=o } o=o+1 } return bi }
112func adam_step(P: *i64, Gd: *i64, Mo: *i64, Vo: *i64, cnt: i64, lr: i64, b1: i64, b2: i64, bc1: i64, bc2: i64, aeps: i64) -> i64 {
113 let one: i64=nx_i32_to_f32(1); var w: i64=0
114 while w<cnt { let g: i64=Gd[w]; Mo[w]=nx_f32_add(nx_f32_mul(b1,Mo[w]),nx_f32_mul(nx_f32_sub(one,b1),g)); Vo[w]=nx_f32_add(nx_f32_mul(b2,Vo[w]),nx_f32_mul(nx_f32_sub(one,b2),nx_f32_mul(g,g))); let mh: i64=nx_f32_div(Mo[w],bc1); let vh: i64=nx_f32_div(Vo[w],bc2); P[w]=nx_f32_sub(P[w], nx_f32_div(nx_f32_mul(lr,mh), nx_f32_add(f32_sqrtx(vh),aeps))); w=w+1 }
115 return 0
116}
117
118// ---- gguf EXPORT: write the trained weights as a full Qwen2 gguf our no-float inference LOADS ----
119// LE writer primitives (pos advances via *i64).
120func w_u8(b: *u8, p: *i64, v: i64) -> i64 { b[p[0]]=(v & 0xff) as u8; p[0]=p[0]+1; return 0 }
121func w_u32(b: *u8, p: *i64, v: i64) -> i64 { w_u8(b,p,v); w_u8(b,p,v>>8); w_u8(b,p,v>>16); w_u8(b,p,v>>24); return 0 }
122func w_u64(b: *u8, p: *i64, v: i64) -> i64 { w_u32(b,p,v & 0xffffffff); w_u32(b,p,(v>>32) & 0xffffffff); return 0 }
123func w_str(b: *u8, p: *i64, s: *u8, n: i64) -> i64 { w_u64(b,p,n); var i: i64=0; while i<n { b[p[0]]=s[i]; p[0]=p[0]+1; i=i+1 } return 0 }
124func w_align(b: *u8, p: *i64) -> i64 { while (p[0] % 32)!=0 { b[p[0]]=0 as u8; p[0]=p[0]+1 } return 0 }
125func w_meta_u32(b: *u8, p: *i64, key: *u8, kl: i64, v: i64) -> i64 { w_str(b,p,key,kl); w_u32(b,p,4); w_u32(b,p,v); return 0 }
126func slen2(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} return n }
127// transpose trainer W[in,out] (W[i*out+j]) -> dst[out,in] (dst[j*in+i]) = the inference [out,in] layout.
128func transpose(src: *i64, dst: *i64, in_dim: i64, out_dim: i64) -> i64 { var i: i64=0; while i<in_dim { var j: i64=0; while j<out_dim { dst[j*in_dim+i]=src[i*out_dim+j]; j=j+1 } i=i+1 } return 0 }
129// one tensor spec into ts[idx*8]: name, ndims, dim0(in), dim1(out), F32, data ptr, ncells
130func setspec(ts: *i64, idx: i64, name: *u8, ndims: i64, d0: i64, d1: i64, data: *i64, ncells: i64) -> i64 {
131 ts[idx*8+0]=name as i64; ts[idx*8+1]=slen2(name); ts[idx*8+2]=ndims; ts[idx*8+3]=d0; ts[idx*8+4]=d1; ts[idx*8+5]=0; ts[idx*8+6]=data as i64; ts[idx*8+7]=ncells; return 0
132}
133func blkname(L: i64, suffix: *u8) -> *u8 { let o: *u8=sys_mmap(48); o[0]=98 as u8; o[1]=108 as u8; o[2]=107 as u8; o[3]=46 as u8; var p: i64=4; if L==0 { o[p]=48 as u8; p=p+1 } else { let t: *u8=sys_mmap(16); var k: i64=0; var m: i64=L; while m>0 { t[k]=(48+(m%10)) as u8; m=m/10; k=k+1 } var i: i64=k-1; while i>=0 { o[p]=t[i]; p=p+1; i=i-1 } } var j: i64=0; while suffix[j]!=(0 as u8) { o[p]=suffix[j]; p=p+1; j=j+1 } o[p]=0 as u8; return o }
134// export M (trained) as a full Qwen2 gguf to `path`. vv=vocab.
135func export_gguf(path: *u8, M: *i64, vv: i64) -> i64 {
136 // build tensor specs (transposing every weight matrix in-major -> out-major)
137 let NT: i64=3+NLAYERS*12
138 let ts: *i64=sys_mmap(NT*8*8) as *i64
139 var ti: i64=0
140 // token_embd.weight: data = E[V,DM] directly (row=token). dims (in=DM, out=V).
141 setspec(ts, ti, "token_embd.weight" as *u8, 2, DM, vv, M[0] as *i64, vv*DM); ti=ti+1
142 // output.weight: transpose Wlm[DM,V] -> [V,DM]. dims (in=DM, out=V).
143 let oT: *i64=sys_mmap(DM*vv*8) as *i64; transpose(M[1] as *i64, oT, DM, vv)
144 setspec(ts, ti, "output.weight" as *u8, 2, DM, vv, oT, vv*DM); ti=ti+1
145 setspec(ts, ti, "output_norm.weight" as *u8, 1, DM, 1, M[2] as *i64, DM); ti=ti+1
146 var L: i64=0
147 while L<NLAYERS {
148 let b: i64=LB(L)
149 setspec(ts, ti, blkname(L, ".attn_norm.weight" as *u8), 1, DM, 1, M[b+7] as *i64, DM); ti=ti+1
150 let qT: *i64=sys_mmap(DM*DM*8) as *i64; transpose(M[b+0] as *i64, qT, DM, DM); setspec(ts, ti, blkname(L, ".attn_q.weight" as *u8), 2, DM, DM, qT, DM*DM); ti=ti+1
151 setspec(ts, ti, blkname(L, ".attn_q.bias" as *u8), 1, DM, 1, M[b+9] as *i64, DM); ti=ti+1
152 let kT: *i64=sys_mmap(DM*DM*8) as *i64; transpose(M[b+1] as *i64, kT, DM, DM); setspec(ts, ti, blkname(L, ".attn_k.weight" as *u8), 2, DM, DM, kT, DM*DM); ti=ti+1
153 setspec(ts, ti, blkname(L, ".attn_k.bias" as *u8), 1, DM, 1, M[b+10] as *i64, DM); ti=ti+1
154 let vT: *i64=sys_mmap(DM*DM*8) as *i64; transpose(M[b+2] as *i64, vT, DM, DM); setspec(ts, ti, blkname(L, ".attn_v.weight" as *u8), 2, DM, DM, vT, DM*DM); ti=ti+1
155 setspec(ts, ti, blkname(L, ".attn_v.bias" as *u8), 1, DM, 1, M[b+11] as *i64, DM); ti=ti+1
156 let woT: *i64=sys_mmap(DM*DM*8) as *i64; transpose(M[b+3] as *i64, woT, DM, DM); setspec(ts, ti, blkname(L, ".attn_output.weight" as *u8), 2, DM, DM, woT, DM*DM); ti=ti+1
157 setspec(ts, ti, blkname(L, ".ffn_norm.weight" as *u8), 1, DM, 1, M[b+8] as *i64, DM); ti=ti+1
158 let gT: *i64=sys_mmap(DM*HF*8) as *i64; transpose(M[b+4] as *i64, gT, DM, HF); setspec(ts, ti, blkname(L, ".ffn_gate.weight" as *u8), 2, DM, HF, gT, DM*HF); ti=ti+1
159 let uT: *i64=sys_mmap(DM*HF*8) as *i64; transpose(M[b+5] as *i64, uT, DM, HF); setspec(ts, ti, blkname(L, ".ffn_up.weight" as *u8), 2, DM, HF, uT, DM*HF); ti=ti+1
160 let dT: *i64=sys_mmap(HF*DM*8) as *i64; transpose(M[b+6] as *i64, dT, HF, DM); setspec(ts, ti, blkname(L, ".ffn_down.weight" as *u8), 2, HF, DM, dT, HF*DM); ti=ti+1
161 L=L+1
162 }
163 // relative data offsets (align 32 between tensors)
164 let reloff: *i64=sys_mmap((NT+1)*8) as *i64; reloff[0]=0
165 var i2: i64=0; while i2<NT { var nx: i64=reloff[i2]+ts[i2*8+7]*4; while (nx % 32)!=0 { nx=nx+1 } reloff[i2+1]=nx; i2=i2+1 }
166 let cap: i64=1048576
167 let buf: *u8=sys_mmap(cap); let p: *i64=sys_mmap(8) as *i64; p[0]=0
168 // header: magic, v3, tensor_count, metadata_count (1 string + 5 u32 = 6)
169 w_u32(buf,p,0x46554747); w_u32(buf,p,3); w_u64(buf,p,NT); w_u64(buf,p,6)
170 let ka: *u8="general.architecture" as *u8; w_str(buf,p,ka,20); w_u32(buf,p,8); w_str(buf,p,"qwen2" as *u8,5)
171 w_meta_u32(buf,p,"qwen2.embedding_length" as *u8,22,DM)
172 w_meta_u32(buf,p,"qwen2.block_count" as *u8,17,NLAYERS)
173 w_meta_u32(buf,p,"qwen2.attention.head_count" as *u8,26,1)
174 w_meta_u32(buf,p,"qwen2.attention.head_count_kv" as *u8,29,1)
175 w_meta_u32(buf,p,"qwen2.feed_forward_length" as *u8,25,HF)
176 // tensor infos
177 i2=0; while i2<NT { w_str(buf,p, ts[i2*8] as *u8, ts[i2*8+1]); let nd: i64=ts[i2*8+2]; w_u32(buf,p,nd); w_u64(buf,p,ts[i2*8+3]); if nd>=2 { w_u64(buf,p,ts[i2*8+4]) } w_u32(buf,p,0); w_u64(buf,p,reloff[i2]); i2=i2+1 }
178 w_align(buf,p)
179 // tensor data (F32 = low 32 of each f32 slot)
180 i2=0; while i2<NT { let data: *i64=ts[i2*8+6] as *i64; let n: i64=ts[i2*8+7]; var c: i64=0; while c<n { w_u32(buf,p, data[c] & 0xffffffff); c=c+1 } w_align(buf,p); i2=i2+1 }
181 let fd: i64=sys_openat_wr(path, 0x1a4); if fd<0 { return 0-1 }
182 sys_write(fd, buf, p[0]); sys_close(fd)
183 return p[0]
184}
185
186func main() -> i64 {
187 gw("=== nx_f32_qwen2_train_gate: FULL-FAITHFUL Qwen2 (RMSNorm-GAMMA + QKV-BIAS + RoPE + causal-attn + SwiGLU) from zero ===\n" as *u8)
188 var pass: i64=0; var total: i64=0
189 let one: i64=nx_i32_to_f32(1); let scale: i64=nx_f32_div(one, f32_sqrtx(nx_i32_to_f32(DM)))
190 let corpus: *u8="the quick brown fox jumps over the lazy dog. " as *u8
191 var clen: i64=0; while corpus[clen]!=(0 as u8){clen=clen+1}
192 let c2id: *i64=sys_mmap(256*8) as *i64; var ci: i64=0; while ci<256 { c2id[ci]=0-1; ci=ci+1 }
193 let id2c: *u8=sys_mmap(VMAX); var vv: i64=0
194 var t: i64=0; while t<clen { let ch: i64=corpus[t]&0xff; if c2id[ch]<0 { c2id[ch]=vv; id2c[vv]=ch as u8; vv=vv+1 } t=t+1 }
195 let alltok: *i64=sys_mmap(TMAX*8) as *i64; t=0; while t<clen { alltok[t]=c2id[corpus[t]&0xff]; t=t+1 }
196 let T: i64=clen-1
197 let toks: *i64=sys_mmap(TMAX*8) as *i64; t=0; while t<T { toks[t]=alltok[t]; t=t+1 }
198 let tgt: *i64=sys_mmap(TMAX*8) as *i64; t=0; while t<T { tgt[t]=alltok[t+1]; t=t+1 }
199 gw(" corpus V=" as *u8); gn(vv); gw(" T=" as *u8); gn(T); gw(" | Qwen2 block: d=" as *u8); gn(DM); gw(" layers=" as *u8); gn(NLAYERS); gw(" 1-head, RMSNorm-gamma + QKV-bias + SwiGLU\n" as *u8)
200
201 // params
202 let NP: i64=3+NLAYERS*13
203 let M: *i64=sys_mmap((NP+2)*8) as *i64; let szs: *i64=sys_mmap((NP+2)*8) as *i64
204 M[0]=sys_mmap(VMAX*DM*8) as i64; szs[0]=vv*DM; det_fill(M[0] as *i64, vv*DM, 1)
205 M[1]=sys_mmap(DM*VMAX*8) as i64; szs[1]=DM*vv; det_fill(M[1] as *i64, DM*vv, 3)
206 M[2]=sys_mmap(DM*8) as i64; szs[2]=DM; var gg: i64=0; while gg<DM { let fp: *i64=M[2] as *i64; fp[gg]=one; gg=gg+1 } // final gamma init 1
207 var L: i64=0
208 while L<NLAYERS {
209 let b: i64=LB(L)
210 M[b+0]=sys_mmap(DM*DM*8) as i64; szs[b+0]=DM*DM; det_fill(M[b+0] as *i64, DM*DM, 5+L*17)
211 M[b+1]=sys_mmap(DM*DM*8) as i64; szs[b+1]=DM*DM; det_fill(M[b+1] as *i64, DM*DM, 7+L*17)
212 M[b+2]=sys_mmap(DM*DM*8) as i64; szs[b+2]=DM*DM; det_fill(M[b+2] as *i64, DM*DM, 9+L*17)
213 M[b+3]=sys_mmap(DM*DM*8) as i64; szs[b+3]=DM*DM; det_fill(M[b+3] as *i64, DM*DM, 11+L*17)
214 M[b+4]=sys_mmap(DM*HF*8) as i64; szs[b+4]=DM*HF; det_fill(M[b+4] as *i64, DM*HF, 13+L*17)
215 M[b+5]=sys_mmap(DM*HF*8) as i64; szs[b+5]=DM*HF; det_fill(M[b+5] as *i64, DM*HF, 19+L*17)
216 M[b+6]=sys_mmap(HF*DM*8) as i64; szs[b+6]=HF*DM; det_fill(M[b+6] as *i64, HF*DM, 23+L*17)
217 M[b+7]=sys_mmap(DM*8) as i64; szs[b+7]=DM; var z1: i64=0; while z1<DM { let gp: *i64=M[b+7] as *i64; gp[z1]=one; z1=z1+1 } // attn gamma=1
218 M[b+8]=sys_mmap(DM*8) as i64; szs[b+8]=DM; z1=0; while z1<DM { let gp: *i64=M[b+8] as *i64; gp[z1]=one; z1=z1+1 } // ffn gamma=1
219 M[b+9]=sys_mmap(DM*8) as i64; szs[b+9]=DM; z1=0; while z1<DM { let bp: *i64=M[b+9] as *i64; bp[z1]=nx_i32_to_f32(0); z1=z1+1 } // bq=0
220 M[b+10]=sys_mmap(DM*8) as i64; szs[b+10]=DM; z1=0; while z1<DM { let bp: *i64=M[b+10] as *i64; bp[z1]=nx_i32_to_f32(0); z1=z1+1 }
221 M[b+11]=sys_mmap(DM*8) as i64; szs[b+11]=DM; z1=0; while z1<DM { let bp: *i64=M[b+11] as *i64; bp[z1]=nx_i32_to_f32(0); z1=z1+1 }
222 M[b+12]=0; szs[b+12]=0
223 L=L+1
224 }
225 let COS: *i64=sys_mmap(TMAX*DM*8) as *i64; let SIN: *i64=sys_mmap(TMAX*DM*8) as *i64; let ROT: *i64=sys_mmap(DM*DM*8) as *i64; build_rope(COS, SIN, ROT)
226 let tape: *i64=sys_mmap(NODES*7*8) as *i64; let vals: *i64=sys_mmap(CELLS*8) as *i64; let grads: *i64=sys_mmap(CELLS*8) as *i64
227 let Xsrc: *i64=sys_mmap(TMAX*DM*8) as *i64; let tsc: *i64=sys_mmap(TMAX*DM*8) as *i64; let st: *i64=sys_mmap(16) as *i64
228 let lv: *i64=sys_mmap((NP+2)*8) as *i64; let gv: *i64=sys_mmap((2*NLAYERS+2)*8) as *i64; let bvn: *i64=sys_mmap((3*NLAYERS+2)*8) as *i64
229
230 // T0
231 st[0]=0; st[1]=0
232 let lg0: i64=build_fwd(tape, vals, st, M, toks, T, vv, scale, Xsrc, COS, SIN, ROT, tsc, lv, gv, bvn)
233 let ls0: i64=ta_softce_rows(tape, vals, st, lg0, tgt)
234 total=total+1; pass=pass+1
235 gw(" [PASS] T0 FORWARD: nodes=" as *u8); gn(st[0]); gw(" cells=" as *u8); gn(st[1]); gw(" initial mean CE=" as *u8); gm(ta_val(tape, vals, ls0, 0)); gw("m\n" as *u8)
236
237 let hh: i64=f32c(1,100); let twoh: i64=nx_f32_mul(nx_i32_to_f32(2),hh); let tol: i64=f32c(5,100)
238 // T1 dWq[L0]
239 st[0]=0; st[1]=0; let lgb: i64=build_fwd(tape, vals, st, M, toks, T, vv, scale, Xsrc, COS, SIN, ROT, tsc, lv, gv, bvn); let lsb: i64=ta_softce_rows(tape, vals, st, lgb, tgt); ta_backward(tape, vals, grads, st[0], lsb)
240 let anaWq: i64=ta_grad(tape, grads, lv[LB(0)], 0); let Wq0: *i64=M[LB(0)] as *i64; let sv: i64=Wq0[0]
241 Wq0[0]=nx_f32_add(sv,hh); let lp: i64=loss_of(tape,vals,M,toks,tgt,T,vv,scale,Xsrc,COS,SIN,ROT,tsc,lv,gv,bvn); Wq0[0]=nx_f32_sub(sv,hh); let lm2: i64=loss_of(tape,vals,M,toks,tgt,T,vv,scale,Xsrc,COS,SIN,ROT,tsc,lv,gv,bvn); Wq0[0]=sv
242 let fdWq: i64=nx_f32_div(nx_f32_sub(lp,lm2), twoh)
243 total=total+1; if f32_lex(f32_absx(nx_f32_sub(anaWq,fdWq)),tol)==1 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) }
244 gw("T1 dL/dWq[L0] GRADCHECK (full net): ana=" as *u8); gm(anaWq); gw("m fd=" as *u8); gm(fdWq); gw("m\n" as *u8)
245
246 // T2 dGamma (attn gamma of layer 0) -- the new param; reduce tiled leaf grad -> compare cell 0
247 st[0]=0; st[1]=0; let lgg: i64=build_fwd(tape, vals, st, M, toks, T, vv, scale, Xsrc, COS, SIN, ROT, tsc, lv, gv, bvn); let lsg: i64=ta_softce_rows(tape, vals, st, lgg, tgt); ta_backward(tape, vals, grads, st[0], lsg)
248 let dgamma: *i64=sys_mmap(DM*8) as *i64; reduce_tiled(tape, grads, gv[0], T, DM, dgamma)
249 let Gam0: *i64=M[LB(0)+7] as *i64; let svg: i64=Gam0[0]
250 Gam0[0]=nx_f32_add(svg,hh); let lpg: i64=loss_of(tape,vals,M,toks,tgt,T,vv,scale,Xsrc,COS,SIN,ROT,tsc,lv,gv,bvn); Gam0[0]=nx_f32_sub(svg,hh); let lmg: i64=loss_of(tape,vals,M,toks,tgt,T,vv,scale,Xsrc,COS,SIN,ROT,tsc,lv,gv,bvn); Gam0[0]=svg
251 let fdg: i64=nx_f32_div(nx_f32_sub(lpg,lmg), twoh)
252 total=total+1; if f32_lex(f32_absx(nx_f32_sub(dgamma[0],fdg)),tol)==1 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) }
253 gw("T2 dL/dGamma[attn L0] GRADCHECK (new param): ana=" as *u8); gm(dgamma[0]); gw("m fd=" as *u8); gm(fdg); gw("m\n" as *u8)
254
255 // T3 dBias (bq of layer 0) -- the other new param
256 st[0]=0; st[1]=0; let lgc: i64=build_fwd(tape, vals, st, M, toks, T, vv, scale, Xsrc, COS, SIN, ROT, tsc, lv, gv, bvn); let lsc: i64=ta_softce_rows(tape, vals, st, lgc, tgt); ta_backward(tape, vals, grads, st[0], lsc)
257 let dbias: *i64=sys_mmap(DM*8) as *i64; reduce_tiled(tape, grads, bvn[0], T, DM, dbias)
258 let Bq0: *i64=M[LB(0)+9] as *i64; let svb: i64=Bq0[0]
259 Bq0[0]=nx_f32_add(svb,hh); let lpb: i64=loss_of(tape,vals,M,toks,tgt,T,vv,scale,Xsrc,COS,SIN,ROT,tsc,lv,gv,bvn); Bq0[0]=nx_f32_sub(svb,hh); let lmb: i64=loss_of(tape,vals,M,toks,tgt,T,vv,scale,Xsrc,COS,SIN,ROT,tsc,lv,gv,bvn); Bq0[0]=svb
260 let fdb: i64=nx_f32_div(nx_f32_sub(lpb,lmb), twoh)
261 total=total+1; if f32_lex(f32_absx(nx_f32_sub(dbias[0],fdb)),tol)==1 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) }
262 gw("T3 dL/dBias[bq L0] GRADCHECK (new param): ana=" as *u8); gm(dbias[0]); gw("m fd=" as *u8); gm(fdb); gw("m\n" as *u8)
263
264 // T4 TRAIN (Adam over all params incl gamma+bias)
265 let Mo: *i64=sys_mmap((NP+2)*8) as *i64; let Vo: *i64=sys_mmap((NP+2)*8) as *i64; let Gd: *i64=sys_mmap((NP+2)*8) as *i64
266 var pi: i64=0; while pi<NP { let sz: i64=szs[pi]; if sz>0 { Mo[pi]=sys_mmap(sz*8) as i64; Vo[pi]=sys_mmap(sz*8) as i64; Gd[pi]=sys_mmap(sz*8) as i64; let mo: *i64=Mo[pi] as *i64; let vo: *i64=Vo[pi] as *i64; var z: i64=0; while z<sz { mo[z]=nx_i32_to_f32(0); vo[z]=nx_i32_to_f32(0); z=z+1 } } pi=pi+1 }
267 let b1: i64=f32c(9,10); let b2: i64=f32c(999,1000); let lr: i64=f32c(2,1000); let aeps: i64=f32c(1,100000000)
268 var b1t: i64=one; var b2t: i64=one; var loss0: i64=nx_i32_to_f32(0); var lossF: i64=nx_i32_to_f32(0)
269 let _t0us: i64=sys_now_us()
270 var ep: i64=1
271 while ep<=EPOCHS {
272 st[0]=0; st[1]=0
273 let lg: i64=build_fwd(tape, vals, st, M, toks, T, vv, scale, Xsrc, COS, SIN, ROT, tsc, lv, gv, bvn)
274 let ls: i64=ta_softce_rows(tape, vals, st, lg, tgt); let el: i64=ta_val(tape, vals, ls, 0); ta_backward(tape, vals, grads, st[0], ls)
275 b1t=nx_f32_mul(b1t,b1); b2t=nx_f32_mul(b2t,b2); let bc1: i64=nx_f32_sub(one,b1t); let bc2: i64=nx_f32_sub(one,b2t)
276 // embed grad scatter
277 let dE: *i64=Gd[0] as *i64; var de: i64=0; while de<vv*DM { dE[de]=nx_i32_to_f32(0); de=de+1 }
278 var tp2: i64=0; while tp2<T { let tk: i64=toks[tp2]; var d: i64=0; while d<DM { dE[tk*DM+d]=nx_f32_add(dE[tk*DM+d], ta_grad(tape, grads, lv[0], tp2*DM+d)); d=d+1 } tp2=tp2+1 }
279 // weight leaves: Wlm(1) + per-layer 7 weights
280 var c: i64=0; while c<DM*vv { let gp: *i64=Gd[1] as *i64; gp[c]=ta_grad(tape, grads, lv[1], c); c=c+1 }
281 var L2: i64=0; while L2<NLAYERS { let b: i64=LB(L2); var wi: i64=0; while wi<7 { let sz: i64=szs[b+wi]; let gp: *i64=Gd[b+wi] as *i64; c=0; while c<sz { gp[c]=ta_grad(tape, grads, lv[b+wi], c); c=c+1 } wi=wi+1 }
282 // gamma (attn=gv[2L], ffn=gv[2L+1]) + bias (bq/bk/bv = bvn[3L..])
283 reduce_tiled(tape, grads, gv[2*L2], T, DM, Gd[b+7] as *i64)
284 reduce_tiled(tape, grads, gv[2*L2+1], T, DM, Gd[b+8] as *i64)
285 reduce_tiled(tape, grads, bvn[3*L2+0], T, DM, Gd[b+9] as *i64)
286 reduce_tiled(tape, grads, bvn[3*L2+1], T, DM, Gd[b+10] as *i64)
287 reduce_tiled(tape, grads, bvn[3*L2+2], T, DM, Gd[b+11] as *i64)
288 L2=L2+1 }
289 reduce_tiled(tape, grads, gv[2*NLAYERS], T, DM, Gd[2] as *i64) // final gamma
290 // Adam all (skip size-0 slots)
291 var ai: i64=0; while ai<NP { if szs[ai]>0 { adam_step(M[ai] as *i64, Gd[ai] as *i64, Mo[ai] as *i64, Vo[ai] as *i64, szs[ai], lr, b1, b2, bc1, bc2, aeps) } ai=ai+1 }
292 if ep==1 { loss0=el } lossF=el
293 ep=ep+1
294 }
295 let _t1us: i64=sys_now_us(); let _dtms: i64=(_t1us-_t0us)/1000
296 gw(" [TIMING] TRAIN loop (fwd+bwd+adam, all params): " as *u8); gn(_dtms); gw(" ms for " as *u8); gn(EPOCHS); gw(" epochs = " as *u8); gn((_dtms*1000)/EPOCHS); gw(" us/epoch <- h2h vs PyTorch 4750 us/epoch (same model, CPU)\n" as *u8)
297 total=total+1; if f32_lex(lossF,loss0)==1 { if ta_f32_to_milli(lossF)<=300 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } } else { gw(" [FAIL] " as *u8) }
298 gw("T4 TRAIN: mean CE " as *u8); gm(loss0); gw("m -> " as *u8); gm(lossF); gw("m over " as *u8); gn(EPOCHS); gw(" epochs (all params incl gamma+bias)\n" as *u8)
299
300 // T5 GENERATE
301 let gseq: *i64=sys_mmap(TMAX*8) as *i64; gseq[0]=alltok[0]; var glen: i64=1
302 while glen<clen { st[0]=0; st[1]=0; let lgz: i64=build_fwd(tape, vals, st, M, gseq, glen, vv, scale, Xsrc, COS, SIN, ROT, tsc, lv, gv, bvn); let nxt: i64=argmax_row(vals, tape, lgz, glen-1, vv); gseq[glen]=nxt; glen=glen+1 }
303 let gstr: *u8=sys_mmap(TMAX+8); var gi: i64=0; while gi<clen { gstr[gi]=id2c[gseq[gi]]; gi=gi+1 }
304 gw(" generated: '" as *u8); sys_write(1, gstr, clen); gw("'\n" as *u8)
305 var matchlen: i64=0; gi=1; while gi<clen { if gstr[gi]==corpus[gi] { matchlen=matchlen+1; gi=gi+1 } else { gi=clen } }
306 total=total+1; if matchlen>=30 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) }
307 gw("T5 GENERATE: reproduced " as *u8); gn(matchlen); gw("/" as *u8); gn(clen-1); gw(" chars (full-faithful Qwen2 learned the sequence)\n" as *u8)
308
309 // T6 EXPORT: write the trained weights as a full Qwen2 gguf (all blk.N.* tensors transposed to inference
310 // layout + Qwen2 metadata) that our no-float arch-config inference LOADS. Also dump the char->id map so the
311 // loop-b loader can feed char ids + decode. This is the artifact that closes loop-b.
312 let nb: i64=export_gguf("/tmp/nx_ours_qwen.gguf" as *u8, M, vv)
313 // write the id->char table (vv bytes) so the loader decodes without a BPE tokenizer
314 let mfd: i64=sys_openat_wr("/tmp/nx_ours_vocab.bin" as *u8, 0x1a4); if mfd>=0 { sys_write(mfd, id2c, vv); sys_close(mfd) }
315 total=total+1; if nb>0 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) }
316 gw("T6 EXPORT: wrote /tmp/nx_ours_qwen.gguf (" as *u8); gn(nb); gw(" bytes, " as *u8); gn(3+NLAYERS*12); gw(" F32 tensors, transposed to inference layout + Qwen2 metadata) + vocab map\n" as *u8)
317
318 gw("\n FULL-FAITHFUL Qwen2: RMSNorm-GAMMA + QKV-BIAS added (via tile->leaf->scatter-reduce, same verified ta_ ops) --\n" as *u8)
319 gw(" so this model has EVERY tensor Qwen2 inference needs. Gradchecked incl the new params, trains, generates,\n" as *u8)
320 gw(" and EXPORTS to gguf. NEXT (loop-b gate): arch-config inference loads it -> ours-on-our-stack generation.\n" as *u8)
321 gw("NX-F32-QWEN2-TRAIN verdict=" as *u8)
322 if pass==total { gw("GREEN passes=" as *u8); gn(pass); gw("/" as *u8); gn(total); gw(" -- full-faithful Qwen2 trainer: gamma+bias gradchecked, trained, generating\n" as *u8); sys_exit(0); return 0 }
323 gw("RED passes=" as *u8); gn(pass); gw("/" as *u8); gn(total); gw("\n" as *u8); sys_exit(1); return 1
324}