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.14.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 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)} 845 846 # Owned layer composition for dimensions read at run time. Each child record is opened once; each closure goes to 847 # exactly one corresponding closure of the composite layer. The template-based Seq above keeps its existing laws. 848 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>: 849 match r: 850 case Fwd{p1, y, s1}: seq_f3(P1, P2, Z, S1, S2, p1, s1, fb(p2, y)) 851 852 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>: 853 match p: 854 case Two{p1, p2}: seq_at_f2(P1, P2, Y, Z, S1, S2, p2, fb, fa(p1, x)) 855 856 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>: 857 match r: 858 case Bwd{p2, dy, g2}: seq_b3(P1, P2, X, p2, g2, ba(p1, s1, dy, g1)) 859 860 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>: 861 match p s g: 862 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)) 863 864 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>: 865 match p g st: 866 case Two{p1, p2} Two{g1, g2} Two{s1, s2}: pair_s3(P1, P2, sa(op, p1, g1, s1), sb(op, p2, g2, s2)) 867 868 def seq_at_dump(-P1: Type, -P2: Type, da: P1 -> PL, db: P2 -> PL, p: Two) -> PL>: 869 match p: 870 case Two{p1, p2}: pl_pair(P1, P2, da(p1), db(p2)) 871 872 def seq_at_load2(-P1: Type, -P2: Type, lb: P2 -> Array -> U32 -> PA, p2: P2, r: PA) -> PA>: 873 match r: 874 case PA{p1, xs, off}: pa_second(P1, P2, p1, lb(p2, xs, off)) 875 876 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>: 877 match p: 878 case Two{p1, p2}: seq_at_load2(P1, P2, lb, p2, la(p1, xs, off)) 879 880 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>: 881 match p: 882 case Two{p1, p2}: seq_at_load2(P1, P2, lb, p2, la(p1, xs, off)) 883 884 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>: 885 match a b: 886 case Layer{fa, ba, sa, za, da, la, ta, na} Layer{fb, bb, sb, zb, db, lb, tb, nb}: 887 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)} 888 889 # ===================================================================================================== 890 # Losses: the mean loss over the batch, and its gradient with respect to the output 891 # ===================================================================================================== 892 893 type LossOut<-Y: Type> is Type: 894 LossOut{loss: F32, dy: Y} 895 896 # ===================================================================================================== 897 # Training: one step of any layer under any loss 898 # ===================================================================================================== 899 900 # Parameters, gradient (zero between steps), optimizer-specific buffers, Adam's 1 - b1^t and 1 - b2^t, last loss. 901 type Train<-P: Type> is Type: 902 Train{p: P, g: P, st: OptState

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

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

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

) -> Train

: 924 match r: 925 case Step{p, g, st}: 926 Train{p, g, st, p1, p2, l} 927 928 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

: 929 match r: 930 case Bwd{p, dx, g}: 931 +q1 = (p1 + (OptCfg.c1(cfg) * (1.0 - p1 : F32) : F32) : F32) 932 +q2 = (p2 + (OptCfg.c2(cfg) * (1.0 - p2 : F32) : F32) : F32) 933 ts5(P, l, q1, q2, Layer.step(P, X, Y, S, m, Opt.at(cfg, q1, q2), p, g, st)) 934 935 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

: 936 match r: 937 case LossOut{l, dy}: 938 ts4(~P, ~X, ~Y, ~S, ~m, cfg, st, p1, p2, l, Layer.bwd(P, X, Y, S, m, p, s, dy, g)) 939 940 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

: 941 match r: 942 case Fwd{p, y, s}: 943 ts3(~P, ~X, ~Y, ~S, ~m, cfg, p, s, g, st, p1, p2, loss(y, t)) 944 945 # forward, loss, backward and one optimizer step: the same code for every model 946 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

: 947 match tr: 948 case Train{p, g, st, p1, p2, l}: 949 ts2(~P, ~X, ~Y, ~S, ~T, ~m, ~loss, cfg, g, st, p1, p2, t, Layer.fwd(P, X, Y, S, m, p, x)) 950 951 # One owned model, including dimensions captured at run time. Its fwd, bwd and step closures each run once. 952 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

: 953 match r: 954 case Bwd{p, dx, g}: 955 +q1 = (p1 + (OptCfg.c1(cfg) * (1.0 - p1 : F32) : F32) : F32) 956 +q2 = (p2 + (OptCfg.c2(cfg) * (1.0 - p2 : F32) : F32) : F32) 957 ts5(P, l, q1, q2, sf(Opt.at(cfg, q1, q2), p, g, st)) 958 959 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

: 960 match r: 961 case LossOut{l, dy}: train_at_b(P, X, sf, cfg, st, p1, p2, l, bf(p, s, dy, g)) 962 963 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

: 964 match r: 965 case Fwd{p, y, s}: train_at_loss(P, X, Y, S, bf, sf, cfg, p, s, g, st, p1, p2, loss(y, t)) 966 967 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

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