nx_distributed_train_gate.nx source
↩ module page · 141 lines · 6518 B
1// nx_distributed_train_gate.nx -- HEADLINE EXCEED PROOF: deterministic DISTRIBUTED training.
2// Composes nx_fabric_collective::ring_allreduce (the byte-exact INTEGER gradient all-reduce) into a minimal
3// no-float data-parallel trainer (forward -> backward -> all-reduce gradient -> integer optimizer step) and
4// PROVES the trained weights are BIT-IDENTICAL regardless of the process count N (the data-parallel sharding).
5// This is the property NO float stack has: NCCL's float all-reduce is non-deterministic across nodes, and the
6// batch-invariant frontier (Thinking Machines) is single-node only. Ours holds because integer reduction is
7// associative + commutative -> the summed gradient is independent of how the global batch is sharded.
8// criteria:
9// 1 N=1 full-batch baseline trains AND W actually MOVES (real work, not all-zeros)
10// 2 N=2 data-parallel -> final W byte-identical to the N=1 baseline
11// 3 N=4 data-parallel -> final W byte-identical to the N=1 baseline
12// 4 N=8 data-parallel -> final W byte-identical to the N=1 baseline (deterministic across ALL shardings)
13// negative control:
14// NC1 N=8 WITHOUT the all-reduce (proc-0 local gradient only) -> W DIFFERS (the all-reduce is load-bearing)
15// expect_exit: 0 license_tier: ORIGINAL
16import "nx_fabric_collective.nx"
17import "nx_gate_verdict.nx"
18
19func dp(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
20func dn(v: i64) -> i64 {
21 let b: *u8 = sys_mmap(28); var m: i64 = v
22 if m < 0 { m = 0 - m; sys_write(1, "-" as *u8, 1) }
23 let t: *u8 = sys_mmap(28); var k: i64 = 0
24 if m == 0 { t[0] = 48 as u8; k = 1 }
25 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
26 var i: i64 = 0
27 while i < k { b[i] = t[k - 1 - i]; i = i + 1 }
28 sys_write(1, b, k); return 0
29}
30func chk(name: *u8, ok: i64) -> i64 {
31 if ok == 1 { dp(" PASS " as *u8); dp(name); dp("\n" as *u8); return 1 }
32 dp(" FAIL " as *u8); dp(name); dp("\n" as *u8); return 0
33}
34func printw(label: *u8, w: *i64, V: i64) -> i64 {
35 dp(label); var i: i64 = 0
36 while i < V { dp(" " as *u8); dn(w[i]); i = i + 1 }
37 dp("\n" as *u8); return 0
38}
39
40// deterministic fixed dataset (small integers)
41func dx(s: i64, d: i64) -> i64 { return (s * 3 + d * 7 + 1) % 11 }
42func dy(s: i64) -> i64 { return (s * 5 + 2) % 13 }
43
44// data-parallel train: N processes, M samples, V weight-dim, T steps, lr_div divisor; do_ar=1 all-reduce else
45// proc-0 local-only (the NC). Writes final W into out_w[0..V).
46func train(N: i64, M: i64, V: i64, T: i64, lr_div: i64, do_ar: i64, out_w: *i64) -> i64 {
47 let W: *i64 = sys_mmap(8 * V) as *i64
48 var z: i64 = 0
49 while z < V { W[z] = 0; z = z + 1 }
50 let gp: *i64 = sys_mmap(8 * N * V) as *i64
51 let gg: *i64 = sys_mmap(8 * V) as *i64
52 let shard: i64 = M / N
53 var step: i64 = 0
54 while step < T {
55 var k: i64 = 0
56 while k < N * V { gp[k] = 0; k = k + 1 }
57 // each process computes its LOCAL gradient over its shard of the global batch
58 var p: i64 = 0
59 while p < N {
60 var s: i64 = p * shard
61 let send: i64 = p * shard + shard
62 while s < send {
63 var pred: i64 = 0
64 var d: i64 = 0
65 while d < V { pred = pred + W[d] * dx(s, d); d = d + 1 }
66 let err: i64 = pred - dy(s)
67 d = 0
68 while d < V { gp[p * V + d] = gp[p * V + d] + err * dx(s, d); d = d + 1 }
69 s = s + 1
70 }
71 p = p + 1
72 }
73 // gradient sync: ring all-reduce -> the GLOBAL gradient (sum over the whole batch), or proc-0 local (NC)
74 if do_ar == 1 { ring_allreduce(gp, N, V, gg) }
75 else { var d2: i64 = 0; while d2 < V { gg[d2] = gp[d2]; d2 = d2 + 1 } }
76 // integer optimizer step (deterministic)
77 var d3: i64 = 0
78 while d3 < V { W[d3] = W[d3] - (gg[d3] / lr_div); d3 = d3 + 1 }
79 step = step + 1
80 }
81 var oi: i64 = 0
82 while oi < V { out_w[oi] = W[oi]; oi = oi + 1 }
83 return 0
84}
85
86func weq(a: *i64, b: *i64, V: i64) -> i64 {
87 var i: i64 = 0
88 var ok: i64 = 1
89 while i < V { if a[i] != b[i] { ok = 0 } i = i + 1 }
90 return ok
91}
92func wmoved(a: *i64, V: i64) -> i64 {
93 var i: i64 = 0
94 var moved: i64 = 0
95 while i < V { if a[i] != 0 { moved = 1 } i = i + 1 }
96 return moved
97}
98
99func main() -> i64 {
100 let V: i64 = 8
101 let M: i64 = 8
102 let T: i64 = 4
103 let LD: i64 = 16
104 dp("=== DETERMINISTIC DISTRIBUTED TRAINING -- no-float DDP via byte-exact integer ring all-reduce ===\n" as *u8)
105
106 let w1: *i64 = sys_mmap(8 * V) as *i64
107 let w2: *i64 = sys_mmap(8 * V) as *i64
108 let w4: *i64 = sys_mmap(8 * V) as *i64
109 let w8: *i64 = sys_mmap(8 * V) as *i64
110 let wnc: *i64 = sys_mmap(8 * V) as *i64
111 train(1, M, V, T, LD, 1, w1)
112 train(2, M, V, T, LD, 1, w2)
113 train(4, M, V, T, LD, 1, w4)
114 train(8, M, V, T, LD, 1, w8)
115 train(8, M, V, T, LD, 0, wnc)
116
117 printw(" W(N=1, full batch) =" as *u8, w1, V)
118 printw(" W(N=8, data-parallel) =" as *u8, w8, V)
119 printw(" W(N=8, NO all-reduce) =" as *u8, wnc, V)
120
121 var pass: i64 = 0
122 var total: i64 = 0
123 total = total + 1; pass = pass + chk("T1 N=1 baseline trains (W moved off zero)" as *u8, wmoved(w1, V))
124 total = total + 1; pass = pass + chk("T2 N=2 == N=1 (bit-identical)" as *u8, weq(w2, w1, V))
125 total = total + 1; pass = pass + chk("T3 N=4 == N=1 (bit-identical)" as *u8, weq(w4, w1, V))
126 total = total + 1; pass = pass + chk("T4 N=8 == N=1 (bit-identical) -- deterministic across ALL shardings" as *u8, weq(w8, w1, V))
127 var nc_ok: i64 = 0
128 if weq(wnc, w1, V) == 0 { nc_ok = 1 }
129 total = total + 1; pass = pass + chk("NC1 N=8 WITHOUT all-reduce DIFFERS (the all-reduce is load-bearing)" as *u8, nc_ok)
130
131 dp("NX-DISTRIBUTED-TRAIN-GATE " as *u8); dn(pass); dp(" / " as *u8); dn(total)
132 // MIGRATED onto nx_gate_verdict by nx_gate_dry_apply (D001, minimal form): every check
133 // row above is untouched, so the PASS/FAIL vector cannot change; only the hand-rolled
134 // verdict emission is replaced by the ONE shared base class. Proven by nx_gate_migrate verify.
135 let ctr__dry: *i64 = gv_ctr()
136 ctr__dry[0] = pass
137 ctr__dry[1] = total
138 let rc__dry: i64 = gv_verdict("DISTRIBUTED-TRAIN-GATE" as *u8, ctr__dry, "deterministic distributed training: weights bit-identical for ANY N -- the property no float stack has)" as *u8)
139 sys_exit(rc__dry)
140 return rc__dry
141}