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

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

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

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

) -> Train

: 918 match r: 919 case Step{p, g, st}: 920 Train{p, g, st, p1, p2, l} 921 922 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

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

: 930 match r: 931 case LossOut{l, dy}: 932 ts4(~P, ~X, ~Y, ~S, ~m, cfg, st, p1, p2, l, Layer.bwd(P, X, Y, S, m, p, s, dy, g)) 933 934 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

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

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

: 947 match r: 948 case Bwd{p, dx, g}: 949 +q1 = (p1 + (OptCfg.c1(cfg) * (1.0 - p1 : F32) : F32) : F32) 950 +q2 = (p2 + (OptCfg.c2(cfg) * (1.0 - p2 : F32) : F32) : F32) 951 ts5(P, l, q1, q2, sf(Opt.at(cfg, q1, q2), p, g, st)) 952 953 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

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

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

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