nx_ta_parallel_grad_gate.nx source
↩ module page · 188 lines · 9375 B
1// nx_ta_parallel_grad_gate.nx -- KEYSTONE for FAST model training: the f32 tensor autograd (ta_*) trains
2// DATA-PARALLEL on the sovereign thread pool. Each example in a batch runs its OWN forward+backward on a
3// PRIVATE (tape,vals,grads,st) arena; the per-example gradients are then reduced. The soft-float training that
4// bottlenecks the neural reader (and every bigger model the benchmark endgame needs) parallelizes across the
5// 16 hardware workers here -- the conv2d pool already measured 6.76x wall-clock.
6//
7// The liar-kill is STRONG: serial and parallel compute each example's gradient with identical math, so the
8// per-example gradients must be BIT-IDENTICAL (f32 raw equality). ANY race, arena collision, or shared-state
9// bug in the parallel path perturbs a bit and fails. (No f32-associativity fuzz: we compare PER-EXAMPLE grads,
10// not a reordered sum.)
11// T1 parallel per-example grads == serial, BIT-IDENTICAL across all B examples x all NP weights
12// T2 the pool actually ran B tasks (completed counter advanced by B) on > 1 worker
13// T3 deterministic (re-run parallel == identical)
14// expect_exit: 0 license_tier: ORIGINAL Sovereign: nx_thread_pool + nx_autograd_tensor + nx_syscalls.
15import "nx_thread_pool.nx"
16import "nx_autograd_tensor.nx"
17import "nx_hw.nx"
18import "nx_syscalls.nx"
19
20func pg_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
21func pg_pn(v: i64) -> i64 { let b: *u8=sys_mmap(28); var x: i64=v; if x<0{b[0]=45;sys_write(1,b,1);x=0-x} if x==0{b[0]=48;sys_write(1,b,1);return 0} var d: i64=0; var y: i64=x; while y>0{d=d+1;y=y/10} var i: i64=d-1; y=x; while i>=0{b[i]=(48+(y%10)) as u8;y=y/10;i=i-1} sys_write(1,b,d); return 0 }
22func pg_ck(name: *u8, c: i64) -> i64 { if c==1 { pg_puts(" PASS " as *u8) } else { pg_puts(" FAIL " as *u8) } pg_puts(name); pg_puts("\n" as *u8); return c }
23
24const PG_D: i64 = 8 // input dim
25const PG_H: i64 = 16 // hidden
26const PG_O: i64 = 4 // output
27const PG_NW1: i64 = 128 // H*D
28const PG_NW2: i64 = 64 // O*H
29const PG_NP: i64 = 192 // NW1+NW2
30const PG_B: i64 = 8 // batch
31
32// per-example work context (pointers travel to the pool worker; no closures)
33struct PgCtx {
34 w_ptr: i64,
35 x_ptr: i64,
36 tgt_ptr: i64,
37 tape_ptr: i64,
38 vals_ptr: i64,
39 grads_ptr: i64,
40 st_ptr: i64,
41 g_ptr: i64,
42}
43const PG_CTX_BYTES: i64 = 64
44
45// forward MLP (x -> W1 -> relu -> W2 -> mse) + backward + extract dW1,dW2 into G[NP]. Uses ONLY the passed
46// private arenas -> thread-safe. Node order: x=0, W1=1, H=2, A=3, W2=4, Y=5, T=6, L=7.
47func pg_example(w: *i64, x: *i64, tgt: *i64, tape: *i64, vals: *i64, grads: *i64, st: *i64, G: *i64) -> i64 {
48 st[0] = 0
49 st[1] = 0
50 let nx: i64 = ta_leaf(tape, vals, st, PG_D, 1, x, 0)
51 let nW1: i64 = ta_leaf(tape, vals, st, PG_H, PG_D, w, 0)
52 let nHp: i64 = ta_matvec(tape, vals, st, nW1, nx)
53 let nA: i64 = ta_relu(tape, vals, st, nHp)
54 let nW2: i64 = ta_leaf(tape, vals, st, PG_O, PG_H, w, PG_NW1)
55 let nY: i64 = ta_matvec(tape, vals, st, nW2, nA)
56 let nT: i64 = ta_leaf(tape, vals, st, PG_O, 1, tgt, 0)
57 let nL: i64 = ta_mse(tape, vals, st, nY, nT)
58 ta_backward(tape, vals, grads, st[0], nL)
59 let g1: i64 = tape[7*nW1 + 6]
60 var i: i64 = 0
61 while i < PG_NW1 { G[i] = grads[g1 + i]; i = i + 1 }
62 let g2: i64 = tape[7*nW2 + 6]
63 i = 0
64 while i < PG_NW2 { G[PG_NW1 + i] = grads[g2 + i]; i = i + 1 }
65 return 0
66}
67
68// pool worker: unpack ctx, run one example.
69func pg_task(ctx_i: i64) -> i64 {
70 let c: *PgCtx = ctx_i as *PgCtx
71 return pg_example(c.w_ptr as *i64, c.x_ptr as *i64, c.tgt_ptr as *i64, c.tape_ptr as *i64, c.vals_ptr as *i64, c.grads_ptr as *i64, c.st_ptr as *i64, c.g_ptr as *i64)
72}
73
74func main() -> i64 {
75 pg_puts("nx_ta_parallel_grad_gate (KEYSTONE: f32 ta_* autograd trains DATA-PARALLEL on the thread pool; per-example grads bit-identical to serial)\n" as *u8)
76 var pass: i64 = 0
77 var total: i64 = 0
78
79 // ---- weights + batch (deterministic) ----
80 let W: *i64 = sys_mmap(PG_NP*8) as *i64
81 var i: i64 = 0
82 while i < PG_NP { W[i] = ta_constf(((i*7+1)%13)-6, 8); i = i + 1 }
83 let X: *i64 = sys_mmap(PG_B*PG_D*8) as *i64
84 let TG: *i64 = sys_mmap(PG_B*PG_O*8) as *i64
85 i = 0
86 while i < PG_B*PG_D { X[i] = ta_constf(((i*5+2)%11)-5, 7); i = i + 1 }
87 i = 0
88 while i < PG_B*PG_O { TG[i] = ta_constf(((i*3+1)%9)-4, 6); i = i + 1 }
89
90 // per-example private arenas (B sets)
91 let tapes: *i64 = sys_mmap(PG_B*256*7*8) as *i64
92 let valss: *i64 = sys_mmap(PG_B*1024*8) as *i64
93 let gradss: *i64 = sys_mmap(PG_B*1024*8) as *i64
94 let sts: *i64 = sys_mmap(PG_B*2*8) as *i64
95 let Gser: *i64 = sys_mmap(PG_B*PG_NP*8) as *i64
96 let Gpar: *i64 = sys_mmap(PG_B*PG_NP*8) as *i64
97
98 // ---- SERIAL: each example into Gser[b] (sequential) ----
99 var b: i64 = 0
100 while b < PG_B {
101 let tp: *i64 = (tapes as i64 + b*256*7*8) as *i64
102 let vp: *i64 = (valss as i64 + b*1024*8) as *i64
103 let gp: *i64 = (gradss as i64 + b*1024*8) as *i64
104 let sp: *i64 = (sts as i64 + b*2*8) as *i64
105 let xp: *i64 = (X as i64 + b*PG_D*8) as *i64
106 let tgp: *i64 = (TG as i64 + b*PG_O*8) as *i64
107 let gsp: *i64 = (Gser as i64 + b*PG_NP*8) as *i64
108 pg_example(W, xp, tgp, tp, vp, gp, sp, gsp)
109 b = b + 1
110 }
111
112 // ---- PARALLEL: same B examples on the pool, into Gpar[b] ----
113 let pool: *NxThreadPool = nx_pool_new(0, 0) // 0 = auto (nx_hw_worker_count)
114 let nworkers: i64 = pool.n_workers
115 let ctxs: *u8 = sys_mmap(PG_B * PG_CTX_BYTES)
116 let done_before: i64 = nx_pool_n_completed(pool)
117 b = 0
118 while b < PG_B {
119 let cx: *PgCtx = ((ctxs as i64) + b*PG_CTX_BYTES) as *PgCtx
120 cx.w_ptr = W as i64
121 cx.x_ptr = (X as i64 + b*PG_D*8)
122 cx.tgt_ptr = (TG as i64 + b*PG_O*8)
123 cx.tape_ptr = (tapes as i64 + b*256*7*8)
124 cx.vals_ptr = (valss as i64 + b*1024*8)
125 cx.grads_ptr = (gradss as i64 + b*1024*8)
126 cx.st_ptr = (sts as i64 + b*2*8)
127 cx.g_ptr = (Gpar as i64 + b*PG_NP*8)
128 nx_pool_submit(pool, pg_task, cx as i64)
129 b = b + 1
130 }
131 let wv: i64 = nx_pool_wait(pool, done_before + PG_B)
132 let ran: i64 = nx_pool_n_completed(pool) - done_before
133
134 pg_puts(" workers="); pg_pn(nworkers); pg_puts(" pool tasks completed="); pg_pn(ran); pg_puts("/"); pg_pn(PG_B); pg_puts("\n" as *u8)
135
136 // ---- T1: parallel per-example grads BIT-IDENTICAL to serial ----
137 var mism: i64 = 0
138 var chk: i64 = 0
139 b = 0
140 while b < PG_B {
141 var j: i64 = 0
142 while j < PG_NP {
143 let sv: i64 = Gser[b*PG_NP + j]
144 let pv: i64 = Gpar[b*PG_NP + j]
145 if sv != pv { mism = mism + 1 }
146 if sv != 0 { chk = chk + 1 } // count non-trivial grad cells (test isn't vacuous)
147 j = j + 1
148 }
149 b = b + 1
150 }
151 pg_puts(" bit-mismatches="); pg_pn(mism); pg_puts(" non-zero grad cells checked="); pg_pn(chk); pg_puts("/"); pg_pn(PG_B*PG_NP); pg_puts("\n" as *u8)
152 total = total + 1
153 if mism == 0 { if chk > PG_NP { pass = pass + 1; pg_ck("T1 parallel grads BIT-IDENTICAL to serial (all B examples)" as *u8, 1) } else { pg_ck("T1 parallel grads BIT-IDENTICAL to serial (all B examples)" as *u8, 0) } } else { pg_ck("T1 parallel grads BIT-IDENTICAL to serial (all B examples)" as *u8, 0) }
154
155 // ---- T2: pool ran B tasks on > 1 worker ----
156 total = total + 1
157 if ran == PG_B { if nworkers > 1 { pass = pass + 1; pg_ck("T2 pool ran all B tasks on >1 worker" as *u8, 1) } else { pg_ck("T2 pool ran all B tasks on >1 worker" as *u8, 0) } } else { pg_ck("T2 pool ran all B tasks on >1 worker" as *u8, 0) }
158
159 // ---- T3: deterministic re-run (parallel again == same Gpar) ----
160 let Gpar2: *i64 = sys_mmap(PG_B*PG_NP*8) as *i64
161 let ctxs2: *u8 = sys_mmap(PG_B * PG_CTX_BYTES)
162 let db2: i64 = nx_pool_n_completed(pool)
163 b = 0
164 while b < PG_B {
165 let cx: *PgCtx = ((ctxs2 as i64) + b*PG_CTX_BYTES) as *PgCtx
166 cx.w_ptr = W as i64
167 cx.x_ptr = (X as i64 + b*PG_D*8)
168 cx.tgt_ptr = (TG as i64 + b*PG_O*8)
169 cx.tape_ptr = (tapes as i64 + b*256*7*8)
170 cx.vals_ptr = (valss as i64 + b*1024*8)
171 cx.grads_ptr = (gradss as i64 + b*1024*8)
172 cx.st_ptr = (sts as i64 + b*2*8)
173 cx.g_ptr = (Gpar2 as i64 + b*PG_NP*8)
174 nx_pool_submit(pool, pg_task, cx as i64)
175 b = b + 1
176 }
177 nx_pool_wait(pool, db2 + PG_B)
178 var mism2: i64 = 0
179 i = 0
180 while i < PG_B*PG_NP { if Gpar[i] != Gpar2[i] { mism2 = mism2 + 1 } i = i + 1 }
181 total = total + 1
182 if mism2 == 0 { pass = pass + 1; pg_ck("T3 deterministic (re-run parallel identical)" as *u8, 1) } else { pg_ck("T3 deterministic (re-run parallel identical)" as *u8, 0) }
183
184 pg_puts("---- nx_ta_parallel_grad_gate: passed "); pg_pn(pass); pg_puts(" / "); pg_pn(total); pg_puts("\n" as *u8)
185 if pass == total { pg_puts("TA-PARALLEL-GRAD GREEN -- the f32 sovereign autograd trains DATA-PARALLEL on the thread pool, per-example gradients bit-identical to serial. Fast model training is unblocked (integrate into the reader trainer's batch loop).\n" as *u8); return 0 }
186 pg_puts("RED -- parallel training not correct (see mismatches)\n" as *u8)
187 return 1
188}