nx_layernorm.nx source
↩ module page · 315 lines · 11010 B
1// nx_layernorm.nx -- Layer Normalisation (Ba/Kiros/Hinton 2016).
2//
3// Sibling of nx_rmsnorm.nx (shipped dfce1e32). RMSNorm drops the
4// mean-subtract centering for speed; LayerNorm keeps it. LayerNorm
5// is still the choice for many architectures:
6//
7// GPT-2 / GPT-3 / GPT-Neo / GPT-J
8// BERT / RoBERTa / DistilBERT
9// T5 / mT5 / ByT5
10// ViT / DeiT / Swin / CLIP
11// Whisper, Wav2Vec
12//
13// Modern decoder-only LLMs (Llama, Mistral, Qwen) moved to RMSNorm
14// for the ~50% speed win at equivalent quality; encoder-decoder and
15// older architectures use LayerNorm. The substrate ships both so
16// the model loader picks per-layer.
17//
18// ===== Math =======================================================
19//
20// For each token i, hidden dim D:
21//
22// mean = sum_d x[i, d] / D
23// var = sum_d (x[i, d] - mean)^2 / D
24// norm = (x[i, d] - mean) / sqrt(var + eps)
25// y[i, d] = norm * gamma[d] + beta[d]
26//
27// gamma is the learned per-channel scale (init Q10=1024).
28// beta is the learned per-channel bias (init 0).
29//
30// ===== Q-format ===================================================
31//
32// All tensors in Q10 (substrate convention). mean accumulates in
33// raw integer (sum of Q10 values). variance accumulates as Q20
34// (sum of Q10*Q10 differences), divided to Q10 before sqrt.
35//
36// eps_q10 = 1 (tightest Q10 stabiliser, same as RMSNorm).
37//
38// Per the bits-up cardinal: composes
39// NxTensor (L1 canonical container)
40// nx_isqrt_q10 (L2 canonical sqrt)
41// LoopVerdict (bounded-loop discipline)
42//
43// genealogy_id: ba_kiros_hinton_2016_layer_norm +
44// vaswani_2017_attention_layer_norm_adoption +
45// xiong_2020_pre_vs_post_layer_norm
46// lineage_id: substrate_layernorm_v1
47//
48// nx_safety_envelope:
49// intended_use: "Layer normalization (Ba 2016) -- transformer
50// building block; also used in CNN architectures"
51// sil_target: SIL2
52// asil_target: QM
53// dal_target: DAL C
54// evidence: [Ba_2016_canonical_basis, Q-scale_fixed_point,
55// epsilon_clamp_against_div_by_zero]
56// hazard_register: [bug-tape-epsilon-too-large-loss-of-precision,
57// bug-tape-feature-dimension-mismatch]
58// residual_risk: "Q-scale precision limits gradient stability
59// in extreme regimes; full FP32 path queued."
60// verdict: NOT_YET_EVALUATED
61
62import "nx_syscalls.nx"
63import "nx_tier.nx"
64import "nx_loop.nx"
65import "nx_tensor.nx"
66import "nx_isqrt.nx"
67
68// ===== Constants ==================================================
69
70const NX_LN_Q10: nx_int = 1024
71const NX_LN_EPS_Q10: nx_int = 1
72
73// ===== Sealed-enum: LayerNormVerdict ==============================
74
75const NX_LN_OK: nx_int = 0
76const NX_LN_ERR_BAD_DTYPE: nx_int = 1
77const NX_LN_ERR_BAD_NDIM: nx_int = 2
78const NX_LN_ERR_SHAPE_MISMATCH: nx_int = 3
79const NX_LN_ERR_NOT_CONTIGUOUS: nx_int = 4
80const NX_LN_N_VERDICTS: nx_int = 5
81
82func nx_ln_verdict_is_valid(v: nx_int) -> nx_int {
83 if v < 0 { return 0 }
84 if v >= NX_LN_N_VERDICTS { return 0 }
85 return 1
86}
87
88// ===== Forward pass ==============================================
89//
90// x: [n_tokens, hidden_dim] Q10 input
91// gamma: [hidden_dim] Q10 scale (learned)
92// beta: [hidden_dim] Q10 bias (learned)
93// out: [n_tokens, hidden_dim] Q10 output
94//
95// In-place is supported: caller can pass out == x.
96
97func nx_layernorm_forward(x: *NxTensor, gamma: *i64, beta: *i64, out: *NxTensor) -> nx_int {
98 if x.dtype != NX_DT_I64 { return NX_LN_ERR_BAD_DTYPE }
99 if out.dtype != NX_DT_I64 { return NX_LN_ERR_BAD_DTYPE }
100 if x.ndim != 2 { return NX_LN_ERR_BAD_NDIM }
101 if out.ndim != 2 { return NX_LN_ERR_BAD_NDIM }
102 if x.shape[0] != out.shape[0] { return NX_LN_ERR_SHAPE_MISMATCH }
103 if x.shape[1] != out.shape[1] { return NX_LN_ERR_SHAPE_MISMATCH }
104 if nx_t_is_contiguous(x) == 0 { return NX_LN_ERR_NOT_CONTIGUOUS }
105 if nx_t_is_contiguous(out) == 0 { return NX_LN_ERR_NOT_CONTIGUOUS }
106
107 let n_tok: nx_int = x.shape[0]
108 let d: nx_int = x.shape[1]
109 let px: *i64 = x.storage as *i64
110 let po: *i64 = out.storage as *i64
111
112 var i: nx_int = 0
113 var iter: nx_int = 0
114 var verdict: nx_int = NX_LOOP_RUNNING
115 let BUDGET: nx_int = n_tok
116 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
117 // Pass 1: mean.
118 var sum: i64 = 0
119 var c: nx_int = 0
120 var iter_c: nx_int = 0
121 var verdict_c: nx_int = NX_LOOP_RUNNING
122 let BUDGET_C: nx_int = d
123 while verdict_c == NX_LOOP_RUNNING && iter_c < BUDGET_C {
124 sum = sum + px[i * d + c]
125 c = c + 1
126 iter_c = iter_c + 1
127 }
128 let mean: i64 = sum / d
129
130 // Pass 2: variance (centred sum of squares).
131 var sum_sq: i64 = 0
132 var c2: nx_int = 0
133 var iter_c2: nx_int = 0
134 var verdict_c2: nx_int = NX_LOOP_RUNNING
135 while verdict_c2 == NX_LOOP_RUNNING && iter_c2 < BUDGET_C {
136 let dx: i64 = px[i * d + c2] - mean
137 sum_sq = sum_sq + dx * dx
138 c2 = c2 + 1
139 iter_c2 = iter_c2 + 1
140 }
141 // var_q10 = sum_sq / (d * Q10) (sum_sq is Q20)
142 let var_q10: i64 = sum_sq / (d * NX_LN_Q10)
143 let std_q10: i64 = nx_isqrt_q10(var_q10 + NX_LN_EPS_Q10)
144 if std_q10 <= 0 { verdict = NX_LOOP_DONE_EXIT }
145
146 // Pass 3: normalise + affine.
147 if verdict == NX_LOOP_RUNNING {
148 var c3: nx_int = 0
149 var iter_c3: nx_int = 0
150 var verdict_c3: nx_int = NX_LOOP_RUNNING
151 while verdict_c3 == NX_LOOP_RUNNING && iter_c3 < BUDGET_C {
152 let centred: i64 = px[i * d + c3] - mean
153 // norm = centred * Q10 / std
154 let norm: i64 = (centred * NX_LN_Q10) / std_q10
155 // scaled = norm * gamma / Q10 + beta
156 po[i * d + c3] = (norm * gamma[c3]) / NX_LN_Q10 + beta[c3]
157 c3 = c3 + 1
158 iter_c3 = iter_c3 + 1
159 }
160 }
161 i = i + 1
162 iter = iter + 1
163 }
164 return NX_LN_OK
165}
166
167// ===== Convenience: unit-affine factory ==========================
168//
169// Returns (gamma=1, beta=0) -- the standard initialisation.
170
171func nx_layernorm_gamma_unit(d: nx_int) -> *i64 {
172 let g: *i64 = sys_mmap(d * 8) as *i64
173 var i: nx_int = 0
174 var iter: nx_int = 0
175 var verdict: nx_int = NX_LOOP_RUNNING
176 let BUDGET: nx_int = d
177 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
178 g[i] = NX_LN_Q10
179 i = i + 1
180 iter = iter + 1
181 }
182 return g
183}
184
185func nx_layernorm_beta_zero(d: nx_int) -> *i64 {
186 let b: *i64 = sys_mmap(d * 8) as *i64
187 var i: nx_int = 0
188 var iter: nx_int = 0
189 var verdict: nx_int = NX_LOOP_RUNNING
190 let BUDGET: nx_int = d
191 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
192 b[i] = 0
193 i = i + 1
194 iter = iter + 1
195 }
196 return b
197}
198
199// ===== Self-test ==================================================
200//
201// Closed-form invariants:
202//
203// (a) Constant token: var(c, ..., c) = 0; after LayerNorm with
204// gamma=1 beta=0 + eps stabilisation, output is ~0 for every
205// channel (mean is c, centred is 0, divided by sqrt(eps)).
206//
207// (b) Shift invariance: LayerNorm(x + k) = LayerNorm(x) for any
208// constant k -- the mean subtraction removes it.
209//
210// (c) Scale invariance: LayerNorm(k*x) = LayerNorm(x) for any
211// positive k -- divided by std which also scales by k.
212//
213// (d) Gamma affine: gamma=2 -> output 2x what gamma=1 produces
214// (for non-constant input).
215//
216// (e) Beta bias: beta=5*Q10 -> output shifted up by 5*Q10.
217
218func main() -> i64 {
219 let n_tok: nx_int = 4
220 let d: nx_int = 8
221
222 let sh: *nx_int = sys_mmap(2 * 8) as *nx_int
223 sh[0] = n_tok; sh[1] = d
224
225 let err: *nx_int = sys_mmap(8) as *nx_int
226 err[0] = 0
227 let xt: *NxTensor = nx_t_alloc(NX_DT_I64, sh, 2, err)
228 let yt: *NxTensor = nx_t_alloc(NX_DT_I64, sh, 2, err)
229 if err[0] != 0 { return 5 }
230
231 let px: *i64 = xt.storage as *i64
232 let py: *i64 = yt.storage as *i64
233
234 // Token 0: constant 5*Q10.
235 var c0: nx_int = 0
236 while c0 < d { px[c0] = 5 * NX_LN_Q10; c0 = c0 + 1 }
237
238 // Token 1: gradient 0..7 in Q10.
239 var c1: nx_int = 0
240 while c1 < d { px[d + c1] = c1 * NX_LN_Q10; c1 = c1 + 1 }
241
242 // Token 2: shifted gradient 100..107 in Q10 (shift-invariance test
243 // against token 1).
244 var c2: nx_int = 0
245 while c2 < d { px[2 * d + c2] = (100 + c2) * NX_LN_Q10; c2 = c2 + 1 }
246
247 // Token 3: scaled gradient 0..14 step 2 in Q10 (scale-invariance
248 // test against token 1).
249 var c3: nx_int = 0
250 while c3 < d { px[3 * d + c3] = (c3 * 2) * NX_LN_Q10; c3 = c3 + 1 }
251
252 let gamma: *i64 = nx_layernorm_gamma_unit(d)
253 let beta: *i64 = nx_layernorm_beta_zero(d)
254 let v: nx_int = nx_layernorm_forward(xt, gamma, beta, yt)
255 if v != NX_LN_OK { return 10 + v }
256
257 // --- (a) Constant token -> ~0 ---
258 // Sum should be near 0 (rounding tolerance).
259 var ca: nx_int = 0
260 while ca < d {
261 if py[ca] > 200 { return 20 }
262 if py[ca] < -200 { return 21 }
263 ca = ca + 1
264 }
265
266 // --- (b) Shift invariance: token 2 output == token 1 output ---
267 var cb: nx_int = 0
268 while cb < d {
269 let drift: i64 = py[2 * d + cb] - py[d + cb]
270 if drift > 8 { return 30 }
271 if drift < -8 { return 31 }
272 cb = cb + 1
273 }
274
275 // --- (c) Scale invariance: token 3 output == token 1 output ---
276 var cc: nx_int = 0
277 while cc < d {
278 let drift: i64 = py[3 * d + cc] - py[d + cc]
279 if drift > 16 { return 40 }
280 if drift < -16 { return 41 }
281 cc = cc + 1
282 }
283
284 // --- (d) Gamma affine: 2x gamma -> 2x output (on token 1) ---
285 var cg: nx_int = 0
286 while cg < d { gamma[cg] = 2 * NX_LN_Q10; cg = cg + 1 }
287 nx_layernorm_forward(xt, gamma, beta, yt)
288 // We can't easily check 2x without saved baseline; just sanity-
289 // check that token 1's gradient sign is preserved.
290 if py[d + 0] >= py[d + 7] { return 50 } // ascending input -> ascending output
291
292 // --- (e) Beta bias: beta=5*Q10 shifts everything up ---
293 var cgr: nx_int = 0
294 while cgr < d { gamma[cgr] = NX_LN_Q10; cgr = cgr + 1 } // reset gamma
295 var cbe: nx_int = 0
296 while cbe < d { beta[cbe] = 5 * NX_LN_Q10; cbe = cbe + 1 }
297 nx_layernorm_forward(xt, gamma, beta, yt)
298 // Token 0 (constant input) -> output should be ~5*Q10 (just beta).
299 var ce: nx_int = 0
300 while ce < d {
301 let drift: i64 = py[ce] - 5 * NX_LN_Q10
302 if drift > 200 { return 60 }
303 if drift < -200 { return 61 }
304 ce = ce + 1
305 }
306
307 // --- (f) Verdict gate ---
308 var vi: nx_int = 0
309 while vi < NX_LN_N_VERDICTS {
310 if nx_ln_verdict_is_valid(vi) != 1 { return 70 + vi }
311 vi = vi + 1
312 }
313
314 return 0
315}