1 # BendTorch: layers with typed forward and backward passes, composed generically, trained in Bend 2. 2 # 3 # import ../torch/main.bend as T 4 # 5 # A layer is a record of four functions over its parameters P, its input X, its output Y and what its forward pass 6 # saves for the backward pass, S: 7 # 8 # fwd : P -> X -> Fwd the output, the parameters handed back, the saved values 9 # bwd : P -> S -> Y -> P -> Bwd from dY (typed Y) to dX (typed X), adding dP into a gradient of type P 10 # step : Opt -> P -> P -> OptState

-> Step

one optimizer step on the parameters; the gradient comes back zeroed 11 # zeros : Unit -> P a zero P (gradients and optimizer state start from it) 12 # 13 # The gradient of the input has the input's type and the gradient of the parameters has the parameters' type, so a 14 # backward pass that returns a transposed or a wrongly sized gradient does not compile. Layers are passed as templates 15 # (`~l`), so `Seq` composes any two layers whose types meet, with no closure called twice and no runtime cost. 16 # 17 # Matrices are bend-ml's `Mat` (bend-ml-tensor-array 0.1.14.0): r*c F32 numbers in one flat Array, row by row, 18 # the dimensions erased type parameters. An Array has one owner, so every operation hands back what it reads. 19 # 20 # What is proved and what is tested: the shapes are in the types (checked); the laws at the end of this file are 21 # checked by Bend's kernel; every F32 result is tested against PyTorch by the scripts in reference/, never proved. 22 23 import Base 24 import bend-ml-tensor-array@0.1.18.0/main.bend as TA 25 import bend-ml-nat-lemmas@0.1.2.0/main.bend as NL 26 27 # ===================================================================================================== 28 # The layer protocol 29 # ===================================================================================================== 30 31 # what a forward pass returns: the parameters, the output, what the backward pass will need 32 type Fwd<-P: Type, -Y: Type, -S: Type> is Type: 33 Fwd{p: P, y: Y, s: S} 34 35 # what a backward pass returns: the parameters, the gradient of the input, the accumulated gradient of the parameters 36 type Bwd<-P: Type, -X: Type> is Type: 37 Bwd{p: P, dx: X, g: P} 38 39 # a pair of owned values (parameters of two layers or their saved values) 40 type Two<-A: Type, -B: Type> is Type: 41 Two{a: A, b: B} 42 43 # Type helpers for the left-associated pairs used by Seq/Fork and for Residual's two parameter-free layers. 44 def Pair3(A: Type, B: Type, C: Type) -> Type: 45 Two, C> 46 47 def ResidualType(A: Type) -> Type: 48 Pair3(Unit, A, Unit) 49 50 # AdamW's moments stored as packed BF16 (CAdamWBF16): a tree that follows the parameters' pairs as the step walks them, 51 # with the buffers of each matrix at its leaf. A packed buffer has half the words of its matrix (two BF16 values a 52 # U32, element 2k in the low half of word k, torch/dtype.bend's layout), so it cannot be a value of the parameters' 53 # type P: the tree holds the words itself. A leaf is allocated on the matrix's first step (PNew marks a part not yet 54 # stepped, and stays at the parameter-free layers), so the full-size buffers never exist. v16 picks v packed as well; 55 # otherwise v is F32 (PLeafV). A leaf keeps its matrix's count n, for checkpoints. 56 type PTree is Type: 57 PNew{v16: Bool} 58 PLeaf{n: U32, m: Array, v: Array} 59 PLeafV{n: U32, m: Array, v: Array} 60 PNode{a: PTree, b: PTree} 61 62 # Owned optimizer buffers: SGD has none, momentum one, Adam/AdamW two (Packed: two, in packed BF16, see PTree). 63 type OptState<-P: Type> is Type: 64 NoMom{} 65 OneMom{m: P} 66 TwoMom{m: P, v: P} 67 Packed{t: PTree} 68 69 # Two affine maps are supplied separately: each is called at most once. A packed tree is the same under any relabelling 70 # of the parameters (Lin <-> Two), so it passes unchanged. 71 def OptState.map(-P: Type, -Q: Type, st: OptState

, f: P -> Q, h: P -> Q) -> OptState: 72 match st: 73 case NoMom{}: NoMom{} 74 case OneMom{m}: OneMom{f(m)} 75 case TwoMom{m, v}: TwoMom{f(m), h(v)} 76 case Packed{t}: Packed{t} 77 78 # A packed tree splits at its node; a part not yet stepped splits into two such parts. A leaf belongs to one matrix and 79 # is never split (it goes to the left, the right side starts empty). 80 def OptState.split(-A: Type, -B: Type, st: OptState>) -> Two, OptState>: 81 match st: 82 case NoMom{}: Two{NoMom{}, NoMom{}} 83 case OneMom{Two{a, b}}: Two{OneMom{a}, OneMom{b}} 84 case TwoMom{Two{a, b}, Two{c, d}}: Two{TwoMom{a, c}, TwoMom{b, d}} 85 case Packed{PNode{a, b}}: Two{Packed{a}, Packed{b}} 86 case Packed{PNew{+f}}: Two{Packed{PNew{f}}, Packed{PNew{f}}} 87 case Packed{t}: Two{Packed{t}, Packed{PNew{False{}}}} 88 89 # Pair the buffers present on both sides. Internal callers split one state, then step both children with the same 90 # optimizer, so the variants agree. No buffer is ever allocated to fill an absent side. 91 def OptState.join(-A: Type, -B: Type, a: OptState, b: OptState) -> OptState>: 92 match a b: 93 case OneMom{x} OneMom{y}: OneMom{Two{x, y}} 94 case TwoMom{x, u} TwoMom{y, v}: TwoMom{Two{x, y}, Two{u, v}} 95 case Packed{x} Packed{y}: Packed{PNode{x, y}} 96 case _ _: NoMom{} 97 98 # what an optimizer step returns: the new parameters, the zeroed gradient, the optimizer state 99 type Step<-P: Type> is Type: 100 Step{p: P, g: P, st: OptState

} 101 102 # The optimizer, with the scalars of the current step. Adam's `step` is lr / (1 - b1^t) and `bc2` is 103 # sqrt(1 - b2^t), as PyTorch computes them; `Opt.at` fills them from an `OptCfg` and the step count. 104 # `dec` picks AdamW's decoupled weight decay (w *= 1 - lr*wd) over Adam's (g += wd*w). 105 type Opt is Data: 106 SGD{lr: F32} 107 Momentum{lr: F32, mu: F32} 108 Adam{lr: F32, b1: F32, b2: F32, eps: F32, wd: F32, dec: Bool, step: F32, bc2: F32, c1: F32, c2: F32} 109 110 # Parameters and their flat dump in a fixed order. Loading has a separate array cursor (PA). 111 type PL<-P: Type> is Type: 112 PL{p: P, xs: List<&2, F32>} 113 114 # Loading cursor: the shared source array and next element offset, handed back without a list or a clone. 115 # The IO boundary checks the source length against Layer.count before calling the pure loaders. 116 type PA<-P: Type> is Type: 117 PA{p: P, xs: Array, off: U32} 118 119 type Layer<-P: Type, -X: Type, -Y: Type, -S: Type> is Type: 120 Layer{fwd: P -> X -> Fwd, bwd: P -> S -> Y -> P -> Bwd, step: Opt -> P -> P -> OptState

-> Step

, zeros: Unit -> P, dump: P -> PL

, load: P -> Array -> U32 -> PA

, store: P -> Array -> U32 -> PA

, count: Unit -> U32} 121 122 def Layer.fwd(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer, p: P, x: X) -> Fwd: 123 match l: 124 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: 125 fwd(p, x) 126 127 def Layer.bwd(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer, p: P, s: S, dy: Y, g: P) -> Bwd: 128 match l: 129 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: 130 bwd(p, s, dy, g) 131 132 def Layer.step(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer, +o: Opt, p: P, g: P, st: OptState

) -> Step

: 133 match l: 134 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: 135 step(o, p, g, st) 136 137 def Layer.zeros(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer) -> P: 138 match l: 139 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: 140 zeros(Unit{}) 141 142 # every number of the parameters, in the layer's fixed order (the parameters handed back) 143 def Layer.dump(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer, p: P) -> PL

: 144 match l: 145 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: 146 dump(p) 147 148 # Parameters overwritten from xs at an element offset, in dump order; xs and the next offset come back. 149 def Layer.load(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer, p: P, xs: Array, off: U32) -> PA

: 150 match l: 151 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: load(p, xs, off) 152 153 # The parameters copied into xs from an element offset, in dump order (load's inverse); xs and the next offset come back. 154 def Layer.store(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer, p: P, xs: Array, off: U32) -> PA

: 155 match l: 156 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: store(p, xs, off) 157 158 def Layer.count(-P: Type, -X: Type, -Y: Type, -S: Type, l: Layer) -> U32: 159 match l: 160 case Layer{fwd, bwd, step, zeros, dump, load, store, count}: count(Unit{}) 161 162 # ===================================================================================================== 163 # Element-wise kernels (templates: one copy per function, no closure at run time) 164 # ===================================================================================================== 165 166 # a[t] = f(a[t]) for the n numbers from t; ra holds the array and a[t], already read 167 def ew1_l(~f: @+a: F32 -> F32, nleft: Nat, +t: U32, ra: Array & F32) -> Array: 168 match nleft ra: 169 case 0n Tuple{a, v}: 170 a 171 case 1n+p Tuple{a, +v}: 172 ew1_l(~f, p, (t + 1 : U32), Array.get(F32, Array.set(F32, a, t, f(v)), (t + 1 : U32))) 173 174 # a[t] = f(a[t], b[t]): the result replaces a; b comes back unchanged 175 def ew2_l(~f: @+a: F32 -> @+b: F32 -> F32, nleft: Nat, +t: U32, ra: Array & F32, rb: Array & F32) -> TA.P2: 176 match nleft ra rb: 177 case 0n Tuple{a, v} Tuple{b, w}: 178 TA.P2{a, b} 179 case 1n+p Tuple{a, +v} Tuple{b, +w}: 180 ew2_l(~f, p, (t + 1 : U32), Array.get(F32, Array.set(F32, a, t, f(v, w)), (t + 1 : U32)), Array.get(F32, b, (t + 1 : U32))) 181 182 def Ew.map(~f: @+a: F32 -> F32, +r: Nat, +c: Nat, m: TA.Mat) -> TA.Mat: 183 match m: 184 case TA.Mat{d}: 185 TA.Mat{ew1_l(~f, Nat.mul(r, c), 0, Array.get(F32, d, 0))} 186 187 def ew2_fin(-r: Nat, -c: Nat, p: TA.P2) -> TA.MPair: 188 match p: 189 case TA.P2{a, b}: 190 TA.MPair{TA.Mat{a}, TA.Mat{b}} 191 192 # (f(a, b), b), element by element 193 def Ew.zip(~f: @+a: F32 -> @+b: F32 -> F32, +r: Nat, +c: Nat, a: TA.Mat, b: TA.Mat) -> TA.MPair: 194 match a b: 195 case TA.Mat{x} TA.Mat{y}: 196 ew2_fin(r, c, ew2_l(~f, Nat.mul(r, c), 0, Array.get(F32, x, 0), Array.get(F32, y, 0))) 197 198 def f32_add(+a: F32, +b: F32) -> F32: 199 (a + b : F32) 200 201 def f32_mul(+a: F32, +b: F32) -> F32: 202 (a * b : F32) 203 204 # (a + b, b) 205 def Mat.add(+r: Nat, +c: Nat, a: TA.Mat, b: TA.Mat) -> TA.MPair: 206 Ew.zip(~f32_add, r, c, a, b) 207 208 def zero_of(+x: F32) -> F32: 209 0.0 210 211 # the smallest d with 2^d >= n (TA.cap_depth's answer), by doubling a U32: log2(n) cheap steps, where cap_depth 212 # compares Nats one unit at a time at each doubling; f is fuel (32 doublings cover every U32) 213 def depth_go(f: Nat, ok: Bool, +n: U32, +p: U32, d: Nat) -> Nat: 214 match f ok: 215 case _ True{}: 216 d 217 case 0n False{}: 218 d 219 case 1n+g False{}: 220 depth_go(g, U32.is_le(n, (p + p : U32)), n, (p + p : U32), 1n+d) 221 222 def depth(+n: Nat) -> Nat: 223 depth_go(32n, U32.is_le(U32.from_nat(n), 1), U32.from_nat(n), 1, 0n) 224 225 # ===================================================================================================== 226 # Optimizers: one pass over a parameter, its gradient and the two moment buffers; the gradient is zeroed 227 # ===================================================================================================== 228 229 type Q3 is Type: 230 Q3{w: Array, g: Array, m: Array} 231 232 type Q4 is Type: 233 Q4{w: Array, g: Array, m: Array, v: Array} 234 235 type MatStep<-r: Nat, -c: Nat> is Type: 236 MatStep{w: TA.Mat, g: TA.Mat, st: OptState>} 237 238 # w <- w + (-lr)*g, as PyTorch's p.add_(g, alpha=-lr) 239 def sgd_l(nleft: Nat, +t: U32, +lr: F32, rw: Array & F32, rg: Array & F32) -> TA.P2: 240 match nleft rw rg: 241 case 0n Tuple{w, x} Tuple{g, y}: 242 TA.P2{w, g} 243 case 1n+p Tuple{w, +x} Tuple{g, +y}: 244 sgd_l(p, (t + 1 : U32), lr, Array.get(F32, Array.set(F32, w, t, (x + (F32.neg(lr) * y : F32) : F32)), (t + 1 : U32)), Array.get(F32, Array.set(F32, g, t, 0.0), (t + 1 : U32))) 245 246 # buf <- buf*mu + g; w <- w + (-lr)*buf, as PyTorch's SGD with momentum (no dampening, not Nesterov) 247 def mom_l(nleft: Nat, +t: U32, +lr: F32, +mu: F32, rw: Array & F32, rg: Array & F32, rm: Array & F32) -> Q3: 248 match nleft rw rg rm: 249 case 0n Tuple{w, x} Tuple{g, y} Tuple{m, b}: 250 Q3{w, g, m} 251 case 1n+p Tuple{w, +x} Tuple{g, +y} Tuple{m, +b}: 252 +nb = ((b * mu : F32) + y : F32) 253 +u = (t + 1 : U32) 254 mom_l(p, u, lr, mu, Array.get(F32, Array.set(F32, w, t, (x + (F32.neg(lr) * nb : F32) : F32)), u), Array.get(F32, Array.set(F32, g, t, 0.0), u), Array.get(F32, Array.set(F32, m, t, nb), u)) 255 256 # the Adam family, in PyTorch's order of operations (torch.optim.adam._single_tensor_adam): 257 # g' = g + wd*w (Adam) or w' = w*(1 - lr*wd) (AdamW); m = m + (1-b1)*(g' - m) (lerp); v = v*b2 + (1-b2)*g'*g'; 258 # w = w' + (-step*m) / (sqrt(v)/bc2 + eps) 259 def adam_grad(dec: Bool, +wd: F32, +w: F32, +g: F32) -> F32: 260 match dec: 261 case True{}: 262 g 263 case False{}: 264 (g + (w * wd : F32) : F32) 265 266 def adam_decay(dec: Bool, +lr: F32, +wd: F32, +w: F32) -> F32: 267 match dec: 268 case True{}: 269 (w * (1.0 - (lr * wd : F32) : F32) : F32) 270 case False{}: 271 w 272 273 # c1 = 1 - b1 and c2 = 1 - b2 are given as their own constants (0.1, 0.001), as PyTorch computes 1 - beta in float64 274 # before rounding: 1 - b2 taken in F32 from b2 = 0.999 (0.99900001287 in F32) is 1.3e-5 too small, and over 100 steps 275 # that alone put the tiny Llama 21x further from float64 than PyTorch's float32 (NOTES, 2026-10-18) 276 def adam_m(+c1: F32, +m: F32, +g: F32) -> F32: 277 (m + (c1 * (g - m : F32) : F32) : F32) 278 279 def adam_v(+b2: F32, +c2: F32, +v: F32, +g: F32) -> F32: 280 ((v * b2 : F32) + ((c2 * g : F32) * g : F32) : F32) 281 282 def adam_w(+step: F32, +bc2: F32, +eps: F32, +w: F32, +m: F32, +v: F32) -> F32: 283 (w + ((F32.neg(step) * m : F32) / ((F32.sqrt(v) / bc2 : F32) + eps : F32) : F32) : F32) 284 285 def adam_l(nleft: Nat, +t: U32, +lr: F32, +c1: F32, +b2: F32, +c2: F32, +eps: F32, +wd: F32, +dec: Bool, +step: F32, +bc2: F32, rw: Array & F32, rg: Array & F32, rm: Array & F32, rv: Array & F32) -> Q4: 286 match nleft rw rg rm rv: 287 case 0n Tuple{w, x} Tuple{g, y} Tuple{m, a} Tuple{v, b}: 288 Q4{w, g, m, v} 289 case 1n+p Tuple{w, +x} Tuple{g, +y} Tuple{m, +a} Tuple{v, +b}: 290 +gg = adam_grad(dec, wd, x, y) 291 +nm = adam_m(c1, a, gg) 292 +nv = adam_v(b2, c2, b, gg) 293 +nw = adam_w(step, bc2, eps, adam_decay(dec, lr, wd, x), nm, nv) 294 +u = (t + 1 : U32) 295 adam_l(p, u, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, Array.set(F32, w, t, nw), u), Array.get(F32, Array.set(F32, g, t, 0.0), u), Array.get(F32, Array.set(F32, m, t, nm), u), Array.get(F32, Array.set(F32, v, t, nv), u)) 296 297 # ---- AdamW over packed BF16 moments (CAdamWBF16) ---- 298 # The same F32 arithmetic as adam_l, element by element: each moment is read from its half of a word (BF16 -> F32 is 299 # exact), updated in F32, rounded to BF16 (to nearest, ties to even; a NaN becomes 0xFFFF, as torch/dtype.bend and 300 # PyTorch's CPU kernel do) and written back into its half; the weight update reads the stored, rounded moments, as 301 # PyTorch's AdamW does with BF16 state. The weights and the gradient stay F32. 302 303 def ph_f2u(x: F32) -> U32: 304 match x: 305 case F32{w}: 306 U32{w} 307 308 def ph_u2f(u: U32) -> F32: 309 match u: 310 case U32{w}: 311 F32{w} 312 313 def ph_pick(b: Bool, +x: U32, +y: U32) -> U32: 314 match b: 315 case True{}: 316 x 317 case False{}: 318 y 319 320 # F32 -> BF16 bits (dtype.bend's BF16.of_f32): add 0x7FFF plus the lowest kept bit, keep the top half; NaN -> 0xFFFF 321 def ph_of_bits(+u: U32) -> U32: 322 ph_pick(U32.is_gt((u .&. 2147483647 : U32), 2139095040), 65535, U32.shrn(((u + 32767 : U32) + (U32.shrn(u, 16n) .&. 1 : U32) : U32), 16n)) 323 324 def ph_of(x: F32) -> U32: 325 ph_of_bits(ph_f2u(x)) 326 327 # BF16 bits -> F32, exact 328 def ph_to(+h: U32) -> F32: 329 ph_u2f(U32.shln(h, 16n)) 330 331 # the half of word w that element t uses (odd t: the high half), and w with that half replaced by h 332 def ph_half(odd: Bool, +w: U32) -> U32: 333 match odd: 334 case True{}: 335 U32.shrn(w, 16n) 336 case False{}: 337 (w .&. 65535 : U32) 338 339 def ph_put(odd: Bool, +w: U32, +h: U32) -> U32: 340 match odd: 341 case True{}: 342 ((w .&. 65535 : U32) .|. U32.shln(h, 16n) : U32) 343 case False{}: 344 ((w .&. 4294901760 : U32) .|. h : U32) 345 346 type PQ is Type: 347 PQ{w: Array, g: Array, m: Array, v: Array} 348 349 type PQV is Type: 350 PQV{w: Array, g: Array, m: Array, v: Array} 351 352 # m and v packed. rm and rv hold the word of element t (index t >> 1), read after the previous element's write. 353 def adamh_l(nleft: Nat, +t: U32, +lr: F32, +c1: F32, +b2: F32, +c2: F32, +eps: F32, +wd: F32, +dec: Bool, +step: F32, +bc2: F32, rw: Array & F32, rg: Array & F32, rm: Array & U32, rv: Array & U32) -> PQ: 354 match nleft rw rg rm rv: 355 case 0n Tuple{w, x} Tuple{g, y} Tuple{m, a} Tuple{v, b}: 356 PQ{w, g, m, v} 357 case 1n+p Tuple{w, +x} Tuple{g, +y} Tuple{m, +a} Tuple{v, +b}: 358 +odd = U32.is_eq((t .&. 1 : U32), 1) 359 +gg = adam_grad(dec, wd, x, y) 360 +hm = ph_of(adam_m(c1, ph_to(ph_half(odd, a)), gg)) 361 +hv = ph_of(adam_v(b2, c2, ph_to(ph_half(odd, b)), gg)) 362 +nw = adam_w(step, bc2, eps, adam_decay(dec, lr, wd, x), ph_to(hm), ph_to(hv)) 363 +u = (t + 1 : U32) 364 +k = U32.shrn(t, 1n) 365 +ku = U32.shrn(u, 1n) 366 adamh_l(p, u, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, Array.set(F32, w, t, nw), u), Array.get(F32, Array.set(F32, g, t, 0.0), u), Array.get(U32, Array.set(U32, m, k, ph_put(odd, a, hm)), ku), Array.get(U32, Array.set(U32, v, k, ph_put(odd, b, hv)), ku)) 367 368 # m packed, v in F32 (one cell an element, as adam_l's) 369 def adamhv_l(nleft: Nat, +t: U32, +lr: F32, +c1: F32, +b2: F32, +c2: F32, +eps: F32, +wd: F32, +dec: Bool, +step: F32, +bc2: F32, rw: Array & F32, rg: Array & F32, rm: Array & U32, rv: Array & F32) -> PQV: 370 match nleft rw rg rm rv: 371 case 0n Tuple{w, x} Tuple{g, y} Tuple{m, a} Tuple{v, b}: 372 PQV{w, g, m, v} 373 case 1n+p Tuple{w, +x} Tuple{g, +y} Tuple{m, +a} Tuple{v, +b}: 374 +odd = U32.is_eq((t .&. 1 : U32), 1) 375 +gg = adam_grad(dec, wd, x, y) 376 +hm = ph_of(adam_m(c1, ph_to(ph_half(odd, a)), gg)) 377 +nv = adam_v(b2, c2, b, gg) 378 +nw = adam_w(step, bc2, eps, adam_decay(dec, lr, wd, x), ph_to(hm), nv) 379 +u = (t + 1 : U32) 380 adamhv_l(p, u, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, Array.set(F32, w, t, nw), u), Array.get(F32, Array.set(F32, g, t, 0.0), u), Array.get(U32, Array.set(U32, m, U32.shrn(t, 1n), ph_put(odd, a, hm)), U32.shrn(u, 1n)), Array.get(F32, Array.set(F32, v, t, nv), u)) 381 382 # the words of n packed values: ceil(n / 2), in a block of the next power of two 383 def ph_words(+n: Nat) -> Array: 384 Array.new(U32, depth(Nat.div(Nat.add(n, 1n), 2n)), 0) 385 386 def ph_fin(-r: Nat, -c: Nat, +n: U32, q: PQ) -> MatStep: 387 match q: 388 case PQ{w, g, m, v}: MatStep{TA.Mat{w}, TA.Mat{g}, Packed{PLeaf{n, m, v}}} 389 390 def phv_fin(-r: Nat, -c: Nat, +n: U32, q: PQV) -> MatStep: 391 match q: 392 case PQV{w, g, m, v}: MatStep{TA.Mat{w}, TA.Mat{g}, Packed{PLeafV{n, m, v}}} 393 394 # one matrix's step on its leaf; anything else leaves the values untouched, as Opt.mat's fallback 395 def adamh_leaf(+r: Nat, +c: Nat, +o: Opt, wa: Array, ga: Array, t: PTree) -> MatStep: 396 match o t: 397 case Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} PLeaf{+n, ma, va}: 398 ph_fin(r, c, n, adamh_l(Nat.mul(r, c), 0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, wa, 0), Array.get(F32, ga, 0), Array.get(U32, ma, 0), Array.get(U32, va, 0))) 399 case Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} PLeafV{+n, ma, va}: 400 phv_fin(r, c, n, adamhv_l(Nat.mul(r, c), 0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, wa, 0), Array.get(F32, ga, 0), Array.get(U32, ma, 0), Array.get(F32, va, 0))) 401 case _ t: 402 MatStep{TA.Mat{wa}, TA.Mat{ga}, Packed{t}} 403 404 # the leaf of a matrix on its first step: zeros (BF16 0x0000 is +0.0) 405 def ph_new(+n: Nat, v16: Bool) -> PTree: 406 match v16: 407 case True{}: 408 PLeaf{U32.from_nat(n), ph_words(n), ph_words(n)} 409 case False{}: 410 PLeafV{U32.from_nat(n), ph_words(n), Array.new(F32, depth(n), 0.0)} 411 412 def adamh_mat(+r: Nat, +c: Nat, +o: Opt, wa: Array, ga: Array, t: PTree) -> MatStep: 413 match t: 414 case PNew{v16}: 415 adamh_leaf(r, c, o, wa, ga, ph_new(Nat.mul(r, c), v16)) 416 case t: 417 adamh_leaf(r, c, o, wa, ga, t) 418 419 # One matrix step, with only the buffers this optimizer reads. 420 def Opt.sgd_fin(-r: Nat, -c: Nat, q: TA.P2) -> MatStep: 421 match q: 422 case TA.P2{w, g}: MatStep{TA.Mat{w}, TA.Mat{g}, NoMom{}} 423 424 def Opt.mom_fin(-r: Nat, -c: Nat, q: Q3) -> MatStep: 425 match q: 426 case Q3{w, g, m}: MatStep{TA.Mat{w}, TA.Mat{g}, OneMom{TA.Mat{m}}} 427 428 def Opt.fin(-r: Nat, -c: Nat, q: Q4) -> MatStep: 429 match q: 430 case Q4{w, g, m, v}: MatStep{TA.Mat{w}, TA.Mat{g}, TwoMom{TA.Mat{m}, TA.Mat{v}}} 431 432 # A mismatched state is initialized only at the training boundary; Opt.mat requires the matching variant. 433 # The fallback leaves the values untouched, so it cannot fabricate a moment buffer. 434 def Opt.mat(+r: Nat, +c: Nat, +o: Opt, w: TA.Mat, g: TA.Mat, st: OptState>) -> MatStep: 435 match o w g st: 436 case SGD{lr} TA.Mat{wd} TA.Mat{gd} _: 437 Opt.sgd_fin(r, c, sgd_l(Nat.mul(r, c), 0, lr, Array.get(F32, wd, 0), Array.get(F32, gd, 0))) 438 case Momentum{lr, mu} TA.Mat{wd} TA.Mat{gd} OneMom{TA.Mat{md}}: 439 Opt.mom_fin(r, c, mom_l(Nat.mul(r, c), 0, lr, mu, Array.get(F32, wd, 0), Array.get(F32, gd, 0), Array.get(F32, md, 0))) 440 case Adam{lr, b1, b2, eps, wd2, dec, step, bc2, c1, c2} TA.Mat{wd} TA.Mat{gd} TwoMom{TA.Mat{md}, TA.Mat{vd}}: 441 Opt.fin(r, c, adam_l(Nat.mul(r, c), 0, lr, c1, b2, c2, eps, wd2, dec, step, bc2, Array.get(F32, wd, 0), Array.get(F32, gd, 0), Array.get(F32, md, 0), Array.get(F32, vd, 0))) 442 case Adam{lr, b1, b2, eps, wd2, dec, step, bc2, c1, c2} TA.Mat{wd} TA.Mat{gd} Packed{t}: 443 adamh_mat(r, c, o, wd, gd, t) 444 case _ w g st: 445 MatStep{w, g, st} 446 447 # What the user picks; `Opt.at` turns it into the step's `Opt` from q1 = 1 - b1^t and q2 = 1 - b2^t (t = 1 on the 448 # first step). Adam's c1 and c2 are 1 - b1 and 1 - b2, written as their own constants (adam_m above). CAdamWBF16 is 449 # AdamW with its first moment (and its second, when v16) stored as packed BF16 (PTree); weights and gradient stay F32. 450 type OptCfg is Data: 451 CSGD{lr: F32} 452 CMomentum{lr: F32, mu: F32} 453 CAdam{lr: F32, b1: F32, b2: F32, eps: F32, wd: F32, c1: F32, c2: F32} 454 CAdamW{lr: F32, b1: F32, b2: F32, eps: F32, wd: F32, c1: F32, c2: F32} 455 CAdamWBF16{lr: F32, b1: F32, b2: F32, eps: F32, wd: F32, c1: F32, c2: F32, v16: Bool} 456 457 def Opt.at(+cfg: OptCfg, +q1: F32, +q2: F32) -> Opt: 458 match cfg: 459 case CSGD{lr}: 460 SGD{lr} 461 case CMomentum{lr, mu}: 462 Momentum{lr, mu} 463 case CAdam{lr, b1, b2, eps, wd, c1, c2}: 464 Adam{lr, b1, b2, eps, wd, False{}, (lr / q1 : F32), F32.sqrt(q2), c1, c2} 465 case CAdamW{lr, b1, b2, eps, wd, c1, c2}: 466 Adam{lr, b1, b2, eps, wd, True{}, (lr / q1 : F32), F32.sqrt(q2), c1, c2} 467 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}: 468 Adam{lr, b1, b2, eps, wd, True{}, (lr / q1 : F32), F32.sqrt(q2), c1, c2} 469 470 def OptCfg.b1(+cfg: OptCfg) -> F32: 471 match cfg: 472 case CAdam{lr, b1, b2, eps, wd, c1, c2}: 473 b1 474 case CAdamW{lr, b1, b2, eps, wd, c1, c2}: 475 b1 476 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}: 477 b1 478 case _: 479 0.0 480 481 def OptCfg.b2(+cfg: OptCfg) -> F32: 482 match cfg: 483 case CAdam{lr, b1, b2, eps, wd, c1, c2}: 484 b2 485 case CAdamW{lr, b1, b2, eps, wd, c1, c2}: 486 b2 487 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}: 488 b2 489 case _: 490 0.0 491 492 def OptCfg.c1(+cfg: OptCfg) -> F32: 493 match cfg: 494 case CAdam{lr, b1, b2, eps, wd, c1, c2}: 495 c1 496 case CAdamW{lr, b1, b2, eps, wd, c1, c2}: 497 c1 498 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}: 499 c1 500 case _: 501 0.0 502 503 def OptCfg.c2(+cfg: OptCfg) -> F32: 504 match cfg: 505 case CAdam{lr, b1, b2, eps, wd, c1, c2}: 506 c2 507 case CAdamW{lr, b1, b2, eps, wd, c1, c2}: 508 c2 509 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}: 510 c2 511 case _: 512 0.0 513 514 # c = 1 - b up to F32's rounding, with 0 < c <= 1 (PyTorch's 0 <= beta < 1). For b in [0, 1), 1 - b taken in F32 is 515 # within half an ulp of the exact value and c's own rounding adds at most as much, so a c written as 1 - b is always 516 # under 2^-23 from it, while a typo (0.01 for 0.001, c1 and c2 swapped) is far above. A NaN fails every comparison. 517 def beta_ok(+b: F32, +c: F32) -> Bool: 518 +d = ((1.0 - b : F32) - c : F32) 519 Bool.and(Bool.and((c > 0.0 : F32), (c <= 1.0 : F32)), ((d * d : F32) <= 0.0000000000000142108547152020037174224853515625 : F32)) 520 521 # True unless an Adam configuration's c1, c2 disagree with its b1, b2. Train.new poisons the state of a configuration 522 # that fails (every loss after the first is NaN); a program checks it first with OptCfg.check to stop with a message. 523 def OptCfg.ok(+cfg: OptCfg) -> Bool: 524 match cfg: 525 case CAdam{lr, b1, b2, eps, wd, c1, c2}: 526 Bool.and(beta_ok(b1, c1), beta_ok(b2, c2)) 527 case CAdamW{lr, b1, b2, eps, wd, c1, c2}: 528 Bool.and(beta_ok(b1, c1), beta_ok(b2, c2)) 529 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}: 530 Bool.and(beta_ok(b1, c1), beta_ok(b2, c2)) 531 case _: 532 True{} 533 534 def cfg_stop(ok: Bool) -> IO(Unit): 535 match ok: 536 case True{}: 537 IO.pure(Unit, Unit{}) 538 case False{}: 539 IO.die(Unit, 1, "CAdam/CAdamW: c1, c2 do not match b1, b2 (c1 = 1 - b1, c2 = 1 - b2, e.g. 0.9, 0.999 with 0.1, 0.001)") 540 541 def OptCfg.check(+cfg: OptCfg) -> IO(Unit): 542 cfg_stop(OptCfg.ok(cfg)) 543 544 # dump and load for a matrix: its r*c numbers, row by row 545 def mat_dump_fin(-r: Nat, -c: Nat, ml: TA.ML) -> PL>: 546 match ml: 547 case TA.ML{m, xs}: 548 PL{m, xs} 549 550 def Mat.dump(+r: Nat, +c: Nat, m: TA.Mat) -> PL>: 551 mat_dump_fin(r, c, TA.Mat.to_list(r, c, m)) 552 553 def load_cells(left: Nat, +i: U32, +off: U32, rx: Array & F32, dst: Array) -> TA.P2: 554 match left rx: 555 case 0n Tuple{xs, v}: TA.P2{xs, dst} 556 case 1n+q Tuple{xs, v}: 557 load_cells(q, (i + 1 : U32), (off + 1 : U32), Array.get(F32, xs, (off + 1 : U32)), Array.set(F32, dst, i, v)) 558 559 def mat_restore_fin(-r: Nat, -c: Nat, +off: U32, q: TA.P2) -> PA>: 560 match q: 561 case TA.P2{xs, dst}: PA{TA.Mat{dst}, xs, off} 562 563 def Mat.restore(+r: Nat, +c: Nat, m: TA.Mat, xs: Array, +off: U32) -> PA>: 564 match m: 565 case TA.Mat{dst}: 566 mat_restore_fin(r, c, (off + U32.from_nat(Nat.mul(r, c)) : U32), load_cells(Nat.mul(r, c), 0, off, Array.get(F32, xs, off), dst)) 567 568 # load_cells' inverse: copies the matrix's cells i.. into xs from `off`; both arrays come back (the matrix's first) 569 def store_cells(left: Nat, +i: U32, +off: U32, rm: Array & F32, xs: Array) -> TA.P2: 570 match left rm: 571 case 0n Tuple{ma, v}: TA.P2{ma, xs} 572 case 1n+q Tuple{ma, v}: 573 store_cells(q, (i + 1 : U32), (off + 1 : U32), Array.get(F32, ma, (i + 1 : U32)), Array.set(F32, xs, off, v)) 574 575 def mat_store_fin(-r: Nat, -c: Nat, +off: U32, q: TA.P2) -> PA>: 576 match q: 577 case TA.P2{ma, xs}: PA{TA.Mat{ma}, xs, off} 578 579 def Mat.store(+r: Nat, +c: Nat, m: TA.Mat, xs: Array, +off: U32) -> PA>: 580 match m: 581 case TA.Mat{ma}: 582 mat_store_fin(r, c, (off + U32.from_nat(Nat.mul(r, c)) : U32), store_cells(Nat.mul(r, c), 0, off, Array.get(F32, ma, 0), xs)) 583 584 def pa_second(-A: Type, -B: Type, a: A, rb: PA) -> PA>: 585 match rb: 586 case PA{b, xs, off}: PA{Two{a, b}, xs, off} 587 588 def pl_pair(-A: Type, -B: Type, ra: PL, rb: PL) -> PL>: 589 match ra rb: 590 case PL{a, xa} PL{b, xb}: 591 PL{Two{a, b}, List.append(&2, F32, xa, xb)} 592 593 def pl_second(-A: Type, -B: Type, a: A, rb: PL) -> PL>: 594 match rb: 595 case PL{b, rest}: 596 PL{Two{a, b}, rest} 597 598 # ===================================================================================================== 599 # Linear: y = x·Wᵀ + b, with W stored out x in (PyTorch's nn.Linear layout); saves its input 600 # ===================================================================================================== 601 602 type Lin<-i: Nat, -o: Nat> is Type: 603 Lin{w: TA.Mat, b: TA.Mat<1n, o>} 604 605 def lin_f3(-n: Nat, -i: Nat, -o: Nat, w: TA.Mat, x: TA.Mat, r: TA.MBias) -> Fwd, TA.Mat, TA.Mat>: 606 match r: 607 case TA.MBias{y, b}: 608 Fwd{Lin{w, b}, y, x} 609 610 def lin_f2(+n: Nat, -i: Nat, +o: Nat, b: TA.Mat<1n, o>, r: TA.MMulNT) -> Fwd, TA.Mat, TA.Mat>: 611 match r: 612 case TA.MMulNT{x, w, z}: 613 lin_f3(n, i, o, w, x, TA.Mat.add_row(n, o, z, b)) 614 615 def lin_fwd(+n: Nat, +i: Nat, +o: Nat, par: Nat, p: Lin, x: TA.Mat) -> Fwd, TA.Mat, TA.Mat>: 616 match p: 617 case Lin{w, b}: 618 lin_f2(n, i, o, b, TA.Mat.matmul_nt(n, i, o, par, x, w)) 619 620 def lin_b4(-n: Nat, -i: Nat, -o: Nat, w: TA.Mat, b: TA.Mat<1n, o>, dx: TA.Mat, gw: TA.Mat, r: TA.MColSum) -> Bwd, TA.Mat>: 621 match r: 622 case TA.MColSum{dy, gb}: 623 Bwd{Lin{w, b}, dx, Lin{gw, gb}} 624 625 def lin_b3(+n: Nat, -i: Nat, +o: Nat, w: TA.Mat, b: TA.Mat<1n, o>, dx: TA.Mat, gb: TA.Mat<1n, o>, r: TA.MMulTN) -> Bwd, TA.Mat>: 626 match r: 627 case TA.MMulTN{dy, x, gw}: 628 lin_b4(n, i, o, w, b, dx, gw, TA.Mat.col_sums_acc(n, o, dy, gb)) 629 630 def lin_b2(+n: Nat, +i: Nat, +o: Nat, par: Nat, b: TA.Mat<1n, o>, x: TA.Mat, g: Lin, r: TA.MMul) -> Bwd, TA.Mat>: 631 match g r: 632 case Lin{gw, gb} TA.MMul{dy, w, dx}: 633 lin_b3(n, i, o, w, b, dx, gb, TA.Mat.matmul_tn_acc(o, n, i, par, dy, x, gw)) 634 635 # dX = dY·W, dW += dYᵀ·X, db += the column sums of dY; the two sums are added straight into the gradient (the 636 # accumulating products of bend-ml-tensor-array 0.1.10.0), with the bits of computing them apart and adding them 637 def lin_bwd(+n: Nat, +i: Nat, +o: Nat, +par: Nat, p: Lin, x: TA.Mat, dy: TA.Mat, g: Lin) -> Bwd, TA.Mat>: 638 match p: 639 case Lin{w, b}: 640 lin_b2(n, i, o, par, b, x, g, TA.Mat.matmul_precise(n, o, i, par, dy, w)) 641 642 def pair_s3(-P1: Type, -P2: Type, r1: Step, r2: Step) -> Step>: 643 match r1 r2: 644 case Step{p1, g1, s1} Step{p2, g2, s2}: 645 Step{Two{p1, p2}, Two{g1, g2}, OptState.join(P1, P2, s1, s2)} 646 647 def lin_pair(-i: Nat, -o: Nat, p: Lin) -> Two, TA.Mat<1n, o>>: 648 match p: 649 case Lin{w, b}: Two{w, b} 650 651 def lin_unpair(-i: Nat, -o: Nat, p: Two, TA.Mat<1n, o>>) -> Lin: 652 match p: 653 case Two{w, b}: Lin{w, b} 654 655 def lin_s2(-i: Nat, -o: Nat, a: MatStep, c: MatStep<1n, o>) -> Step>: 656 match a c: 657 case MatStep{w, gw, sw} MatStep{b, gb, sb}: 658 Step{Lin{w, b}, Lin{gw, gb}, OptState.map(Two, TA.Mat<1n, o>>, Lin, OptState.join(TA.Mat, TA.Mat<1n, o>, sw, sb), x => lin_unpair(i, o, x), x => lin_unpair(i, o, x))} 659 660 def lin_step2(+i: Nat, +o: Nat, +op: Opt, p: Lin, g: Lin, st: Two>, OptState>>) -> Step>: 661 match p g st: 662 case Lin{w, b} Lin{gw, gb} Two{sw, sb}: 663 lin_s2(i, o, Opt.mat(o, i, op, w, gw, sw), Opt.mat(1n, o, op, b, gb, sb)) 664 665 def lin_step(+i: Nat, +o: Nat, +op: Opt, p: Lin, g: Lin, st: OptState>) -> Step>: 666 lin_step2(i, o, op, p, g, OptState.split(TA.Mat, TA.Mat<1n, o>, OptState.map(Lin, Two, TA.Mat<1n, o>>, st, x => lin_pair(i, o, x), x => lin_pair(i, o, x)))) 667 668 def lin_dump2(-i: Nat, -o: Nat, r: PL, TA.Mat<1n, o>>>) -> PL>: 669 match r: 670 case PL{Two{w, b}, xs}: 671 PL{Lin{w, b}, xs} 672 673 def lin_dump(+i: Nat, +o: Nat, p: Lin) -> PL>: 674 match p: 675 case Lin{w, b}: 676 lin_dump2(i, o, pl_pair(TA.Mat, TA.Mat<1n, o>, Mat.dump(o, i, w), Mat.dump(1n, o, b))) 677 678 def lin_load3(-i: Nat, -o: Nat, r: PA, TA.Mat<1n, o>>>) -> PA>: 679 match r: 680 case PA{Two{w, b}, xs, off}: PA{Lin{w, b}, xs, off} 681 682 def lin_load2(+i: Nat, +o: Nat, b: TA.Mat<1n, o>, r: PA>) -> PA>: 683 match r: 684 case PA{w, xs, off}: 685 lin_load3(i, o, pa_second(TA.Mat, TA.Mat<1n, o>, w, Mat.restore(1n, o, b, xs, off))) 686 687 def lin_load(+i: Nat, +o: Nat, p: Lin, xs: Array, off: U32) -> PA>: 688 match p: 689 case Lin{w, b}: lin_load2(i, o, b, Mat.restore(o, i, w, xs, off)) 690 691 def lin_store2(+i: Nat, +o: Nat, b: TA.Mat<1n, o>, r: PA>) -> PA>: 692 match r: 693 case PA{w, xs, off}: 694 lin_load3(i, o, pa_second(TA.Mat, TA.Mat<1n, o>, w, Mat.store(1n, o, b, xs, off))) 695 696 def lin_store(+i: Nat, +o: Nat, p: Lin, xs: Array, off: U32) -> PA>: 697 match p: 698 case Lin{w, b}: lin_store2(i, o, b, Mat.store(o, i, w, xs, off)) 699 700 def lin_zeros(+i: Nat, +o: Nat) -> Lin: 701 Lin{TA.Mat.zeros(o, i), TA.Mat.zeros(1n, o)} 702 703 # a linear layer from i to o features on batches of n rows; its products run in 2^par blocks of rows 704 def Linear(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> Layer, TA.Mat, TA.Mat, TA.Mat>: 705 Layer{p => x => lin_fwd(n, i, o, par, p, x), p => s => dy => g => lin_bwd(n, i, o, par, p, s, dy, g), op => p => g => st => lin_step(i, o, op, p, g, st), u => lin_zeros(i, o), p => lin_dump(i, o, p), p => xs => off => lin_load(i, o, p, xs, off), p => xs => off => lin_store(i, o, p, xs, off), u => U32.from_nat(Nat.mul(o, 1n+i))} 706 707 # the first layer of a model, whose input (the data) needs no gradient, as an input without requires_grad in PyTorch: 708 # the backward pass computes dW and db only, skipping dX = dY·W (a third of a linear layer's work), and returns zeros 709 # as dX. Use it only where nothing before the layer has parameters. 710 def linf_b3(+n: Nat, +i: Nat, +o: Nat, w: TA.Mat, b: TA.Mat<1n, o>, gb: TA.Mat<1n, o>, r: TA.MMulTN) -> Bwd, TA.Mat>: 711 match r: 712 case TA.MMulTN{dy, x, gw}: 713 lin_b4(n, i, o, w, b, TA.Mat.zeros(n, i), gw, TA.Mat.col_sums_acc(n, o, dy, gb)) 714 715 def linf_bwd(+n: Nat, +i: Nat, +o: Nat, +par: Nat, p: Lin, x: TA.Mat, dy: TA.Mat, g: Lin) -> Bwd, TA.Mat>: 716 match p g: 717 case Lin{w, b} Lin{gw, gb}: 718 linf_b3(n, i, o, w, b, gb, TA.Mat.matmul_tn_acc(o, n, i, par, dy, x, gw)) 719 720 def LinearFirst(+n: Nat, +i: Nat, +o: Nat, +par: Nat) -> Layer, TA.Mat, TA.Mat, TA.Mat>: 721 Layer{p => x => lin_fwd(n, i, o, par, p, x), p => s => dy => g => linf_bwd(n, i, o, par, p, s, dy, g), op => p => g => st => lin_step(i, o, op, p, g, st), u => lin_zeros(i, o), p => lin_dump(i, o, p), p => xs => off => lin_load(i, o, p, xs, off), p => xs => off => lin_store(i, o, p, xs, off), u => U32.from_nat(Nat.mul(o, 1n+i))} 722 723 # ===================================================================================================== 724 # Layers without parameters 725 # ===================================================================================================== 726 727 def pick(b: Bool, +x: F32, +y: F32) -> F32: 728 match b: 729 case True{}: 730 x 731 case False{}: 732 y 733 734 def no_step(-P: Type, +op: Opt, p: P, g: P, st: OptState

) -> Step

: 735 Step{p, g, st} 736 737 # an element-wise activation that saves a copy of its output (relu, sigmoid, tanh: their derivative is a function of 738 # the output). `f` is the activation, `d(dy, y)` the gradient of its input. 739 def act_f2(-n: Nat, -c: Nat, r: TA.MM2) -> Fwd, TA.Mat>: 740 match r: 741 case TA.MM2{y, s}: 742 Fwd{Unit{}, y, s} 743 744 def act_b2(-n: Nat, -c: Nat, r: TA.MPair) -> Bwd>: 745 match r: 746 case TA.MPair{dx, y}: 747 Bwd{Unit{}, dx, Unit{}} 748 749 def ActOut(~f: @+a: F32 -> F32, ~d: @+a: F32 -> @+b: F32 -> F32, +n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 750 Layer{p => x => act_f2(n, c, TA.Mat.clone(n, c, Ew.map(~f, n, c, x))), p => s => dy => g => act_b2(n, c, Ew.zip(~d, n, c, dy, s)), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0} 751 752 # relu(x) = x unless x < 0 (so relu(NaN) = NaN and relu(-0) = -0); its gradient is +0 where the output is <= 0 and 753 # passes everywhere else, NaN included (PyTorch's threshold_backward) 754 def relu_v(+x: F32) -> F32: 755 pick(F32.is_lt(x, 0.0), 0.0, x) 756 757 def relu_d(+dy: F32, +y: F32) -> F32: 758 pick(F32.is_le(y, 0.0), 0.0, dy) 759 760 def ReLU(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 761 ActOut(~relu_v, ~relu_d, n, c) 762 763 # ===================================================================================================== 764 # Seq: la then lb. Parameters, gradients, saved values and optimizer state are pairs of the two layers'. 765 # ===================================================================================================== 766 767 def seq_f3(-P1: Type, -P2: Type, -Z: Type, -S1: Type, -S2: Type, p1: P1, s1: S1, r: Fwd) -> Fwd, Z, Two>: 768 match r: 769 case Fwd{p2, z, s2}: 770 Fwd{Two{p1, p2}, z, Two{s1, s2}} 771 772 def seq_f2(~P2: Type, ~Y: Type, ~Z: Type, ~S2: Type, ~lb: Layer, -P1: Type, -S1: Type, p2: P2, r: Fwd) -> Fwd, Z, Two>: 773 match r: 774 case Fwd{p1, y, s1}: 775 seq_f3(P1, P2, Z, S1, S2, p1, s1, Layer.fwd(P2, Y, Z, S2, lb, p2, y)) 776 777 def seq_fwd(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p: Two, x: X) -> Fwd, Z, Two>: 778 match p: 779 case Two{p1, p2}: 780 seq_f2(~P2, ~Y, ~Z, ~S2, ~lb, P1, S1, p2, Layer.fwd(P1, X, Y, S1, la, p1, x)) 781 782 def seq_b3(-P1: Type, -P2: Type, -X: Type, p2: P2, g2: P2, r: Bwd) -> Bwd, X>: 783 match r: 784 case Bwd{p1, dx, g1}: 785 Bwd{Two{p1, p2}, dx, Two{g1, g2}} 786 787 def seq_b2(~P1: Type, ~X: Type, ~Y: Type, ~S1: Type, ~la: Layer, -P2: Type, p1: P1, s1: S1, g1: P1, r: Bwd) -> Bwd, X>: 788 match r: 789 case Bwd{p2, dy, g2}: 790 seq_b3(P1, P2, X, p2, g2, Layer.bwd(P1, X, Y, S1, la, p1, s1, dy, g1)) 791 792 def seq_bwd(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p: Two, s: Two, dz: Z, g: Two) -> Bwd, X>: 793 match p s g: 794 case Two{p1, p2} Two{s1, s2} Two{g1, g2}: 795 seq_b2(~P1, ~X, ~Y, ~S1, ~la, P2, p1, s1, g1, Layer.bwd(P2, Y, Z, S2, lb, p2, s2, dz, g2)) 796 797 def seq_s3(-P1: Type, -P2: Type, r1: Step, r2: Step) -> Step>: 798 pair_s3(P1, P2, r1, r2) 799 800 def seq_s2(~P2: Type, ~Y: Type, ~Z: Type, ~S2: Type, ~lb: Layer, -P1: Type, +op: Opt, p2: P2, g2: P2, s2: OptState, r1: Step) -> Step>: 801 seq_s3(P1, P2, r1, Layer.step(P2, Y, Z, S2, lb, op, p2, g2, s2)) 802 803 def seq_step2(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer, +op: Opt, p: Two, g: Two, st: Two, OptState>) -> Step>: 804 match p g st: 805 case Two{p1, p2} Two{g1, g2} Two{s1, s2}: 806 seq_s2(~P2, ~Y, ~Z, ~S2, ~lb, P1, op, p2, g2, s2, Layer.step(P1, X, Y, S1, la, op, p1, g1, s1)) 807 808 def seq_step(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer, +op: Opt, p: Two, g: Two, st: OptState>) -> Step>: 809 seq_step2(~P1, ~P2, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, op, p, g, OptState.split(P1, P2, st)) 810 811 # dump and load for a pair of layers: the first layer's numbers, then the second's 812 # a pair's zeros, the two halves in a parallel let: a model's gradient and optimizer state are many arrays, and making 813 # them one after another left the training state's zeroing (page faults, mostly) on one thread 814 def pair_zeros(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer, ~lb: Layer) -> Two: 815 a b = Layer.zeros(P1, X1, Y1, S1, la) Layer.zeros(P2, X2, Y2, S2, lb) 816 Two{a, b} 817 818 def pair_dump(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p: Two) -> PL>: 819 match p: 820 case Two{p1, p2}: 821 pl_pair(P1, P2, Layer.dump(P1, X1, Y1, S1, la, p1), Layer.dump(P2, X2, Y2, S2, lb, p2)) 822 823 def pair_load2(~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~lb: Layer, -P1: Type, p2: P2, r: PA) -> PA>: 824 match r: 825 case PA{p1, xs, off}: 826 pa_second(P1, P2, p1, Layer.load(P2, X2, Y2, S2, lb, p2, xs, off)) 827 828 def pair_load(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p: Two, xs: Array, off: U32) -> PA>: 829 match p: 830 case Two{p1, p2}: 831 pair_load2(~P2, ~X2, ~Y2, ~S2, ~lb, P1, p2, Layer.load(P1, X1, Y1, S1, la, p1, xs, off)) 832 833 def pair_store2(~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~lb: Layer, -P1: Type, p2: P2, r: PA) -> PA>: 834 match r: 835 case PA{p1, xs, off}: 836 pa_second(P1, P2, p1, Layer.store(P2, X2, Y2, S2, lb, p2, xs, off)) 837 838 def pair_store(~P1: Type, ~X1: Type, ~Y1: Type, ~S1: Type, ~P2: Type, ~X2: Type, ~Y2: Type, ~S2: Type, ~la: Layer, ~lb: Layer, p: Two, xs: Array, off: U32) -> PA>: 839 match p: 840 case Two{p1, p2}: 841 pair_store2(~P2, ~X2, ~Y2, ~S2, ~lb, P1, p2, Layer.store(P1, X1, Y1, S1, la, p1, xs, off)) 842 843 def Seq(~P1: Type, ~P2: Type, ~X: Type, ~Y: Type, ~Z: Type, ~S1: Type, ~S2: Type, ~la: Layer, ~lb: Layer) -> Layer, X, Z, Two>: 844 # bend-lint: disable-next bend-lint-rules/performance/gpu_fit -- zeros are made once per run (Train.new, OptState.new), and pair_zeros already forks not-expected-because: the combinators nest, so a bang here would run inside a bang 845 Layer{p => x => seq_fwd(~P1, ~P2, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, p, x), p => s => dz => g => seq_bwd(~P1, ~P2, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, p, s, dz, g), op => p => g => st => seq_step(~P1, ~P2, ~X, ~Y, ~Z, ~S1, ~S2, ~la, ~lb, op, p, g, st), u => pair_zeros(~P1, ~X, ~Y, ~S1, ~P2, ~Y, ~Z, ~S2, ~la, ~lb), p => pair_dump(~P1, ~X, ~Y, ~S1, ~P2, ~Y, ~Z, ~S2, ~la, ~lb, p), p => xs => off => pair_load(~P1, ~X, ~Y, ~S1, ~P2, ~Y, ~Z, ~S2, ~la, ~lb, p, xs, off), p => xs => off => pair_store(~P1, ~X, ~Y, ~S1, ~P2, ~Y, ~Z, ~S2, ~la, ~lb, p, xs, off), u => (Layer.count(P1, X, Y, S1, la) + Layer.count(P2, Y, Z, S2, lb) : U32)} 846 847 # Owned layer composition for dimensions read at run time. Each child record is opened once; each closure goes to 848 # exactly one corresponding closure of the composite layer. The template-based Seq above keeps its existing laws. 849 def seq_at_f2(-P1: Type, -P2: Type, -Y: Type, -Z: Type, -S1: Type, -S2: Type, p2: P2, fb: P2 -> Y -> Fwd, r: Fwd) -> Fwd, Z, Two>: 850 match r: 851 case Fwd{p1, y, s1}: seq_f3(P1, P2, Z, S1, S2, p1, s1, fb(p2, y)) 852 853 def seq_at_f(-P1: Type, -P2: Type, -X: Type, -Y: Type, -Z: Type, -S1: Type, -S2: Type, fa: P1 -> X -> Fwd, fb: P2 -> Y -> Fwd, p: Two, x: X) -> Fwd, Z, Two>: 854 match p: 855 case Two{p1, p2}: seq_at_f2(P1, P2, Y, Z, S1, S2, p2, fb, fa(p1, x)) 856 857 def seq_at_b2(-P1: Type, -P2: Type, -X: Type, -Y: Type, -S1: Type, ba: P1 -> S1 -> Y -> P1 -> Bwd, p1: P1, s1: S1, g1: P1, r: Bwd) -> Bwd, X>: 858 match r: 859 case Bwd{p2, dy, g2}: seq_b3(P1, P2, X, p2, g2, ba(p1, s1, dy, g1)) 860 861 def seq_at_b(-P1: Type, -P2: Type, -X: Type, -Y: Type, -Z: Type, -S1: Type, -S2: Type, ba: P1 -> S1 -> Y -> P1 -> Bwd, bb: P2 -> S2 -> Z -> P2 -> Bwd, p: Two, s: Two, dz: Z, g: Two) -> Bwd, X>: 862 match p s g: 863 case Two{p1, p2} Two{s1, s2} Two{g1, g2}: seq_at_b2(P1, P2, X, Y, S1, ba, p1, s1, g1, bb(p2, s2, dz, g2)) 864 865 def seq_at_step(-P1: Type, -P2: Type, sa: Opt -> P1 -> P1 -> OptState -> Step, sb: Opt -> P2 -> P2 -> OptState -> Step, +op: Opt, p: Two, g: Two, st: Two, OptState>) -> Step>: 866 match p g st: 867 case Two{p1, p2} Two{g1, g2} Two{s1, s2}: pair_s3(P1, P2, sa(op, p1, g1, s1), sb(op, p2, g2, s2)) 868 869 def seq_at_dump(-P1: Type, -P2: Type, da: P1 -> PL, db: P2 -> PL, p: Two) -> PL>: 870 match p: 871 case Two{p1, p2}: pl_pair(P1, P2, da(p1), db(p2)) 872 873 def seq_at_load2(-P1: Type, -P2: Type, lb: P2 -> Array -> U32 -> PA, p2: P2, r: PA) -> PA>: 874 match r: 875 case PA{p1, xs, off}: pa_second(P1, P2, p1, lb(p2, xs, off)) 876 877 def seq_at_load(-P1: Type, -P2: Type, la: P1 -> Array -> U32 -> PA, lb: P2 -> Array -> U32 -> PA, p: Two, xs: Array, off: U32) -> PA>: 878 match p: 879 case Two{p1, p2}: seq_at_load2(P1, P2, lb, p2, la(p1, xs, off)) 880 881 def seq_at_store(-P1: Type, -P2: Type, la: P1 -> Array -> U32 -> PA, lb: P2 -> Array -> U32 -> PA, p: Two, xs: Array, off: U32) -> PA>: 882 match p: 883 case Two{p1, p2}: seq_at_load2(P1, P2, lb, p2, la(p1, xs, off)) 884 885 def Seq.at(-P1: Type, -P2: Type, -X: Type, -Y: Type, -Z: Type, -S1: Type, -S2: Type, a: Layer, b: Layer) -> Layer, X, Z, Two>: 886 match a b: 887 case Layer{fa, ba, sa, za, da, la, ta, na} Layer{fb, bb, sb, zb, db, lb, tb, nb}: 888 Layer{p => x => seq_at_f(P1, P2, X, Y, Z, S1, S2, fa, fb, p, x), p => s => dz => g => seq_at_b(P1, P2, X, Y, Z, S1, S2, ba, bb, p, s, dz, g), op => p => g => st => seq_at_step(P1, P2, sa, sb, op, p, g, OptState.split(P1, P2, st)), u => Two{za(Unit{}), zb(Unit{})}, p => seq_at_dump(P1, P2, da, db, p), p => xs => off => seq_at_load(P1, P2, la, lb, p, xs, off), p => xs => off => seq_at_store(P1, P2, ta, tb, p, xs, off), u => (na(Unit{}) + nb(Unit{}) : U32)} 889 890 # ===================================================================================================== 891 # Losses: the mean loss over the batch, and its gradient with respect to the output 892 # ===================================================================================================== 893 894 type LossOut<-Y: Type> is Type: 895 LossOut{loss: F32, dy: Y} 896 897 # ===================================================================================================== 898 # Training: one step of any layer under any loss 899 # ===================================================================================================== 900 901 # Parameters, gradient (zero between steps), optimizer-specific buffers, Adam's 1 - b1^t and 1 - b2^t, last loss. 902 type Train<-P: Type> is Type: 903 Train{p: P, g: P, st: OptState

, p1: F32, p2: F32, loss: F32} 904 905 def OptState.new(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer, cfg: OptCfg) -> OptState

: 906 match cfg: 907 case CSGD{lr}: NoMom{} 908 case CMomentum{lr, mu}: OneMom{Layer.zeros(P, X, Y, S, m)} 909 case CAdamWBF16{lr, b1, b2, eps, wd, c1, c2, v16}: Packed{PNew{v16}} 910 case _: TwoMom{Layer.zeros(P, X, Y, S, m), Layer.zeros(P, X, Y, S, m)} 911 912 # 1 - b^t starts at 0; a configuration OptCfg.ok rejects starts at NaN, so every loss after the first is NaN. 913 def q_start(ok: Bool) -> F32: 914 match ok: 915 case True{}: 916 0.0 917 case False{}: 918 (0.0 / 0.0 : F32) 919 920 def Train.new(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer, +cfg: OptCfg, p: P) -> Train

: 921 +q = q_start(OptCfg.ok(cfg)) 922 Train{p, Layer.zeros(P, X, Y, S, m), OptState.new(~P, ~X, ~Y, ~S, ~m, cfg), q, q, 0.0} 923 924 def ts5(-P: Type, +l: F32, +p1: F32, +p2: F32, r: Step

) -> Train

: 925 match r: 926 case Step{p, g, st}: 927 Train{p, g, st, p1, p2, l} 928 929 def ts4(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer, +cfg: OptCfg, st: OptState

, +p1: F32, +p2: F32, +l: F32, r: Bwd) -> Train

: 930 match r: 931 case Bwd{p, dx, g}: 932 +q1 = (p1 + (OptCfg.c1(cfg) * (1.0 - p1 : F32) : F32) : F32) 933 +q2 = (p2 + (OptCfg.c2(cfg) * (1.0 - p2 : F32) : F32) : F32) 934 ts5(P, l, q1, q2, Layer.step(P, X, Y, S, m, Opt.at(cfg, q1, q2), p, g, st)) 935 936 def ts3(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: Layer, +cfg: OptCfg, p: P, s: S, g: P, st: OptState

, +p1: F32, +p2: F32, r: LossOut) -> Train

: 937 match r: 938 case LossOut{l, dy}: 939 ts4(~P, ~X, ~Y, ~S, ~m, cfg, st, p1, p2, l, Layer.bwd(P, X, Y, S, m, p, s, dy, g)) 940 941 def ts2(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~T: Type, ~m: Layer, ~loss: Y -> T -> LossOut, +cfg: OptCfg, g: P, st: OptState

, +p1: F32, +p2: F32, t: T, r: Fwd) -> Train

: 942 match r: 943 case Fwd{p, y, s}: 944 ts3(~P, ~X, ~Y, ~S, ~m, cfg, p, s, g, st, p1, p2, loss(y, t)) 945 946 # forward, loss, backward and one optimizer step: the same code for every model 947 def Train.step(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~T: Type, ~m: Layer, ~loss: Y -> T -> LossOut, +cfg: OptCfg, tr: Train

, x: X, t: T) -> Train

: 948 match tr: 949 case Train{p, g, st, p1, p2, l}: 950 ts2(~P, ~X, ~Y, ~S, ~T, ~m, ~loss, cfg, g, st, p1, p2, t, Layer.fwd(P, X, Y, S, m, p, x)) 951 952 # One owned model, including dimensions captured at run time. Its fwd, bwd and step closures each run once. 953 def train_at_b(-P: Type, -X: Type, sf: Opt -> P -> P -> OptState

-> Step

, +cfg: OptCfg, st: OptState

, +p1: F32, +p2: F32, +l: F32, r: Bwd) -> Train

: 954 match r: 955 case Bwd{p, dx, g}: 956 +q1 = (p1 + (OptCfg.c1(cfg) * (1.0 - p1 : F32) : F32) : F32) 957 +q2 = (p2 + (OptCfg.c2(cfg) * (1.0 - p2 : F32) : F32) : F32) 958 ts5(P, l, q1, q2, sf(Opt.at(cfg, q1, q2), p, g, st)) 959 960 def train_at_loss(-P: Type, -X: Type, -Y: Type, -S: Type, bf: P -> S -> Y -> P -> Bwd, sf: Opt -> P -> P -> OptState

-> Step

, +cfg: OptCfg, p: P, s: S, g: P, st: OptState

, +p1: F32, +p2: F32, r: LossOut) -> Train

: 961 match r: 962 case LossOut{l, dy}: train_at_b(P, X, sf, cfg, st, p1, p2, l, bf(p, s, dy, g)) 963 964 def train_at_f(-P: Type, -X: Type, -Y: Type, -S: Type, -Target: Type, bf: P -> S -> Y -> P -> Bwd, sf: Opt -> P -> P -> OptState

-> Step

, loss: Y -> Target -> LossOut, +cfg: OptCfg, g: P, st: OptState

, +p1: F32, +p2: F32, t: Target, r: Fwd) -> Train

: 965 match r: 966 case Fwd{p, y, s}: train_at_loss(P, X, Y, S, bf, sf, cfg, p, s, g, st, p1, p2, loss(y, t)) 967 968 def Train.step_at(-P: Type, -X: Type, -Y: Type, -S: Type, -Target: Type, m: Layer, loss: Y -> Target -> LossOut, +cfg: OptCfg, tr: Train

, x: X, t: Target) -> Train

: 969 match m tr: 970 case Layer{ff, bf, sf, zeros, dump, load, store, count} Train{p, g, st, p1, p2, l}: 971 train_at_f(P, X, Y, S, Target, bf, sf, loss, cfg, g, st, p1, p2, t, ff(p, x)) 972 973 # ===================================================================================================== 974 # Reading and printing numbers (bend-ml's and llama.bend's loaders, MIT) 975 # ===================================================================================================== 976 977 # U32 and F32 hold the same Word(32n): these are bitcasts 978 def bitcast(u: U32) -> F32: 979 match u: 980 case U32{w}: 981 F32{w} 982 983 def bits(x: F32) -> U32: 984 match x: 985 case F32{w}: 986 U32{w} 987 988 def f32_le(+b0: U32, +b1: U32, +b2: U32, +b3: U32) -> F32: 989 bitcast((b0 + (b1 * 256 : U32) + (b2 * 65536 : U32) + (b3 * 16777216 : U32) : U32)) 990 991 def f32s(bs: List<&2, U32>) -> List<&2, F32>: 992 match bs: 993 case Con{b0, Con{b1, Con{b2, Con{b3, t}}}}: 994 f32_le(b0, b1, b2, b3) <> f32s(t) 995 case _: 996 Nil{} 997 998 def rd_done(m: File & Result<&1, &1, U32 & String, List<&2, U32>>) -> IO(List<&2, U32>): 999 match m: 1000 case Tuple{f, r}: 1001 do IO>: 1002 bs : List<&2, U32> <- IO.pass(List<&2, U32>, r) 1003 File.close(f) 1004 return bs 1005 1006 # n bytes from offset off 1007 def read_at(path: String, +off: U32, +n: U32) -> IO(List<&2, U32>): 1008 do IO>: 1009 f : File <- IO.try(File, File.open(path, "r")) 1010 m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, off, n) 1011 rd_done(m) 1012 1013 def mat_or_die(-r: Nat, -c: Nat, +path: String, m: Maybe<&1, TA.Mat>) -> IO(TA.Mat): 1014 match m: 1015 case None{}: 1016 IO.die(TA.Mat, 1, path ++ ": wrong number of values for this shape") 1017 case Some{x}: 1018 IO.pure(TA.Mat, x) 1019 1020 # an r x c matrix stored as r*c little-endian F32 from byte off of a file 1021 def Mat.load(+r: Nat, +c: Nat, +path: String, +off: U32) -> IO(TA.Mat): 1022 do IO>: 1023 bs : List<&2, U32> <- read_at(path, off, (U32.from_nat(Nat.mul(r, c)) * 4 : U32)) 1024 mat_or_die(r, c, path, TA.Mat.of_list(r, c, f32s(bs))) 1025 1026 def words(xs: List<&2, F32>) -> String: 1027 match xs: 1028 case Nil{}: 1029 "" 1030 case Con{h, t}: 1031 " " ++ U32.show(bits(h)) ++ words(t) 1032 1033 def ml_words(-r: Nat, -c: Nat, m: TA.ML) -> String: 1034 match m: 1035 case TA.ML{d, xs}: 1036 words(xs) 1037 1038 # one line: the name, then each number as the decimal of its 32 bits (exact, for the tests in reference/) 1039 def Mat.print(+r: Nat, +c: Nat, +name: String, m: TA.Mat) -> IO(Unit): 1040 IO.print(name ++ ml_words(r, c, TA.Mat.to_list(r, c, m))) 1041 1042 def F32.print(+name: String, x: F32) -> IO(Unit): 1043 IO.print(name ++ " " ++ U32.show(bits(x))) 1044 1045 # n zeros, in an array with room for them 1046 def zeros_a(+n: Nat) -> Array: 1047 Array.new(F32, depth(n), 0.0) 1048 1049 # ===================================================================================================== 1050 # More element-wise activations 1051 # ===================================================================================================== 1052 1053 # an activation whose derivative needs its input (silu, gelu): saves a copy of the input 1054 def acti_f3(~f: @+a: F32 -> F32, +n: Nat, +c: Nat, m: TA.MM2) -> Fwd, TA.Mat>: 1055 match m: 1056 case TA.MM2{x, s}: 1057 Fwd{Unit{}, Ew.map(~f, n, c, x), s} 1058 1059 def ActIn(~f: @+a: F32 -> F32, ~d: @+a: F32 -> @+b: F32 -> F32, +n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 1060 Layer{p => x => acti_f3(~f, n, c, TA.Mat.clone(n, c, x)), p => s => dy => g => act_b2(n, c, Ew.zip(~d, n, c, dy, s)), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0} 1061 1062 def sigmoid_v(+x: F32) -> F32: 1063 (1.0 / (1.0 + F32.exp(F32.neg(x)) : F32) : F32) 1064 1065 # PyTorch's sigmoid_backward: dy * (1 - y) * y 1066 def sigmoid_d(+dy: F32, +y: F32) -> F32: 1067 ((dy * (1.0 - y : F32) : F32) * y : F32) 1068 1069 def Sigmoid(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 1070 ActOut(~sigmoid_v, ~sigmoid_d, n, c) 1071 1072 def tanh_v(+x: F32) -> F32: 1073 F32.tanh(x) 1074 1075 def tanh_d(+dy: F32, +y: F32) -> F32: 1076 (dy * (1.0 - (y * y : F32) : F32) : F32) 1077 1078 def Tanh(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 1079 ActOut(~tanh_v, ~tanh_d, n, c) 1080 1081 # silu(x) = x * sigmoid(x), computed as PyTorch's CPU kernel: x / (1 + exp(-x)); 1082 # its gradient dy * s * (1 + x * (1 - s)) with s = sigmoid(x) 1083 def silu_v(+x: F32) -> F32: 1084 (x / (1.0 + F32.exp(F32.neg(x)) : F32) : F32) 1085 1086 def silu_d(+dy: F32, +x: F32) -> F32: 1087 +s = sigmoid_v(x) 1088 ((dy * s : F32) * (1.0 + (x * (1.0 - s : F32) : F32) : F32) : F32) 1089 1090 def SiLU(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 1091 ActIn(~silu_v, ~silu_d, n, c) 1092 1093 # GELU in its tanh form (PyTorch's approximate="tanh"): 0.5 x (1 + tanh(sqrt(2/pi) (x + 0.044715 x^3))). 1094 # The exact form, with erf, is GELUErf below. 1095 def gelu_inner(+x: F32) -> F32: 1096 (0.7978845608028654 * (x + (0.044715 * ((x * x : F32) * x : F32) : F32) : F32) : F32) 1097 1098 def gelu_v(+x: F32) -> F32: 1099 (((0.5 * x : F32) * (1.0 + F32.tanh(gelu_inner(x)) : F32)) : F32) 1100 1101 def gelu_d(+dy: F32, +x: F32) -> F32: 1102 +th = F32.tanh(gelu_inner(x)) 1103 +left = (0.5 * x : F32) 1104 +dleft = (0.5 * (1.0 + th : F32) : F32) 1105 +dright = ((left * (1.0 - (th * th : F32) : F32) : F32) * (0.7978845608028654 * (1.0 + ((3.0 * 0.044715 : F32) * (x * x : F32) : F32) : F32) : F32) : F32) 1106 (dy * (dleft + dright : F32) : F32) 1107 1108 def GELU(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 1109 ActIn(~gelu_v, ~gelu_d, n, c) 1110 1111 # ===================================================================================================== 1112 # erf in plain Bend, and GELU in its exact form 1113 # ===================================================================================================== 1114 1115 # erf(x) = 2/sqrt(pi) times the integral of exp(-t^2) from 0 to x. Base has no erf, so it is written here with F32 1116 # additions, multiplications and comparisons only: no libm call, so every target runs the same operations (the emitted C 1117 # turns fused multiply-adds off); the C and JavaScript programs give the same bits (CUDA is not checked yet). erf is 1118 # odd, erf(-x) = -erf(x), so the work is done on a = |x| and the sign put back, in three kinds of pieces: 1119 # - a < 0.84375: erf(x) = x + x P(x^2), with P of degree 6 (fdlibm's erff has this x + x y shape there); it takes 1120 # x itself, so erf(-0) = -0, and NaN reaches it (every comparison with NaN is false) and comes out NaN; 1121 # - 0.84375 <= a < 3.92, in six intervals around c = 1, 1.5, ..., 3.5: erf(a) = K + S(a - c), where K is erf(c) 1122 # rounded to F32 and S a polynomial of degree 7 or 8 whose value is small next to K, so the last addition rounds 1123 # once and carries almost all of the error (a - c is exact: a and c are within a factor of 2); 1124 # - a >= 3.92, infinity included: 1, since erf(a) rounds to 1 in F32 from a = 3.9192 on. 1125 # The coefficients are least-squares fits against a 50-digit erf on Chebyshev points, rounded to F32. Measured on all 1126 # 2^32 inputs against glibc's double erf: at most 0.93 ulp off, the correctly rounded F32 on 97.1% of the inputs, and 1 1127 # ulp from it on the rest; NaN for every NaN. 1128 1129 # one Horner step, c + t s; erf_hm is the step for a negative coefficient -c: t s - c gives the bits of -c + t s 1130 def erf_hp(+c: F32, +t: F32, +s: F32) -> F32: 1131 (c + (t * s : F32) : F32) 1132 1133 def erf_hm(+c: F32, +t: F32, +s: F32) -> F32: 1134 ((t * s : F32) - c : F32) 1135 1136 def erf_small(+x: F32) -> F32: 1137 +z = (x * x : F32) 1138 (x + (x * erf_hp(0.12837916612625122, z, erf_hm(0.37612634897232056, z, erf_hp(0.11283732205629349, z, erf_hm(0.026861993595957756, z, erf_hp(0.005208558402955532, z, erf_hm(0.000823960464913398, z, 8.829222497297451e-05)))))) : F32) : F32) 1139 1140 # [0.84375, 1.25): t = a - 1.0 1141 def erf_p0(+t: F32) -> F32: 1142 (0.8427007794380188 + erf_hp(1.3508899776581984e-08, t, erf_hp(0.41510748863220215, t, erf_hm(0.41510751843452454, t, erf_hp(0.13837093114852905, t, erf_hp(0.06918490678071976, t, erf_hm(0.06929104775190353, t, erf_hp(0.004778923466801643, t, erf_hp(0.017144912853837013, t, F32.neg(0.010648881085216999))))))))) : F32) 1143 1144 # [1.25, 1.75): t = a - 1.5 1145 def erf_p1(+t: F32) -> F32: 1146 (0.9661051630973816 + erf_hm(1.6622220044837377e-08, t, erf_hp(0.11893028765916824, t, erf_hm(0.17839543521404266, t, erf_hp(0.1387522667646408, t, erf_hm(0.04459875077009201, t, erf_hm(0.014879402704536915, t, erf_hp(0.019323812797665596, t, erf_hm(0.004488371312618256, t, F32.neg(0.0023522074334323406))))))))) : F32) 1147 1148 # [1.75, 2.25): t = a - 2.0 1149 def erf_p2(+t: F32) -> F32: 1150 (0.9953222870826721 + erf_hm(2.2064206461891445e-08, t, erf_hp(0.02066698484122753, t, erf_hm(0.04133396968245506, t, erf_hp(0.04822301119565964, t, erf_hm(0.03444507718086243, t, erf_hp(0.013088003732264042, t, erf_hp(0.00046277992078103125, t, erf_hm(0.0033736424520611763, t, 0.0015402256976813078)))))))) : F32) 1151 1152 # [2.25, 2.75): t = a - 2.5 1153 def erf_p3(+t: F32) -> F32: 1154 (0.9995930194854736 + erf_hp(2.846470792405853e-08, t, erf_hp(0.002178284339606762, t, erf_hm(0.005445694085210562, t, erf_hp(0.008350063115358353, t, erf_hm(0.008623690344393253, t, erf_hp(0.00611893879249692, t, erf_hm(0.002765072975307703, t, 0.0005077228415757418))))))) : F32) 1155 1156 # [2.75, 3.25): t = a - 3.0 1157 def erf_p4(+t: F32) -> F32: 1158 (0.9999778866767883 + erf_hp(2.2842012725732275e-08, t, erf_hp(0.00013925303937867284, t, erf_hm(0.0004177672672085464, t, erf_hp(0.0007891026325523853, t, erf_hm(0.0010437482269480824, t, erf_hp(0.001016448251903057, t, erf_hm(0.0007547215791419148, t, 0.0003925197524949908))))))) : F32) 1159 1160 # [3.25, 3.92): t = a - 3.5 1161 def erf_p5(+t: F32) -> F32: 1162 (0.9999992847442627 + erf_hm(2.7855115547481546e-08, t, erf_hp(5.400151167123113e-06, t, erf_hm(1.8894650565925986e-05, t, erf_hp(4.221926792524755e-05, t, erf_hm(6.778862734790891e-05, t, erf_hp(8.422814426012337e-05, t, erf_hm(7.95367595856078e-05, t, 4.151461325818673e-05))))))) : F32) 1163 1164 def erf_sign(+x: F32, +e: F32) -> F32: 1165 pick(F32.is_lt(x, 0.0), F32.neg(e), e) 1166 1167 # the pieces are chosen from the top down, so that only one polynomial is computed (each test is a match on a Bool 1168 # parameter); b is "a is at least the piece's lower end" 1169 def erf_i0(b: Bool, +x: F32, +a: F32) -> F32: 1170 match b: 1171 case True{}: erf_sign(x, erf_p0((a - 1.0 : F32))) 1172 case False{}: erf_small(x) 1173 1174 def erf_i1(b: Bool, +x: F32, +a: F32) -> F32: 1175 match b: 1176 case True{}: erf_sign(x, erf_p1((a - 1.5 : F32))) 1177 case False{}: erf_i0(F32.is_ge(a, 0.84375), x, a) 1178 1179 def erf_i2(b: Bool, +x: F32, +a: F32) -> F32: 1180 match b: 1181 case True{}: erf_sign(x, erf_p2((a - 2.0 : F32))) 1182 case False{}: erf_i1(F32.is_ge(a, 1.25), x, a) 1183 1184 def erf_i3(b: Bool, +x: F32, +a: F32) -> F32: 1185 match b: 1186 case True{}: erf_sign(x, erf_p3((a - 2.5 : F32))) 1187 case False{}: erf_i2(F32.is_ge(a, 1.75), x, a) 1188 1189 def erf_i4(b: Bool, +x: F32, +a: F32) -> F32: 1190 match b: 1191 case True{}: erf_sign(x, erf_p4((a - 3.0 : F32))) 1192 case False{}: erf_i3(F32.is_ge(a, 2.25), x, a) 1193 1194 def erf_i5(b: Bool, +x: F32, +a: F32) -> F32: 1195 match b: 1196 case True{}: erf_sign(x, erf_p5((a - 3.5 : F32))) 1197 case False{}: erf_i4(F32.is_ge(a, 2.75), x, a) 1198 1199 def erf_top(b: Bool, +x: F32, +a: F32) -> F32: 1200 match b: 1201 case True{}: erf_sign(x, 1.0) 1202 case False{}: erf_i5(F32.is_ge(a, 3.25), x, a) 1203 1204 def erf(+x: F32) -> F32: 1205 +a = F32.abs(x) 1206 erf_top(F32.is_ge(a, 3.92), x, a) 1207 1208 # GELU in its exact form, PyTorch's default (approximate="none"): gelu(x) = x/2 (1 + erf(x / sqrt 2)), in PyTorch's CPU 1209 # order, (x * 0.5) * (1 + erf(x * 0.70710677)). Its gradient is cdf + x pdf, with cdf = 0.5 (1 + erf(x / sqrt 2)) and 1210 # pdf = exp(-x^2 / 2) / sqrt(2 pi), computed as PyTorch does: 0.5 (1 + erf(x * 0.70710677)) and 1211 # 0.3989423 * exp((x * x) * -0.5). The exponential is F32.exp, as in SiLU (so on a GPU the gradient would take CUDA's 1212 # float exp). 1213 def gelu_erf_cdf(+x: F32) -> F32: 1214 (0.5 * (1.0 + erf((x * 0.7071067811865476 : F32)) : F32) : F32) 1215 1216 def gelu_erf_v(+x: F32) -> F32: 1217 ((x * 0.5 : F32) * (1.0 + erf((x * 0.7071067811865476 : F32)) : F32) : F32) 1218 1219 def gelu_erf_d(+dy: F32, +x: F32) -> F32: 1220 +pdf = (0.3989422804014327 * F32.exp(((x * x : F32) * F32.neg(0.5) : F32)) : F32) 1221 (dy * (gelu_erf_cdf(x) + (x * pdf : F32) : F32) : F32) 1222 1223 def GELUErf(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 1224 ActIn(~gelu_erf_v, ~gelu_erf_d, n, c) 1225 1226 # ===================================================================================================== 1227 # Dropout: each number is kept with probability 1 - p and scaled by 1 / (1 - p), or set to 0. The mask comes from a 1228 # counter-based hash of (seed, step, index), so a run is reproducible and a twin can compute the same mask. The step 1229 # counter is the layer's "parameter": the optimizer step only increments it, and checkpoints keep it. 1230 # ===================================================================================================== 1231 1232 type DropP is Data: 1233 DropP{t: U32} 1234 1235 # lowbias32 (Chris Wellons): an invertible mix of the 32 bits 1236 def mix32(+x: U32) -> U32: 1237 +a = U32.xor(x, U32.shrn(x, 16n)) 1238 +b = (a * 2146121005 : U32) 1239 +c = U32.xor(b, U32.shrn(b, 15n)) 1240 +d = (c * 2221713035 : U32) 1241 U32.xor(d, U32.shrn(d, 16n)) 1242 1243 # the factor of number i: 1 / (1 - p) if its uniform u = (h >> 8) / 2^24 is >= p, else 0 1244 def drop_k(+key: U32, +i: U32, +p: F32, +scale: F32) -> F32: 1245 pick(F32.is_lt((U32.to_f32(U32.shrn(mix32((key + i : U32)), 8n)) / 16777216.0 : F32), p), 0.0, scale) 1246 1247 def drop_l(nleft: Nat, +i: U32, +key: U32, +p: F32, +scale: F32, a: Array) -> Array: 1248 match nleft: 1249 case 0n: 1250 a 1251 case 1n+q: 1252 drop_l(q, (i + 1 : U32), key, p, scale, Array.set(F32, a, i, drop_k(key, i, p, scale))) 1253 1254 def drop_mask(+n: Nat, +c: Nat, +seed: U32, +t: U32, +p: F32) -> TA.Mat: 1255 TA.Mat{drop_l(Nat.mul(n, c), 0, mix32(U32.xor(seed, mix32(t))), p, (1.0 / (1.0 - p : F32) : F32), zeros_a(Nat.mul(n, c)))} 1256 1257 def drop_f2(-n: Nat, -c: Nat, +t: U32, r: TA.MPair) -> Fwd, TA.Mat>: 1258 match r: 1259 case TA.MPair{y, k}: 1260 Fwd{DropP{t}, y, k} 1261 1262 def drop_fwd(+n: Nat, +c: Nat, +seed: U32, +p: F32, d: DropP, x: TA.Mat) -> Fwd, TA.Mat>: 1263 match d: 1264 case DropP{+t}: 1265 drop_f2(n, c, t, Ew.zip(~f32_mul, n, c, x, drop_mask(n, c, seed, t, p))) 1266 1267 def drop_b2(-n: Nat, -c: Nat, d: DropP, g: DropP, r: TA.MPair) -> Bwd>: 1268 match r: 1269 case TA.MPair{dx, k}: 1270 Bwd{d, dx, g} 1271 1272 def drop_step(+op: Opt, d: DropP, g: DropP, st: OptState) -> Step: 1273 match d: 1274 case DropP{t}: 1275 Step{DropP{(t + 1 : U32)}, g, st} 1276 1277 def drop_dump(d: DropP) -> PL: 1278 match d: 1279 case DropP{+t}: 1280 PL{DropP{t}, [bitcast(t)]} 1281 1282 def drop_load_fin(+off: U32, r: Array & F32) -> PA: 1283 match r: 1284 case Tuple{a, v}: PA{DropP{bits(v)}, a, (off + 1 : U32)} 1285 1286 def drop_load(d: DropP, xs: Array, +off: U32) -> PA: 1287 drop_load_fin(off, Array.get(F32, xs, off)) 1288 1289 # the counter, bitcast to one F32 as dump writes it 1290 def drop_store(d: DropP, xs: Array, +off: U32) -> PA: 1291 match d: 1292 case DropP{+c}: PA{DropP{c}, Array.set(F32, xs, off, bitcast(c)), (off + 1 : U32)} 1293 1294 def Dropout(+n: Nat, +c: Nat, +p: F32, +seed: U32) -> Layer, TA.Mat, TA.Mat>: 1295 Layer{d => x => drop_fwd(n, c, seed, p, d, x), d => s => dy => g => drop_b2(n, c, d, g, Ew.zip(~f32_mul, n, c, dy, s)), op => d => g => st => drop_step(op, d, g, st), u => DropP{0}, d => drop_dump(d), d => xs => off => drop_load(d, xs, off), d => xs => off => drop_store(d, xs, off), u => 1} 1296 1297 # ===================================================================================================== 1298 # Row kernels over arrays: loops over a row of a flat array by index. Every sum runs left to right from +0.0 and every 1299 # expression keeps its order of operations, the ones the row-wise layers had when they worked on lists of rows, so 1300 # they give the same bits (NOTES.md, M6) 1301 # ===================================================================================================== 1302 1303 type D2 is Type: 1304 D2{a: Array, b: Array, s: F32} 1305 1306 type A2 is Type: 1307 A2{a: Array, b: Array} 1308 1309 # a[ia]·b[ib] + a[ia + sa]·b[ib + sb] + ... (k terms) added to acc left to right; ra and rb hold the 1310 # first pair, already read 1311 def sdot(k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, acc: F32, ra: Array & F32, rb: Array & F32) -> D2: 1312 match k ra rb: 1313 case 0n Tuple{a, x} Tuple{b, y}: 1314 D2{a, b, acc} 1315 case 1n+p Tuple{a, +x} Tuple{b, +y}: 1316 sdot(p, (ia + sa : U32), (ib + sb : U32), sa, sb, (acc + (x * y : F32) : F32), Array.get(F32, a, (ia + sa : U32)), Array.get(F32, b, (ib + sb : U32))) 1317 1318 def dot0(+k: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, a: Array, b: Array) -> D2: 1319 sdot(k, ia, ib, sa, sb, 0.0, Array.get(F32, a, ia), Array.get(F32, b, ib)) 1320 1321 1322 type AS is Type: 1323 AS{a: Array, s: F32} 1324 1325 # the largest of a[idx] (held in r) and the k numbers after it: max(max(a0, a1), a2) ..., from the first 1326 def pmax(k: Nat, +idx: U32, cur: F32, r: Array & F32) -> AS: 1327 match k r: 1328 case 0n Tuple{a, +x}: 1329 AS{a, F32.max(cur, x)} 1330 case 1n+q Tuple{a, +x}: 1331 pmax(q, (idx + 1 : U32), F32.max(cur, x), Array.get(F32, a, (idx + 1 : U32))) 1332 1333 def pmax0(k: Nat, +idx: U32, r: Array & F32) -> AS: 1334 match k r: 1335 case 0n Tuple{a, +x}: 1336 AS{a, x} 1337 case 1n+q Tuple{a, +x}: 1338 pmax(q, (idx + 1 : U32), x, Array.get(F32, a, (idx + 1 : U32))) 1339 1340 # a[j] = exp(a[j] - m) for j = idx..idx + k, their sum added to acc in order 1341 def pexp(k: Nat, +idx: U32, +m: F32, acc: F32, r: Array & F32) -> AS: 1342 match k r: 1343 case 0n Tuple{a, +x}: 1344 +e = F32.exp((x - m : F32)) 1345 AS{Array.set(F32, a, idx, e), (acc + e : F32)} 1346 case 1n+q Tuple{a, +x}: 1347 +e = F32.exp((x - m : F32)) 1348 pexp(q, (idx + 1 : U32), m, (acc + e : F32), Array.get(F32, Array.set(F32, a, idx, e), (idx + 1 : U32))) 1349 1350 # the array and how many numbers the last loop of a softmax divided 1351 type PC is Type: 1352 PC{a: Array, c: Nat} 1353 1354 def pc_a(r: PC) -> Array: 1355 match r: 1356 case PC{a, c}: 1357 a 1358 1359 # a[j] = a[j] / s for j = idx..idx + k, counted onto c 1360 def pdiv(k: Nat, +idx: U32, +s: F32, c: Nat, r: Array & F32) -> PC: 1361 match k r: 1362 case 0n Tuple{a, +x}: 1363 PC{Array.set(F32, a, idx, (x / s : F32)), 1n+c} 1364 case 1n+q Tuple{a, +x}: 1365 pdiv(q, (idx + 1 : U32), s, 1n+c, Array.get(F32, Array.set(F32, a, idx, (x / s : F32)), (idx + 1 : U32))) 1366 1367 def psm3(k: Nat, +off: U32, r: AS) -> PC: 1368 match r: 1369 case AS{a, +s}: 1370 pdiv(k, off, s, 0n, Array.get(F32, a, off)) 1371 1372 def psm2(+k: Nat, +off: U32, r: AS) -> PC: 1373 match r: 1374 case AS{a, +m}: 1375 psm3(k, off, pexp(k, off, m, 0.0, Array.get(F32, a, off))) 1376 1377 # softmax in place over the n numbers of p from off (n >= 1), as PyTorch: exp(x - max) / sum; with the count of the 1378 # numbers it divided (the law causal reads it) 1379 def psmc(+n: Nat, +off: U32, p: Array) -> PC: 1380 psm2(Nat.sub(n, 1n), off, pmax0(Nat.sub(n, 1n), off, Array.get(F32, p, off))) 1381 1382 def psm(+n: Nat, +off: U32, p: Array) -> Array: 1383 pc_a(psmc(n, off, p)) 1384 1385 # dst[di..] = src[si..], cnt + 1 numbers; r holds src[si], already read 1386 def copy_n(cnt: Nat, +si: U32, +di: U32, o: Array, r: Array & F32) -> A2: 1387 match cnt r: 1388 case 0n Tuple{s, +x}: 1389 A2{s, Array.set(F32, o, di, x)} 1390 case 1n+n Tuple{s, +x}: 1391 copy_n(n, (si + 1 : U32), (di + 1 : U32), Array.set(F32, o, di, x), Array.get(F32, s, (si + 1 : U32))) 1392 1393 # ===================================================================================================== 1394 # Row-wise layers: softmax, log-softmax, RMSNorm, LayerNorm (each row of an n x c matrix on its own) 1395 # ===================================================================================================== 1396 1397 # rows i.. of x (c wide, left more after i), each softmaxed in place (psm) 1398 def smf_rows(left: Nat, +i: U32, +c: Nat, x: Array) -> Array: 1399 match left: 1400 case 0n: 1401 psm(c, (i * U32.from_nat(c) : U32), x) 1402 case 1n+q: 1403 smf_rows(q, (i + 1 : U32), c, psm(c, (i * U32.from_nat(c) : U32), x)) 1404 1405 # the output, and a copy of it saved for the backward pass 1406 def out_saved(-n: Nat, -c: Nat, r: TA.MM2) -> Fwd, TA.Mat>: 1407 match r: 1408 case TA.MM2{a, b}: 1409 Fwd{Unit{}, a, b} 1410 1411 def sm_fwd(+n: Nat, +c: Nat, x: TA.Mat) -> Fwd, TA.Mat>: 1412 match x: 1413 case TA.Mat{a}: 1414 out_saved(n, c, TA.Mat.clone(n, c, TA.Mat{smf_rows(Nat.sub(n, 1n), 0, c, a)})) 1415 1416 # dy[j] = y[j] * (dy[j] - s), j = idx..idx + k; rd and ry hold dy[idx] and y[idx] 1417 def smb_fix(k: Nat, +idx: U32, +s: F32, rd: Array & F32, ry: Array & F32) -> A2: 1418 match k rd ry: 1419 case 0n Tuple{d, +g} Tuple{y, +v}: 1420 A2{Array.set(F32, d, idx, (v * (g - s : F32) : F32)), y} 1421 case 1n+q Tuple{d, +g} Tuple{y, +v}: 1422 smb_fix(q, (idx + 1 : U32), s, Array.get(F32, Array.set(F32, d, idx, (v * (g - s : F32) : F32)), (idx + 1 : U32)), Array.get(F32, y, (idx + 1 : U32))) 1423 1424 def smb_row(+c: Nat, +off: U32, r: D2) -> A2: 1425 match r: 1426 case D2{d, y, +s}: 1427 smb_fix(Nat.sub(c, 1n), off, s, Array.get(F32, d, off), Array.get(F32, y, off)) 1428 1429 # softmax backward, as PyTorch's: dx = y * (dy - sum(dy * y)), row after row, in place over dy 1430 def smb_rows(left: Nat, +i: U32, +c: Nat, st: A2) -> A2: 1431 match left st: 1432 case 0n A2{d, y}: 1433 smb_row(c, (i * U32.from_nat(c) : U32), dot0(c, (i * U32.from_nat(c) : U32), (i * U32.from_nat(c) : U32), 1, 1, d, y)) 1434 case 1n+q A2{d, y}: 1435 smb_rows(q, (i + 1 : U32), c, smb_row(c, (i * U32.from_nat(c) : U32), dot0(c, (i * U32.from_nat(c) : U32), (i * U32.from_nat(c) : U32), 1, 1, d, y))) 1436 1437 def dx_fin(-n: Nat, -c: Nat, r: A2) -> Bwd>: 1438 match r: 1439 case A2{d, y}: 1440 Bwd{Unit{}, TA.Mat{d}, Unit{}} 1441 1442 def sm_bwd(+n: Nat, +c: Nat, s: TA.Mat, dy: TA.Mat) -> Bwd>: 1443 match s dy: 1444 case TA.Mat{y} TA.Mat{d}: 1445 dx_fin(n, c, smb_rows(Nat.sub(n, 1n), 0, c, A2{d, y})) 1446 1447 # softmax over each row; saves its output 1448 def Softmax(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 1449 Layer{p => x => sm_fwd(n, c, x), p => s => dy => g => sm_bwd(n, c, s, dy), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0} 1450 1451 # log_softmax(x) = (x - max) - log(sum(exp(x - max))), as PyTorch's CPU kernel 1452 1453 # acc + exp(a[j] - m) for j = idx..idx + k, in order 1454 def lsum(k: Nat, +idx: U32, +m: F32, acc: F32, r: Array & F32) -> AS: 1455 match k r: 1456 case 0n Tuple{a, +x}: 1457 AS{a, (acc + F32.exp((x - m : F32)) : F32)} 1458 case 1n+q Tuple{a, +x}: 1459 lsum(q, (idx + 1 : U32), m, (acc + F32.exp((x - m : F32)) : F32), Array.get(F32, a, (idx + 1 : U32))) 1460 1461 # a[j] = (a[j] - m) - l for j = idx..idx + k 1462 def lw(k: Nat, +idx: U32, +m: F32, +l: F32, r: Array & F32) -> Array: 1463 match k r: 1464 case 0n Tuple{a, +x}: 1465 Array.set(F32, a, idx, ((x - m : F32) - l : F32)) 1466 case 1n+q Tuple{a, +x}: 1467 lw(q, (idx + 1 : U32), m, l, Array.get(F32, Array.set(F32, a, idx, ((x - m : F32) - l : F32)), (idx + 1 : U32))) 1468 1469 def lsm_r3(+c: Nat, +off: U32, +m: F32, r: AS) -> Array: 1470 match r: 1471 case AS{a, +s}: 1472 lw(Nat.sub(c, 1n), off, m, F32.log(s), Array.get(F32, a, off)) 1473 1474 def lsm_r2(+c: Nat, +off: U32, r: AS) -> Array: 1475 match r: 1476 case AS{a, +m}: 1477 lsm_r3(c, off, m, lsum(Nat.sub(c, 1n), off, m, 0.0, Array.get(F32, a, off))) 1478 1479 def lsm_row(+c: Nat, +off: U32, x: Array) -> Array: 1480 lsm_r2(c, off, pmax0(Nat.sub(c, 1n), off, Array.get(F32, x, off))) 1481 1482 def lsmf_rows(left: Nat, +i: U32, +c: Nat, x: Array) -> Array: 1483 match left: 1484 case 0n: 1485 lsm_row(c, (i * U32.from_nat(c) : U32), x) 1486 case 1n+q: 1487 lsmf_rows(q, (i + 1 : U32), c, lsm_row(c, (i * U32.from_nat(c) : U32), x)) 1488 1489 def lsm_fwd(+n: Nat, +c: Nat, x: TA.Mat) -> Fwd, TA.Mat>: 1490 match x: 1491 case TA.Mat{a}: 1492 out_saved(n, c, TA.Mat.clone(n, c, TA.Mat{lsmf_rows(Nat.sub(n, 1n), 0, c, a)})) 1493 1494 # acc + a[j] for j = idx..idx + k, in order 1495 def asum(k: Nat, +idx: U32, acc: F32, r: Array & F32) -> AS: 1496 match k r: 1497 case 0n Tuple{a, +x}: 1498 AS{a, (acc + x : F32)} 1499 case 1n+q Tuple{a, +x}: 1500 asum(q, (idx + 1 : U32), (acc + x : F32), Array.get(F32, a, (idx + 1 : U32))) 1501 1502 # dy[j] = dy[j] - exp(y[j]) * s, j = idx..idx + k 1503 def lsmb_fix(k: Nat, +idx: U32, +s: F32, rd: Array & F32, ry: Array & F32) -> A2: 1504 match k rd ry: 1505 case 0n Tuple{d, +g} Tuple{y, +v}: 1506 A2{Array.set(F32, d, idx, (g - (F32.exp(v) * s : F32) : F32)), y} 1507 case 1n+q Tuple{d, +g} Tuple{y, +v}: 1508 lsmb_fix(q, (idx + 1 : U32), s, Array.get(F32, Array.set(F32, d, idx, (g - (F32.exp(v) * s : F32) : F32)), (idx + 1 : U32)), Array.get(F32, y, (idx + 1 : U32))) 1509 1510 def lsmb_row(+c: Nat, +off: U32, y: Array, r: AS) -> A2: 1511 match r: 1512 case AS{d, +s}: 1513 lsmb_fix(Nat.sub(c, 1n), off, s, Array.get(F32, d, off), Array.get(F32, y, off)) 1514 1515 # dx = dy - exp(y) * sum(dy), row after row, in place over dy 1516 def lsmb_rows(left: Nat, +i: U32, +c: Nat, st: A2) -> A2: 1517 match left st: 1518 case 0n A2{d, y}: 1519 lsmb_row(c, (i * U32.from_nat(c) : U32), y, asum(Nat.sub(c, 1n), (i * U32.from_nat(c) : U32), 0.0, Array.get(F32, d, (i * U32.from_nat(c) : U32)))) 1520 case 1n+q A2{d, y}: 1521 lsmb_rows(q, (i + 1 : U32), c, lsmb_row(c, (i * U32.from_nat(c) : U32), y, asum(Nat.sub(c, 1n), (i * U32.from_nat(c) : U32), 0.0, Array.get(F32, d, (i * U32.from_nat(c) : U32))))) 1522 1523 def lsm_bwd(+n: Nat, +c: Nat, s: TA.Mat, dy: TA.Mat) -> Bwd>: 1524 match s dy: 1525 case TA.Mat{y} TA.Mat{d}: 1526 dx_fin(n, c, lsmb_rows(Nat.sub(n, 1n), 0, c, A2{d, y})) 1527 1528 def LogSoftmax(+n: Nat, +c: Nat) -> Layer, TA.Mat, TA.Mat>: 1529 Layer{p => x => lsm_fwd(n, c, x), p => s => dy => g => lsm_bwd(n, c, s, dy), op => p => g => st => no_step(Unit, op, p, g, st), u => Unit{}, p => PL{p, Nil{}}, p => xs => off => PA{p, xs, off}, p => xs => off => PA{p, xs, off}, u => 0} 1530 1531 # ---- RMSNorm: y = g * (x / sqrt(mean(x^2) + eps)); its parameter is the gain g (1 x d) ---- 1532 1533 # one optimizer step on a layer whose parameter is a single matrix 1534 def mat_step_fin(-r: Nat, -c: Nat, q: MatStep) -> Step>: 1535 match q: 1536 case MatStep{w, g, st}: Step{w, g, st} 1537 1538 def mat_step(+r: Nat, +c: Nat, +op: Opt, p: TA.Mat, g: TA.Mat, st: OptState>) -> Step>: 1539 mat_step_fin(r, c, Opt.mat(r, c, op, p, g, st)) 1540 1541 def g_add_fin(-r: Nat, -c: Nat, r2: TA.MPair) -> TA.Mat: 1542 match r2: 1543 case TA.MPair{a, b}: 1544 a 1545 1546 # a + b, dropping b 1547 def Mat.sum2(+r: Nat, +c: Nat, a: TA.Mat, b: TA.Mat) -> TA.Mat: 1548 g_add_fin(r, c, Mat.add(r, c, a, b)) 1549 1550 def dup_fin(-r: Nat, -c: Nat, m: TA.MM2) -> Two, TA.Mat>: 1551 match m: 1552 case TA.MM2{a, b}: 1553 Two{a, b} 1554 1555 # two copies of a matrix 1556 def Mat.dup(+r: Nat, +c: Nat, m: TA.Mat) -> Two, TA.Mat>: 1557 dup_fin(r, c, TA.Mat.clone(r, c, m)) 1558 1559 def ml_list(-r: Nat, -c: Nat, ml: TA.ML) -> List<&2, F32>: 1560 match ml: 1561 case TA.ML{m, xs}: 1562 xs 1563 1564 # ---- RMSNorm layer, a row at a time over the arrays: x (n x d), the gain g (1 x d), the output y. Each sum is added 1565 # left to right from +0.0, each product in the order the list version multiplied, so the bits are the list version's. 1566 # ---- 1567 1568 # acc + a[i]·a[i] + ... over k + 1 numbers from i, in order; r holds a[i], already read 1569 def ssq(k: Nat, +i: U32, acc: F32, r: Array & F32) -> AS: 1570 match k r: 1571 case 0n Tuple{a, +x}: 1572 AS{a, (acc + (x * x : F32) : F32)} 1573 case 1n+p Tuple{a, +x}: 1574 ssq(p, (i + 1 : U32), (acc + (x * x : F32) : F32), Array.get(F32, a, (i + 1 : U32))) 1575 1576 # 1 / rms of a row whose sum of squares is s 1577 def rms_inv(+s: F32, +df: F32, +eps: F32) -> F32: 1578 (1.0 / F32.sqrt(((s / df : F32) + eps : F32)) : F32) 1579 1580 type R3 is Type: 1581 R3{x: Array, g: Array, y: Array} 1582 1583 # y[off + j] = g[j] * (x[off + j] * r), j = jj..jj + k; rx and rg hold x[off + jj] and g[jj] 1584 def rms_w(k: Nat, +j: U32, +off: U32, +r: F32, y: Array, rx: Array & F32, rg: Array & F32) -> R3: 1585 match k rx rg: 1586 case 0n Tuple{x, +a} Tuple{g, +b}: 1587 R3{x, g, Array.set(F32, y, (off + j : U32), (b * (a * r : F32) : F32))} 1588 case 1n+p Tuple{x, +a} Tuple{g, +b}: 1589 rms_w(p, (j + 1 : U32), off, r, Array.set(F32, y, (off + j : U32), (b * (a * r : F32) : F32)), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32))) 1590 1591 def rmsf_row2(+d: Nat, +off: U32, +df: F32, +eps: F32, g: Array, y: Array, r: AS) -> R3: 1592 match r: 1593 case AS{x, +s}: 1594 rms_w(Nat.sub(d, 1n), 0, off, rms_inv(s, df, eps), y, Array.get(F32, x, off), Array.get(F32, g, 0)) 1595 1596 def rmsf_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R3) -> R3: 1597 match s: 1598 case R3{x, g, y}: 1599 rmsf_row2(d, off, df, eps, g, y, ssq(Nat.sub(d, 1n), off, 0.0, Array.get(F32, x, off))) 1600 1601 # rows i.. (left + 1 of them, d wide) 1602 def rmsf_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R3) -> R3: 1603 match left: 1604 case 0n: 1605 rmsf_row(d, (i * U32.from_nat(d) : U32), df, eps, s) 1606 case 1n+p: 1607 rmsf_rows(p, (i + 1 : U32), d, df, eps, rmsf_row(d, (i * U32.from_nat(d) : U32), df, eps, s)) 1608 1609 def rms_f2(-n: Nat, -d: Nat, s: R3) -> Fwd, TA.Mat, TA.Mat>: 1610 match s: 1611 case R3{x, g, y}: 1612 Fwd{TA.Mat{g}, TA.Mat{y}, TA.Mat{x}} 1613 1614 def rms_fwd(+n: Nat, +d: Nat, +eps: F32, p: TA.Mat<1n, d>, x: TA.Mat) -> Fwd, TA.Mat, TA.Mat>: 1615 match p x: 1616 case TA.Mat{g} TA.Mat{xa}: 1617 rms_f2(n, d, rmsf_rows(Nat.sub(n, 1n), 0, d, F32.from_nat(d), eps, R3{xa, g, zeros_a(Nat.mul(n, d))})) 1618 1619 # with r = 1/rms, xh = x*r, dxh = dy*g: dx = r * (dxh - xh * mean(dxh * xh)); dg += dy * xh, row after row 1620 type R5 is Type: 1621 R5{x: Array, dy: Array, g: Array, dx: Array, dg: Array} 1622 1623 type R3S is Type: 1624 R3S{x: Array, dy: Array, g: Array, s: F32} 1625 1626 # acc + dxh[j]·xh[j] + ..., in order; rx, rd, rg hold x[off + jj], dy[off + jj], g[jj] 1627 def rmsb_dot(k: Nat, +j: U32, +off: U32, +r: F32, acc: F32, rx: Array & F32, rd: Array & F32, rg: Array & F32) -> R3S: 1628 match k rx rd rg: 1629 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1630 R3S{x, dy, g, (acc + ((b * c : F32) * (a * r : F32) : F32) : F32)} 1631 case 1n+p Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1632 rmsb_dot(p, (j + 1 : U32), off, r, (acc + ((b * c : F32) * (a * r : F32) : F32) : F32), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32))) 1633 1634 def dg_add(+j: U32, +v: F32, r: Array & F32) -> Array: 1635 match r: 1636 case Tuple{dg, +old}: 1637 Array.set(F32, dg, j, (old + v : F32)) 1638 1639 # dx[off + j] = ((dy[j]·g[j]) - (x[j]·r)·m)·r and dg[j] += dy[j]·(x[j]·r), j = jj..jj + k 1640 def rmsb_w(k: Nat, +j: U32, +off: U32, +r: F32, +m: F32, dx: Array, dg: Array, rx: Array & F32, rd: Array & F32, rg: Array & F32) -> R5: 1641 match k rx rd rg: 1642 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1643 +xh = (a * r : F32) 1644 R5{x, dy, g, Array.set(F32, dx, (off + j : U32), (((b * c : F32) - (xh * m : F32) : F32) * r : F32)), dg_add(j, (b * xh : F32), Array.get(F32, dg, j))} 1645 case 1n+p Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1646 +xh = (a * r : F32) 1647 rmsb_w(p, (j + 1 : U32), off, r, m, Array.set(F32, dx, (off + j : U32), (((b * c : F32) - (xh * m : F32) : F32) * r : F32)), dg_add(j, (b * xh : F32), Array.get(F32, dg, j)), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32))) 1648 1649 def rmsb_row3(+d: Nat, +off: U32, +df: F32, +r: F32, dx: Array, dg: Array, q: R3S) -> R5: 1650 match q: 1651 case R3S{x, dy, g, +s}: 1652 rmsb_w(Nat.sub(d, 1n), 0, off, r, (s / df : F32), dx, dg, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0)) 1653 1654 def rmsb_row2(+d: Nat, +off: U32, +df: F32, +eps: F32, dy: Array, g: Array, dx: Array, dg: Array, q: AS) -> R5: 1655 match q: 1656 case AS{x, +s}: 1657 +r = rms_inv(s, df, eps) 1658 rmsb_row3(d, off, df, r, dx, dg, rmsb_dot(Nat.sub(d, 1n), 0, off, r, 0.0, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0))) 1659 1660 def rmsb_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R5) -> R5: 1661 match s: 1662 case R5{x, dy, g, dx, dg}: 1663 rmsb_row2(d, off, df, eps, dy, g, dx, dg, ssq(Nat.sub(d, 1n), off, 0.0, Array.get(F32, x, off))) 1664 1665 def rmsb_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R5) -> R5: 1666 match left: 1667 case 0n: 1668 rmsb_row(d, (i * U32.from_nat(d) : U32), df, eps, s) 1669 case 1n+p: 1670 rmsb_rows(p, (i + 1 : U32), d, df, eps, rmsb_row(d, (i * U32.from_nat(d) : U32), df, eps, s)) 1671 1672 def gain_bwd_fin(-n: Nat, -d: Nat, -P: Type, p: P, dx: TA.Mat, g: P) -> Bwd>: 1673 Bwd{p, dx, g} 1674 1675 def rms_b2(+d: Nat, -n: Nat, gacc: TA.Mat<1n, d>, s: R5) -> Bwd, TA.Mat>: 1676 match s: 1677 case R5{x, dy, g, dx, dg}: 1678 gain_bwd_fin(n, d, TA.Mat<1n, d>, TA.Mat{g}, TA.Mat{dx}, Mat.sum2(1n, d, gacc, TA.Mat{dg})) 1679 1680 def rms_bwd(+n: Nat, +d: Nat, +eps: F32, p: TA.Mat<1n, d>, x: TA.Mat, dy: TA.Mat, gacc: TA.Mat<1n, d>) -> Bwd, TA.Mat>: 1681 match p x dy: 1682 case TA.Mat{g} TA.Mat{xa} TA.Mat{da}: 1683 rms_b2(d, n, gacc, rmsb_rows(Nat.sub(n, 1n), 0, d, F32.from_nat(d), eps, R5{xa, da, g, zeros_a(Nat.mul(n, d)), zeros_a(d)})) 1684 1685 # RMSNorm over rows of d numbers, with a learned gain (Llama's norm); saves its input 1686 def RMSNorm(+n: Nat, +d: Nat, +eps: F32) -> Layer, TA.Mat, TA.Mat, TA.Mat>: 1687 Layer{p => x => rms_fwd(n, d, eps, p, x), p => s => dy => g => rms_bwd(n, d, eps, p, s, dy, g), op => p => g => st => mat_step(1n, d, op, p, g, st), u => TA.Mat.zeros(1n, d), p => Mat.dump(1n, d, p), p => xs => off => Mat.restore(1n, d, p, xs, off), p => xs => off => Mat.store(1n, d, p, xs, off), u => U32.from_nat(d)} 1688 1689 # ---- LayerNorm: y = (x - mean) / sqrt(var + eps) * g + b, var the biased variance ---- 1690 1691 type LN<-d: Nat> is Type: 1692 LN{g: TA.Mat<1n, d>, b: TA.Mat<1n, d>} 1693 1694 # acc + (a[j] - mu)·(a[j] - mu) for j = idx..idx + k, in order 1695 def csq(k: Nat, +idx: U32, +mu: F32, acc: F32, r: Array & F32) -> AS: 1696 match k r: 1697 case 0n Tuple{a, +x}: 1698 AS{a, (acc + ((x - mu : F32) * (x - mu : F32) : F32) : F32)} 1699 case 1n+q Tuple{a, +x}: 1700 csq(q, (idx + 1 : U32), mu, (acc + ((x - mu : F32) * (x - mu : F32) : F32) : F32), Array.get(F32, a, (idx + 1 : U32))) 1701 1702 type ASM is Type: 1703 ASM{a: Array, mu: F32, r: F32} 1704 1705 def ln_mr2(+df: F32, +eps: F32, +mu: F32, q: AS) -> ASM: 1706 match q: 1707 case AS{a, +v}: 1708 ASM{a, mu, (1.0 / F32.sqrt(((v / df : F32) + eps : F32)) : F32)} 1709 1710 def ln_mr1(+d: Nat, +off: U32, +df: F32, +eps: F32, q: AS) -> ASM: 1711 match q: 1712 case AS{a, +s}: 1713 ln_mr2(df, eps, (s / df : F32), csq(Nat.sub(d, 1n), off, (s / df : F32), 0.0, Array.get(F32, a, off))) 1714 1715 # a row's mean and 1 / sqrt(var + eps), as ln_xh and ln_r computed them 1716 def ln_mr(+d: Nat, +off: U32, +df: F32, +eps: F32, x: Array) -> ASM: 1717 ln_mr1(d, off, df, eps, asum(Nat.sub(d, 1n), off, 0.0, Array.get(F32, x, off))) 1718 1719 type R4 is Type: 1720 R4{x: Array, g: Array, b: Array, y: Array} 1721 1722 # y[off + j] = ((x - mu) * r) * g[j] + b[j], j = jj..jj + k 1723 def lnf_w(k: Nat, +j: U32, +off: U32, +mu: F32, +r: F32, y: Array, rx: Array & F32, rg: Array & F32, rb: Array & F32) -> R4: 1724 match k rx rg rb: 1725 case 0n Tuple{x, +a} Tuple{g, +gv} Tuple{b, +bv}: 1726 R4{x, g, b, Array.set(F32, y, (off + j : U32), ((((a - mu : F32) * r : F32) * gv : F32) + bv : F32))} 1727 case 1n+q Tuple{x, +a} Tuple{g, +gv} Tuple{b, +bv}: 1728 lnf_w(q, (j + 1 : U32), off, mu, r, Array.set(F32, y, (off + j : U32), ((((a - mu : F32) * r : F32) * gv : F32) + bv : F32)), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32)), Array.get(F32, b, (j + 1 : U32))) 1729 1730 def lnf_row2(+d: Nat, +off: U32, g: Array, b: Array, y: Array, m: ASM) -> R4: 1731 match m: 1732 case ASM{x, +mu, +r}: 1733 lnf_w(Nat.sub(d, 1n), 0, off, mu, r, y, Array.get(F32, x, off), Array.get(F32, g, 0), Array.get(F32, b, 0)) 1734 1735 def lnf_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R4) -> R4: 1736 match s: 1737 case R4{x, g, b, y}: 1738 lnf_row2(d, off, g, b, y, ln_mr(d, off, df, eps, x)) 1739 1740 def lnf_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R4) -> R4: 1741 match left: 1742 case 0n: 1743 lnf_row(d, (i * U32.from_nat(d) : U32), df, eps, s) 1744 case 1n+q: 1745 lnf_rows(q, (i + 1 : U32), d, df, eps, lnf_row(d, (i * U32.from_nat(d) : U32), df, eps, s)) 1746 1747 def ln_f2(-n: Nat, -d: Nat, s: R4) -> Fwd, TA.Mat, TA.Mat>: 1748 match s: 1749 case R4{x, g, b, y}: 1750 Fwd{LN{TA.Mat{g}, TA.Mat{b}}, TA.Mat{y}, TA.Mat{x}} 1751 1752 def ln_fwd(+n: Nat, +d: Nat, +eps: F32, p: LN, x: TA.Mat) -> Fwd, TA.Mat, TA.Mat>: 1753 match p x: 1754 case LN{TA.Mat{g}, TA.Mat{b}} TA.Mat{xa}: 1755 ln_f2(n, d, lnf_rows(Nat.sub(n, 1n), 0, d, F32.from_nat(d), eps, R4{xa, g, b, zeros_a(Nat.mul(n, d))})) 1756 1757 # with xh = (x - mu) * r and dxh = dy * g: dx = r * ((dxh - mean(dxh)) - xh * mean(dxh * xh)); dg += dy * xh; db += dy 1758 type R3SS is Type: 1759 R3SS{x: Array, dy: Array, g: Array, sa: F32, sb: F32} 1760 1761 # sa + dxh[j] and sb + dxh[j]·xh[j], j = jj..jj + k, each in order 1762 def lnb_sums(k: Nat, +j: U32, +off: U32, +mu: F32, +r: F32, sa: F32, sb: F32, rx: Array & F32, rd: Array & F32, rg: Array & F32) -> R3SS: 1763 match k rx rd rg: 1764 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1765 R3SS{x, dy, g, (sa + (b * c : F32) : F32), (sb + ((b * c : F32) * ((a - mu : F32) * r : F32) : F32) : F32)} 1766 case 1n+q Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1767 lnb_sums(q, (j + 1 : U32), off, mu, r, (sa + (b * c : F32) : F32), (sb + ((b * c : F32) * ((a - mu : F32) * r : F32) : F32) : F32), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32))) 1768 1769 def ln_dw(+j: U32, +v: F32, +w: F32, st: A2) -> A2: 1770 match st: 1771 case A2{dg, db}: 1772 A2{dg_add(j, v, Array.get(F32, dg, j)), dg_add(j, w, Array.get(F32, db, j))} 1773 1774 type R4B is Type: 1775 R4B{x: Array, dy: Array, g: Array, dx: Array, w: A2} 1776 1777 # dx[off + j] = (((dxh - ma) - xh * mb) * r), dg[j] += dy * xh, db[j] += dy, j = jj..jj + k (ma, mb the two means) 1778 def lnb_w(k: Nat, +j: U32, +off: U32, +mu: F32, +r: F32, +ma: F32, +mb: F32, dx: Array, w: A2, rx: Array & F32, rd: Array & F32, rg: Array & F32) -> R4B: 1779 match k rx rd rg: 1780 case 0n Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1781 +xh = ((a - mu : F32) * r : F32) 1782 R4B{x, dy, g, Array.set(F32, dx, (off + j : U32), ((((b * c : F32) - ma : F32) - (xh * mb : F32) : F32) * r : F32)), ln_dw(j, (b * xh : F32), b, w)} 1783 case 1n+q Tuple{x, +a} Tuple{dy, +b} Tuple{g, +c}: 1784 +xh = ((a - mu : F32) * r : F32) 1785 lnb_w(q, (j + 1 : U32), off, mu, r, ma, mb, Array.set(F32, dx, (off + j : U32), ((((b * c : F32) - ma : F32) - (xh * mb : F32) : F32) * r : F32)), ln_dw(j, (b * xh : F32), b, w), Array.get(F32, x, ((off + j : U32) + 1 : U32)), Array.get(F32, dy, ((off + j : U32) + 1 : U32)), Array.get(F32, g, (j + 1 : U32))) 1786 1787 def lnb_row3(+d: Nat, +off: U32, +df: F32, +mu: F32, +r: F32, dx: Array, w: A2, q: R3SS) -> R4B: 1788 match q: 1789 case R3SS{x, dy, g, +sa, +sb}: 1790 lnb_w(Nat.sub(d, 1n), 0, off, mu, r, (sa / df : F32), (sb / df : F32), dx, w, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0)) 1791 1792 def lnb_row2(+d: Nat, +off: U32, +df: F32, dy: Array, g: Array, dx: Array, w: A2, m: ASM) -> R4B: 1793 match m: 1794 case ASM{x, +mu, +r}: 1795 lnb_row3(d, off, df, mu, r, dx, w, lnb_sums(Nat.sub(d, 1n), 0, off, mu, r, 0.0, 0.0, Array.get(F32, x, off), Array.get(F32, dy, off), Array.get(F32, g, 0))) 1796 1797 def lnb_row(+d: Nat, +off: U32, +df: F32, +eps: F32, s: R4B) -> R4B: 1798 match s: 1799 case R4B{x, dy, g, dx, w}: 1800 lnb_row2(d, off, df, dy, g, dx, w, ln_mr(d, off, df, eps, x)) 1801 1802 def lnb_rows(left: Nat, +i: U32, +d: Nat, +df: F32, +eps: F32, s: R4B) -> R4B: 1803 match left: 1804 case 0n: 1805 lnb_row(d, (i * U32.from_nat(d) : U32), df, eps, s) 1806 case 1n+q: 1807 lnb_rows(q, (i + 1 : U32), d, df, eps, lnb_row(d, (i * U32.from_nat(d) : U32), df, eps, s)) 1808 1809 def ln_b2(+d: Nat, -n: Nat, b: TA.Mat<1n, d>, gg: LN, s: R4B) -> Bwd, TA.Mat>: 1810 match gg s: 1811 case LN{ga, gb} R4B{x, dy, g, dx, A2{dg, db}}: 1812 gain_bwd_fin(n, d, LN, LN{TA.Mat{g}, b}, TA.Mat{dx}, LN{Mat.sum2(1n, d, ga, TA.Mat{dg}), Mat.sum2(1n, d, gb, TA.Mat{db})}) 1813 1814 def ln_bwd(+n: Nat, +d: Nat, +eps: F32, p: LN, x: TA.Mat, dy: TA.Mat, gg: LN) -> Bwd, TA.Mat>: 1815 match p x dy: 1816 case LN{TA.Mat{g}, b} TA.Mat{xa} TA.Mat{da}: 1817 ln_b2(d, n, b, gg, lnb_rows(Nat.sub(n, 1n), 0, d, F32.from_nat(d), eps, R4B{xa, da, g, zeros_a(Nat.mul(n, d)), A2{zeros_a(d), zeros_a(d)}})) 1818 1819 def ln_pair(-d: Nat, p: LN) -> Two, TA.Mat<1n, d>>: 1820 match p: 1821 case LN{w, b}: Two{w, b} 1822 1823 def ln_unpair(-d: Nat, p: Two, TA.Mat<1n, d>>) -> LN: 1824 match p: 1825 case Two{w, b}: LN{w, b} 1826 1827 def ln_s2(-d: Nat, a: MatStep<1n, d>, c: MatStep<1n, d>) -> Step>: 1828 match a c: 1829 case MatStep{w, gw, sw} MatStep{b, gb, sb}: 1830 Step{LN{w, b}, LN{gw, gb}, OptState.map(Two, TA.Mat<1n, d>>, LN, OptState.join(TA.Mat<1n, d>, TA.Mat<1n, d>, sw, sb), x => ln_unpair(d, x), x => ln_unpair(d, x))} 1831 1832 def ln_step2(+d: Nat, +op: Opt, p: LN, g: LN, st: Two>, OptState>>) -> Step>: 1833 match p g st: 1834 case LN{w, b} LN{gw, gb} Two{sw, sb}: 1835 ln_s2(d, Opt.mat(1n, d, op, w, gw, sw), Opt.mat(1n, d, op, b, gb, sb)) 1836 1837 def ln_step(+d: Nat, +op: Opt, p: LN, g: LN, st: OptState>) -> Step>: 1838 ln_step2(d, op, p, g, OptState.split(TA.Mat<1n, d>, TA.Mat<1n, d>, OptState.map(LN, Two, TA.Mat<1n, d>>, st, x => ln_pair(d, x), x => ln_pair(d, x)))) 1839 1840 def ln_dump2(-d: Nat, r: PL, TA.Mat<1n, d>>>) -> PL>: 1841 match r: 1842 case PL{Two{g, b}, xs}: 1843 PL{LN{g, b}, xs} 1844 1845 def ln_dump(+d: Nat, p: LN) -> PL>: 1846 match p: 1847 case LN{g, b}: 1848 ln_dump2(d, pl_pair(TA.Mat<1n, d>, TA.Mat<1n, d>, Mat.dump(1n, d, g), Mat.dump(1n, d, b))) 1849 1850 def ln_load3(-d: Nat, r: PA, TA.Mat<1n, d>>>) -> PA>: 1851 match r: 1852 case PA{Two{g, b}, xs, off}: PA{LN{g, b}, xs, off} 1853 1854 def ln_load2(+d: Nat, b: TA.Mat<1n, d>, r: PA>) -> PA>: 1855 match r: 1856 case PA{g, xs, off}: 1857 ln_load3(d, pa_second(TA.Mat<1n, d>, TA.Mat<1n, d>, g, Mat.restore(1n, d, b, xs, off))) 1858 1859 def ln_load(+d: Nat, p: LN, xs: Array, off: U32) -> PA>: 1860 match p: 1861 case LN{g, b}: ln_load2(d, b, Mat.restore(1n, d, g, xs, off)) 1862 1863 def ln_store2(+d: Nat, b: TA.Mat<1n, d>, r: PA>) -> PA>: 1864 match r: 1865 case PA{g, xs, off}: 1866 ln_load3(d, pa_second(TA.Mat<1n, d>, TA.Mat<1n, d>, g, Mat.store(1n, d, b, xs, off))) 1867 1868 def ln_store(+d: Nat, p: LN, xs: Array, off: U32) -> PA>: 1869 match p: 1870 case LN{g, b}: ln_store2(d, b, Mat.store(1n, d, g, xs, off)) 1871 1872 def LayerNorm(+n: Nat, +d: Nat, +eps: F32) -> Layer, TA.Mat, TA.Mat, TA.Mat>: 1873 Layer{p => x => ln_fwd(n, d, eps, p, x), p => s => dy => g => ln_bwd(n, d, eps, p, s, dy, g), op => p => g => st => ln_step(d, op, p, g, st), u => LN{TA.Mat.zeros(1n, d), TA.Mat.zeros(1n, d)}, p => ln_dump(d, p), p => xs => off => ln_load(d, p, xs, off), p => xs => off => ln_store(d, p, xs, off), u => U32.from_nat(Nat.mul(2n, d))} 1874 1875 # ---- GroupNorm: LayerNorm over one image's group of channels, with a gain and a bias per channel ---- 1876 # Images are channels last: x is (n·hw) x c with c = ng·cg, image s owns rows s·hw .. s·hw + hw - 1 and group k columns 1877 # k·cg .. k·cg + cg - 1. Each (image, group) block of hw·cg numbers is normalized by its own mean and biased variance, 1878 # then y = xh * g[ch] + b[ch], as nn.GroupNorm. Writing c as ng·cg in the type makes "ng divides c" hold by 1879 # construction. A block is walked one row segment at a time with LayerNorm's loops: `off` is the row's first cell and 1880 # `j` the group's first channel, so a segment's cells are off + j .. off + j + cg - 1 and its gains g[j .. j + cg - 1]. 1881 1882 # acc + the segments of the block's remaining `left` rows, each one row (c cells) below the last, in order 1883 def gn_sum(left: Nat, +off: U32, +c: U32, +cg: Nat, q: AS) -> AS: 1884 match left q: 1885 case 0n r: 1886 r 1887 case 1n+l AS{a, +s}: 1888 gn_sum(l, (off + c : U32), c, cg, asum(Nat.sub(cg, 1n), (off + c : U32), s, Array.get(F32, a, (off + c : U32)))) 1889 1890 # acc + (x - mu)·(x - mu) over the same segments 1891 def gn_csq(left: Nat, +off: U32, +c: U32, +cg: Nat, +mu: F32, q: AS) -> AS: 1892 match left q: 1893 case 0n r: 1894 r 1895 case 1n+l AS{a, +s}: 1896 gn_csq(l, (off + c : U32), c, cg, mu, csq(Nat.sub(cg, 1n), (off + c : U32), mu, s, Array.get(F32, a, (off + c : U32)))) 1897 1898 def gn_mr1(+hw: Nat, +off: U32, +c: U32, +cg: Nat, +nf: F32, +eps: F32, q: AS) -> ASM: 1899 match q: 1900 case AS{a, +s}: 1901 ln_mr2(nf, eps, (s / nf : F32), gn_csq(Nat.sub(hw, 1n), off, c, cg, (s / nf : F32), csq(Nat.sub(cg, 1n), off, (s / nf : F32), 0.0, Array.get(F32, a, off)))) 1902 1903 # a block's mean and 1 / sqrt(var + eps); off is its first cell, nf = hw·cg 1904 def gn_mr(+hw: Nat, +off: U32, +c: U32, +cg: Nat, +nf: F32, +eps: F32, x: Array) -> ASM: 1905 gn_mr1(hw, off, c, cg, nf, eps, gn_sum(Nat.sub(hw, 1n), off, c, cg, asum(Nat.sub(cg, 1n), off, 0.0, Array.get(F32, x, off)))) 1906 1907 # y over the block's rows, one segment each (LayerNorm's lnf_w with the group's gains) 1908 def gnf_rows(left: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +mu: F32, +r: F32, s: R4) -> R4: 1909 match left s: 1910 case 0n R4{x, g, b, y}: 1911 lnf_w(Nat.sub(cg, 1n), j, off, mu, r, y, Array.get(F32, x, (off + j : U32)), Array.get(F32, g, j), Array.get(F32, b, j)) 1912 case 1n+l R4{x, g, b, y}: 1913 gnf_rows(l, (off + c : U32), j, c, cg, mu, r, lnf_w(Nat.sub(cg, 1n), j, off, mu, r, y, Array.get(F32, x, (off + j : U32)), Array.get(F32, g, j), Array.get(F32, b, j))) 1914 1915 def gnf_grp2(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, g: Array, b: Array, y: Array, m: ASM) -> R4: 1916 match m: 1917 case ASM{x, +mu, +r}: 1918 gnf_rows(Nat.sub(hw, 1n), off, j, c, cg, mu, r, R4{x, g, b, y}) 1919 1920 # one (image, group) block: off is the image's first row's first cell, j the group's first channel 1921 def gnf_grp(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +nf: F32, +eps: F32, s: R4) -> R4: 1922 match s: 1923 case R4{x, g, b, y}: 1924 gnf_grp2(hw, off, j, c, cg, g, b, y, gn_mr(hw, (off + j : U32), c, cg, nf, eps, x)) 1925 1926 def gnf_groups(left: Nat, +off: U32, +j: U32, +hw: Nat, +c: U32, +cg: Nat, +nf: F32, +eps: F32, s: R4) -> R4: 1927 match left: 1928 case 0n: 1929 gnf_grp(hw, off, j, c, cg, nf, eps, s) 1930 case 1n+q: 1931 gnf_groups(q, off, (j + U32.from_nat(cg) : U32), hw, c, cg, nf, eps, gnf_grp(hw, off, j, c, cg, nf, eps, s)) 1932 1933 def gnf_images(left: Nat, +off: U32, +ng: Nat, +hw: Nat, +c: U32, +cg: Nat, +nf: F32, +eps: F32, s: R4) -> R4: 1934 match left: 1935 case 0n: 1936 gnf_groups(Nat.sub(ng, 1n), off, 0, hw, c, cg, nf, eps, s) 1937 case 1n+q: 1938 gnf_images(q, (off + (U32.from_nat(hw) * c : U32) : U32), ng, hw, c, cg, nf, eps, gnf_groups(Nat.sub(ng, 1n), off, 0, hw, c, cg, nf, eps, s)) 1939 1940 def gn_fwd(+n: Nat, +hw: Nat, +ng: Nat, +cg: Nat, +eps: F32, p: LN, x: TA.Mat) -> Fwd, TA.Mat, TA.Mat>: 1941 match p x: 1942 case LN{TA.Mat{g}, TA.Mat{b}} TA.Mat{xa}: 1943 ln_f2(Nat.mul(n, hw), Nat.mul(ng, cg), gnf_images(Nat.sub(n, 1n), 0, ng, hw, U32.from_nat(Nat.mul(ng, cg)), cg, F32.from_nat(Nat.mul(hw, cg)), eps, R4{xa, g, b, zeros_a(Nat.mul(Nat.mul(n, hw), Nat.mul(ng, cg)))})) 1944 1945 # the backward pass is LayerNorm's with the block in place of the row: the two means (of dxh and of dxh·xh) are taken 1946 # over the whole block, then each cell's dx, dg[ch] += dy·xh and db[ch] += dy, in order 1947 def gnb_sums(left: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +mu: F32, +r: F32, q: R3SS) -> R3SS: 1948 match left q: 1949 case 0n t: 1950 t 1951 case 1n+l R3SS{x, dy, g, +sa, +sb}: 1952 gnb_sums(l, (off + c : U32), j, c, cg, mu, r, lnb_sums(Nat.sub(cg, 1n), j, (off + c : U32), mu, r, sa, sb, Array.get(F32, x, ((off + c : U32) + j : U32)), Array.get(F32, dy, ((off + c : U32) + j : U32)), Array.get(F32, g, j))) 1953 1954 def gnb_rows(left: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +mu: F32, +r: F32, +ma: F32, +mb: F32, s: R4B) -> R4B: 1955 match left s: 1956 case 0n R4B{x, dy, g, dx, w}: 1957 lnb_w(Nat.sub(cg, 1n), j, off, mu, r, ma, mb, dx, w, Array.get(F32, x, (off + j : U32)), Array.get(F32, dy, (off + j : U32)), Array.get(F32, g, j)) 1958 case 1n+l R4B{x, dy, g, dx, w}: 1959 gnb_rows(l, (off + c : U32), j, c, cg, mu, r, ma, mb, lnb_w(Nat.sub(cg, 1n), j, off, mu, r, ma, mb, dx, w, Array.get(F32, x, (off + j : U32)), Array.get(F32, dy, (off + j : U32)), Array.get(F32, g, j))) 1960 1961 def gnb_grp3(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +nf: F32, +mu: F32, +r: F32, dx: Array, w: A2, q: R3SS) -> R4B: 1962 match q: 1963 case R3SS{x, dy, g, +sa, +sb}: 1964 gnb_rows(Nat.sub(hw, 1n), off, j, c, cg, mu, r, (sa / nf : F32), (sb / nf : F32), R4B{x, dy, g, dx, w}) 1965 1966 def gnb_grp2(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +nf: F32, dy: Array, g: Array, dx: Array, w: A2, m: ASM) -> R4B: 1967 match m: 1968 case ASM{x, +mu, +r}: 1969 gnb_grp3(hw, off, j, c, cg, nf, mu, r, dx, w, gnb_sums(Nat.sub(hw, 1n), off, j, c, cg, mu, r, lnb_sums(Nat.sub(cg, 1n), j, off, mu, r, 0.0, 0.0, Array.get(F32, x, (off + j : U32)), Array.get(F32, dy, (off + j : U32)), Array.get(F32, g, j)))) 1970 1971 def gnb_grp(+hw: Nat, +off: U32, +j: U32, +c: U32, +cg: Nat, +nf: F32, +eps: F32, s: R4B) -> R4B: 1972 match s: 1973 case R4B{x, dy, g, dx, w}: 1974 gnb_grp2(hw, off, j, c, cg, nf, dy, g, dx, w, gn_mr(hw, (off + j : U32), c, cg, nf, eps, x)) 1975 1976 def gnb_groups(left: Nat, +off: U32, +j: U32, +hw: Nat, +c: U32, +cg: Nat, +nf: F32, +eps: F32, s: R4B) -> R4B: 1977 match left: 1978 case 0n: 1979 gnb_grp(hw, off, j, c, cg, nf, eps, s) 1980 case 1n+q: 1981 gnb_groups(q, off, (j + U32.from_nat(cg) : U32), hw, c, cg, nf, eps, gnb_grp(hw, off, j, c, cg, nf, eps, s)) 1982 1983 def gnb_images(left: Nat, +off: U32, +ng: Nat, +hw: Nat, +c: U32, +cg: Nat, +nf: F32, +eps: F32, s: R4B) -> R4B: 1984 match left: 1985 case 0n: 1986 gnb_groups(Nat.sub(ng, 1n), off, 0, hw, c, cg, nf, eps, s) 1987 case 1n+q: 1988 gnb_images(q, (off + (U32.from_nat(hw) * c : U32) : U32), ng, hw, c, cg, nf, eps, gnb_groups(Nat.sub(ng, 1n), off, 0, hw, c, cg, nf, eps, s)) 1989 1990 def gn_bwd(+n: Nat, +hw: Nat, +ng: Nat, +cg: Nat, +eps: F32, p: LN, x: TA.Mat, dy: TA.Mat, gg: LN) -> Bwd, TA.Mat>: 1991 match p x dy: 1992 case LN{TA.Mat{g}, b} TA.Mat{xa} TA.Mat{da}: 1993 ln_b2(Nat.mul(ng, cg), Nat.mul(n, hw), b, gg, gnb_images(Nat.sub(n, 1n), 0, ng, hw, U32.from_nat(Nat.mul(ng, cg)), cg, F32.from_nat(Nat.mul(hw, cg)), eps, R4B{xa, da, g, zeros_a(Nat.mul(Nat.mul(n, hw), Nat.mul(ng, cg))), A2{zeros_a(Nat.mul(ng, cg)), zeros_a(Nat.mul(ng, cg))}})) 1994 1995 # GroupNorm over n images of hw pixels and ng groups of cg channels, channels last; its parameters are LayerNorm's 1996 def GroupNorm(+n: Nat, +hw: Nat, +ng: Nat, +cg: Nat, +eps: F32) -> Layer, TA.Mat, TA.Mat, TA.Mat>: 1997 Layer{p => x => gn_fwd(n, hw, ng, cg, eps, p, x), p => s => dy => g => gn_bwd(n, hw, ng, cg, eps, p, s, dy, g), op => p => g => st => ln_step(Nat.mul(ng, cg), op, p, g, st), u => LN{TA.Mat.zeros(1n, Nat.mul(ng, cg)), TA.Mat.zeros(1n, Nat.mul(ng, cg))}, p => ln_dump(Nat.mul(ng, cg), p), p => xs => off => ln_load(Nat.mul(ng, cg), p, xs, off), p => xs => off => ln_store(Nat.mul(ng, cg), p, xs, off), u => U32.from_nat(Nat.mul(2n, Nat.mul(ng, cg)))} 1998 1999 # ===================================================================================================== 2000 # Embedding: row i of the output is row ids[i] of the table (v x d); the gradient adds each row back Lines 1-2000 of 3707. More: https://jend.leeeet.dev/read.txt?path=bendtorch@1.3.0.0/main.bend&start_line=2001