code wiki / (root) / nx_ta_transformer_gradcheck_gate.nx
`nx_ta_transformer_gradcheck_gate.nx` source

↩ module page · 207 lines · 10564 B
1// nx_ta_transformer_gradcheck_gate.nx -- LIAR-KILL for the 9 f32 transformer ops added to nx_autograd_tensor
2// (2026-07-09, for the R3e neural reader). Each backward identity was HAND-PORTED from its Q16 nfa_* twin;
3// a transposed index or flipped sign would train the reader to garbage and waste the slow f32 run. This gate
4// verifies every one by FINITE DIFFERENCE: build op -> scalar loss (mse-vs-0, or the CE itself for softce) ->
5// analytic grad (ta_backward) vs numeric grad ((L(+h)-L(-h))/2h). A wrong identity => order-of-magnitude/sign
6// mismatch. Requires each op to show a NON-TRIVIAL gradient (|g|>=20 milli) so it cannot pass at 0==0.
7//   PASS <op> for each of: matmul, matmul_nt, cmul, softmax_rows(causal), rope_tab, hadamard, silu,
8//   rmsnorm_rows, softce_rows.
9// expect_exit: 0   license_tier: ORIGINAL   Sovereign: nx_autograd_tensor + nx_syscalls.
10import "nx_autograd_tensor.nx"
11import "nx_syscalls.nx"
12
13func gt_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
14func gt_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 }
15
16func gt_abs(x: i64) -> i64 { if x < 0 { return 0 - x } return x }
17// finite-diff compare in milli: pass if |ana-num|*100 <= |ana|*20 + 400 (20% rel + 0.4 abs, generous for f32
18// finite diff but catches any sign/index bug). returns pass(0/1); *nt set to 1 if this cell is non-trivial.
19func gt_cell(Lp: i64, Lm: i64, h2: i64, ana: i64, nt: *i64) -> i64 {
  let num: i64 = nx_f32_div(nx_f32_sub(Lp, Lm), h2)
  let am: i64 = ta_f32_to_milli(ana)
  let nm: i64 = ta_f32_to_milli(num)
  gt_puts("      ana="); gt_pn(am); gt_puts(" num="); gt_pn(nm); gt_puts("\n" as *u8)
  let aa: i64 = gt_abs(am)
  if aa >= 20 { nt[0] = 1 }
  let diff: i64 = gt_abs(am - nm)
  if diff*100 <= aa*20 + 400 { return 1 }
  return 0
29}
30
31// ---- per-op tape builders: return the scalar-root node. leaf 0 = the checked input (grad cells at arena 0). ----
32func bld_matmul(tape: *i64, vals: *i64, st: *i64, A: *i64, B: *i64, Z: *i64) -> i64 {
  st[0]=0; st[1]=0
  let nA: i64 = ta_leaf(tape,vals,st,2,3,A,0)
  let nB: i64 = ta_leaf(tape,vals,st,3,2,B,0)
  let nC: i64 = ta_matmul(tape,vals,st,nA,nB)
  let nZ: i64 = ta_leaf(tape,vals,st,2,2,Z,0)
  return ta_mse(tape,vals,st,nC,nZ)
39}
40func bld_matmul_nt(tape: *i64, vals: *i64, st: *i64, A: *i64, B: *i64, Z: *i64) -> i64 {
  st[0]=0; st[1]=0
  let nA: i64 = ta_leaf(tape,vals,st,2,3,A,0)
  let nB: i64 = ta_leaf(tape,vals,st,2,3,B,0)          // B[p=2, k=3] -> S[2,2]
  let nS: i64 = ta_matmul_nt(tape,vals,st,nA,nB)
  let nZ: i64 = ta_leaf(tape,vals,st,2,2,Z,0)
  return ta_mse(tape,vals,st,nS,nZ)
47}
48func bld_cmul(tape: *i64, vals: *i64, st: *i64, A: *i64, Z: *i64, c: i64) -> i64 {
  st[0]=0; st[1]=0
  let nA: i64 = ta_leaf(tape,vals,st,2,3,A,0)
  let nY: i64 = ta_cmul(tape,vals,st,nA,c)
  let nZ: i64 = ta_leaf(tape,vals,st,2,3,Z,0)
  return ta_mse(tape,vals,st,nY,nZ)
54}
55func bld_smrows(tape: *i64, vals: *i64, st: *i64, A: *i64, Z: *i64) -> i64 {
  st[0]=0; st[1]=0
  let nA: i64 = ta_leaf(tape,vals,st,3,3,A,0)
  let nY: i64 = ta_softmax_rows(tape,vals,st,nA,1)     // causal
  let nZ: i64 = ta_leaf(tape,vals,st,3,3,Z,0)
  return ta_mse(tape,vals,st,nY,nZ)
61}
62func bld_rope(tape: *i64, vals: *i64, st: *i64, A: *i64, Z: *i64, tab: *i64) -> i64 {
  st[0]=0; st[1]=0
  let nA: i64 = ta_leaf(tape,vals,st,2,4,A,0)          // T=2, hd=4
  let nY: i64 = ta_rope_tab(tape,vals,st,nA,tab)
  let nZ: i64 = ta_leaf(tape,vals,st,2,4,Z,0)
  return ta_mse(tape,vals,st,nY,nZ)
68}
69func bld_had(tape: *i64, vals: *i64, st: *i64, A: *i64, B: *i64, Z: *i64) -> i64 {
  st[0]=0; st[1]=0
  let nA: i64 = ta_leaf(tape,vals,st,2,3,A,0)
  let nB: i64 = ta_leaf(tape,vals,st,2,3,B,0)
  let nY: i64 = ta_hadamard(tape,vals,st,nA,nB)
  let nZ: i64 = ta_leaf(tape,vals,st,2,3,Z,0)
  return ta_mse(tape,vals,st,nY,nZ)
76}
77func bld_silu(tape: *i64, vals: *i64, st: *i64, A: *i64, Z: *i64) -> i64 {
  st[0]=0; st[1]=0
  let nA: i64 = ta_leaf(tape,vals,st,2,3,A,0)
  let nY: i64 = ta_silu(tape,vals,st,nA)
  let nZ: i64 = ta_leaf(tape,vals,st,2,3,Z,0)
  return ta_mse(tape,vals,st,nY,nZ)
83}
84func bld_rms(tape: *i64, vals: *i64, st: *i64, A: *i64, Z: *i64) -> i64 {
  st[0]=0; st[1]=0
  let nA: i64 = ta_leaf(tape,vals,st,2,3,A,0)
  let nY: i64 = ta_rmsnorm_rows(tape,vals,st,nA)
  let nZ: i64 = ta_leaf(tape,vals,st,2,3,Z,0)
  return ta_mse(tape,vals,st,nY,nZ)
90}
91func bld_softce(tape: *i64, vals: *i64, st: *i64, A: *i64, ids: *i64) -> i64 {
  st[0]=0; st[1]=0
  let nA: i64 = ta_leaf(tape,vals,st,2,3,A,0)          // logits [T=2, V=3]
  return ta_softce_rows(tape,vals,st,nA,ids)
95}
96
97func main() -> i64 {
  gt_puts("nx_ta_transformer_gradcheck_gate (finite-diff gradcheck of the 9 f32 transformer ops behind the R3e reader)\n" as *u8)
  let tape: *i64 = sys_mmap(512*7*8) as *i64
  let vals: *i64 = sys_mmap(4096*8) as *i64
  let grads: *i64 = sys_mmap(4096*8) as *i64
  let st: *i64 = sys_mmap(2*8) as *i64
  let A: *i64 = sys_mmap(64*8) as *i64
  let B: *i64 = sys_mmap(64*8) as *i64
  let Z: *i64 = sys_mmap(64*8) as *i64
  let ids: *i64 = sys_mmap(8*8) as *i64
  let tab: *i64 = sys_mmap((2 + 2*2*2)*8) as *i64       // T=2, np=2
  var i: i64 = 0
  // Z = a NON-ZERO mse target. mse-vs-ZERO is scale-invariant under rmsnorm/softmax (‖normalized‖²≈const ->
  // gradient structurally ~0 -> a degenerate test), so use a real target: the gradient is then non-trivial
  // and the analytic-vs-numeric match becomes a meaningful liar-kill for every op.
  while i < 64 { A[i] = ta_constf(((i*7+1)%9)-4, 6); B[i] = ta_constf(((i*5+2)%9)-4, 6); Z[i] = ta_constf(((i*3+1)%7)-3, 5); i = i + 1 }
  ids[0] = 1; ids[1] = 2
  ta_rope_build_tab(tab, 2, 2)
  let h: i64 = ta_constf(1, 256)
  let h2: i64 = nx_f32_mul(nx_i32_to_f32(2), h)
  let chalf: i64 = ta_constf(1, 2)                      // cmul test constant (0.5; not the tiny perturbation h)
  var pass: i64 = 0
  var total: i64 = 0
120
  // ---- helper: run one op's gradcheck over `ncell` cells of leaf-0's source `src`. opid selects the builder. ----
  // (dispatch inline; no fn-ptr. cells checked = indices cs[0..ncell).)
  let cs: *i64 = sys_mmap(8*8) as *i64
  cs[0]=0; cs[1]=2; cs[2]=4
125
  // ===== per-op driver, unrolled by opid (0..8) =====
  var opid: i64 = 0
  while opid < 9 {
      // pick name + src + ncell
      var name: *u8 = "?" as *u8
      var src: *i64 = A
      var ncell: i64 = 3
      if opid==0 { name = "matmul" as *u8 }
      if opid==1 { name = "matmul_nt" as *u8 }
      if opid==2 { name = "cmul" as *u8 }
      if opid==3 { name = "softmax_rows" as *u8 }
      if opid==4 { name = "rope_tab" as *u8 }
      if opid==5 { name = "hadamard" as *u8 }
      if opid==6 { name = "silu" as *u8 }
      if opid==7 { name = "rmsnorm_rows" as *u8 }
      if opid==8 { name = "softce_rows" as *u8 }
      gt_puts("  ["); gt_puts(name); gt_puts("]\n" as *u8)
143
      // analytic: build once, backward, read grads of leaf-0 (arena offset 0)
      var root: i64 = 0
      if opid==0 { root = bld_matmul(tape,vals,st,A,B,Z) }
      if opid==1 { root = bld_matmul_nt(tape,vals,st,A,B,Z) }
      if opid==2 { root = bld_cmul(tape,vals,st,A,Z,chalf) }
      if opid==3 { root = bld_smrows(tape,vals,st,A,Z) }
      if opid==4 { root = bld_rope(tape,vals,st,A,Z,tab) }
      if opid==5 { root = bld_had(tape,vals,st,A,B,Z) }
      if opid==6 { root = bld_silu(tape,vals,st,A,Z) }
      if opid==7 { root = bld_rms(tape,vals,st,A,Z) }
      if opid==8 { root = bld_softce(tape,vals,st,A,ids) }
      ta_backward(tape, vals, grads, st[0], root)
156
      var allok: i64 = 1
      var nt: *i64 = sys_mmap(8)
      nt[0] = 0
      var ci: i64 = 0
      while ci < ncell {
          let c: i64 = cs[ci]
          let ana: i64 = grads[0 + c]                  // leaf-0 grad cell c
          let sv: i64 = src[c]
          // L(+h)
          src[c] = nx_f32_add(sv, h)
          var rp: i64 = 0
          if opid==0 { rp = bld_matmul(tape,vals,st,A,B,Z) }
          if opid==1 { rp = bld_matmul_nt(tape,vals,st,A,B,Z) }
          if opid==2 { rp = bld_cmul(tape,vals,st,A,Z,chalf) }
          if opid==3 { rp = bld_smrows(tape,vals,st,A,Z) }
          if opid==4 { rp = bld_rope(tape,vals,st,A,Z,tab) }
          if opid==5 { rp = bld_had(tape,vals,st,A,B,Z) }
          if opid==6 { rp = bld_silu(tape,vals,st,A,Z) }
          if opid==7 { rp = bld_rms(tape,vals,st,A,Z) }
          if opid==8 { rp = bld_softce(tape,vals,st,A,ids) }
          let Lp: i64 = ta_val(tape, vals, rp, 0)
          // L(-h)
          src[c] = nx_f32_sub(sv, h)
          var rm: i64 = 0
          if opid==0 { rm = bld_matmul(tape,vals,st,A,B,Z) }
          if opid==1 { rm = bld_matmul_nt(tape,vals,st,A,B,Z) }
          if opid==2 { rm = bld_cmul(tape,vals,st,A,Z,chalf) }
          if opid==3 { rm = bld_smrows(tape,vals,st,A,Z) }
          if opid==4 { rm = bld_rope(tape,vals,st,A,Z,tab) }
          if opid==5 { rm = bld_had(tape,vals,st,A,B,Z) }
          if opid==6 { rm = bld_silu(tape,vals,st,A,Z) }
          if opid==7 { rm = bld_rms(tape,vals,st,A,Z) }
          if opid==8 { rm = bld_softce(tape,vals,st,A,ids) }
          let Lm: i64 = ta_val(tape, vals, rm, 0)
          src[c] = sv
          let ok: i64 = gt_cell(Lp, Lm, h2, ana, nt)
          if ok == 0 { allok = 0 }
          ci = ci + 1
      }
      total = total + 1
      var v: i64 = 0
      if allok == 1 { if nt[0] == 1 { v = 1 } }
      if v == 1 { pass = pass + 1; gt_puts("  PASS " as *u8); gt_puts(name); gt_puts("\n" as *u8) } else { gt_puts("  FAIL " as *u8); gt_puts(name); gt_puts("\n" as *u8) }
      opid = opid + 1
  }
202
  gt_puts("---- nx_ta_transformer_gradcheck_gate: passed "); gt_pn(pass); gt_puts(" / "); gt_pn(total); gt_puts("\n" as *u8)
  if pass == total { gt_puts("TA-TRANSFORMER GRADCHECK GREEN -- all 9 f32 transformer backward identities match finite-difference; the R3e f32 reader's gradients are CORRECT.\n" as *u8); return 0 }
  gt_puts("RED -- a backward identity is WRONG (see the ana/num pair that diverged) -- fix before trusting the f32 reader run.\n" as *u8)
  return 1
207}
nx_ta_transformer_gradcheck_gate.nx source

`nx_ta_transformer_gradcheck_gate.nx` source
