1 # Temporary loading-only exception: bend-ml-binary-read is foreign C/JavaScript. No file with laws imports this module. 2 # Replace it with native Bend when the language supports dense reads (Renan, 2026-10-17). 3 # The reader is bend-ml-binary-read@0.2.0.0, with 64-bit offsets, BF16 files and reads into a cell offset 4 import Base 5 import bend-ml-binary-read@0.2.0.0/main.bend as R 6 import bendtorch@1.2.0.0/main.bend as T 7 8 def require(ok: Bool, +message: String) -> IO(Unit): 9 match ok: 10 case True{}: IO.pure(Unit, Unit{}) 11 case False{}: IO.die(Unit, 1, message) 12 13 # ===================================================================================================== 14 # W64: a 64-bit byte count or offset as two U32s, hi * 2^32 + lo. Bend (2.0.39) has no 64-bit integer, and a file past 15 # 4 GiB needs one (SmolLM2-1.7B's F32 weights, a 360M AdamW checkpoint). Plain Bend: none of these reaches the reader. 16 # ===================================================================================================== 17 18 type W64 is Data: 19 W64{hi: U32, lo: U32} 20 21 def W64.of(x: U32) -> W64: 22 W64{0, x} 23 24 def W64.of_pair(p: U32 & U32) -> W64: 25 match p: 26 case Tuple{hi, lo}: W64{hi, lo} 27 28 def W64.bit(b: Bool) -> U32: 29 match b: 30 case True{}: 1 31 case False{}: 0 32 33 def W64.eq(a: W64, b: W64) -> Bool: 34 match a b: 35 case W64{ah, al} W64{bh, bl}: Bool.and(U32.is_eq(ah, bh), U32.is_eq(al, bl)) 36 37 def W64.add_go(+ah: U32, +bh: U32, +s: U32, +al: U32) -> W64: 38 W64{((ah + bh : U32) + W64.bit(U32.is_lt(s, al)) : U32), s} 39 40 # a + b, wrapping at 2^64: the low words' sum carries one when it wrapped (it is then below either term) 41 def W64.add(a: W64, b: W64) -> W64: 42 match a b: 43 case W64{ah, +al} W64{bh, bl}: W64.add_go(ah, bh, (al + bl : U32), al) 44 45 def W64.mul_go(+h: U32, +k: U32, +p0: U32, +p1: U32) -> W64: 46 W64.add(W64{((h * k : U32) + U32.shrn(p1, 16n) : U32), p0}, W64{0, U32.shln(p1, 16n)}) 47 48 # a * k for k < 2^16 (enough for bytes a number and decimal digits), wrapping at 2^64: lo = l1 2^16 + l0, so 49 # lo * k = (l1 k) 2^16 + l0 k, where both products fit a U32 50 def W64.mul_small(a: W64, +k: U32) -> W64: 51 match a: 52 case W64{h, +l}: W64.mul_go(h, k, (U32.and(l, 65535) * k : U32), (U32.shrn(l, 16n) * k : U32)) 53 54 def W64.div10_go(+qh: U32, +r1: U32, +q1: U32, +l: U32) -> W64 & U32: 55 +t = (U32.shln(r1, 16n) + U32.and(l, 65535) : U32) 56 Tuple{W64{qh, U32.or(U32.shln(q1, 16n), U32.div(t, 10))}, U32.mod(t, 10)} 57 58 # (a / 10, a % 10), by long division in 16-bit steps so no partial remainder overflows a U32 59 def W64.div10(a: W64) -> W64 & U32: 60 match a: 61 case W64{+h, +l}: 62 +t = (U32.shln(U32.mod(h, 10), 16n) + U32.shrn(l, 16n) : U32) 63 W64.div10_go(U32.div(h, 10), U32.mod(t, 10), U32.div(t, 10), l) 64 65 def W64.is_zero(a: W64) -> Bool: 66 match a: 67 case W64{h, l}: Bool.and(U32.is_zero(h), U32.is_zero(l)) 68 69 def W64.show_go(f: Nat, z: Bool, qr: W64 & U32, acc: String) -> String: 70 match f z qr: 71 case _ True{} Tuple{q, r}: acc 72 case 0n False{} Tuple{q, r}: acc 73 case 1n+g False{} Tuple{+q, r}: W64.show_go(g, W64.is_zero(q), W64.div10(q), U32.show(r) ++ acc) 74 75 def W64.show_if(z: Bool, +a: W64) -> String: 76 match z: 77 case True{}: "0" 78 case False{}: W64.show_go(20n, False{}, W64.div10(a), "") 79 80 # the decimal digits, as U32.show writes them 81 def W64.show(+a: W64) -> String: 82 W64.show_if(W64.is_zero(a), a) 83 84 def W64.digits(s: String, f: Nat) -> Bool: 85 match s f: 86 case SNil{} _: True{} 87 case SCon{c, t} 0n: False{} 88 case SCon{c, t} 1n+g: Bool.and(Char.is_digit(c), W64.digits(t, g)) 89 90 def W64.read_go(s: String, +acc: W64) -> W64: 91 match s: 92 case SNil{}: acc 93 case SCon{Chr{x}, t}: W64.read_go(t, W64.add(W64.mul_small(acc, 10), W64.of(U32.sub(x, 48)))) 94 95 def W64.read_if(ok: Bool, s: String) -> Maybe<&2, W64>: 96 match ok: 97 case True{}: Some{W64.read_go(s, W64{0, 0})} 98 case False{}: None{} 99 100 # A decimal number of 1 to 19 digits (every such number fits 64 bits); anything else is None 101 def W64.read(s: String) -> Maybe<&2, W64>: 102 match s: 103 case SNil{}: None{} 104 case SCon{+h, +t}: W64.read_if(W64.digits(SCon{h, t}, 19n), SCon{h, t}) 105 106 # ===================================================================================================== 107 # Reading 108 # ===================================================================================================== 109 110 def read_done(+path: String, +n: U32, r: File & (Array & Result<&1, &1, U32 & String, U32>)) -> IO(Array): 111 match r: 112 case Tuple{f, Tuple{a, res}}: 113 do IO>: 114 +got : U32 <- IO.pass(U32, res) 115 File.close(f) 116 require(U32.is_eq(got, n), path ++ ": short read (" ++ U32.show(got) ++ " of " ++ U32.show(n) ++ " numbers)") 117 return a 118 119 def read_go(+path: String, +n: U32, f: File, +size: W64) -> IO(Array): 120 do IO>: 121 require(W64.eq(size, W64.mul_small(W64.of(n), 4)), path ++ ": expected " ++ W64.show(W64.mul_small(W64.of(n), 4)) ++ " bytes, got " ++ W64.show(size)) 122 rr : File & (Array & Result<&1, &1, U32 & String, U32>) <- R.read_f32_at(f, 0, 0, n, Array.new(F32, T.depth(U32.to_nat(n)), 0.0)) 123 read_done(path, n, rr) 124 125 def read_sized(+path: String, +n: U32, r: File & Result<&1, &1, U32 & String, U32 & U32>) -> IO(Array): 126 match r: 127 case Tuple{f, res}: 128 do IO>: 129 sz : U32 & U32 <- IO.pass(U32 & U32, res) 130 read_go(path, n, f, W64.of_pair(sz)) 131 132 # A raw, exact-size little-endian F32 file, read once into one source array: up to 2^32 - 1 numbers (16 GiB). 133 def read(+path: String, +n: U32) -> IO(Array): 134 do IO>: 135 f : File <- IO.try(File, File.open(path, "r")) 136 rr : File & Result<&1, &1, U32 & String, U32 & U32> <- R.size64(f) 137 read_sized(path, n, rr) 138 139 # n little-endian F32 numbers at byte `off` of a file of any size (a tensor inside a larger file) 140 def read_f32_at(+path: String, off: W64, +n: U32) -> IO(Array): 141 match off: 142 case W64{hi, lo}: 143 do IO>: 144 f : File <- IO.try(File, File.open(path, "r")) 145 rr : File & (Array & Result<&1, &1, U32 & String, U32>) <- R.read_f32_at(f, hi, lo, n, Array.new(F32, T.depth(U32.to_nat(n)), 0.0)) 146 read_done(path, n, rr) 147 148 # n little-endian BF16 numbers at byte `off` of a file of any size, each widened exactly to F32 (bf16.to(torch.float32)) 149 def read_bf16_at(+path: String, off: W64, +n: U32) -> IO(Array): 150 match off: 151 case W64{hi, lo}: 152 do IO>: 153 f : File <- IO.try(File, File.open(path, "r")) 154 rr : File & (Array & Result<&1, &1, U32 & String, U32>) <- R.read_bf16_at(f, hi, lo, n, Array.new(F32, T.depth(U32.to_nat(n)), 0.0)) 155 read_done(path, n, rr) 156 157 def size_done(r: File & Result<&1, &1, U32 & String, U32 & U32>) -> IO(W64): 158 match r: 159 case Tuple{f, res}: 160 do IO: 161 sz : U32 & U32 <- IO.pass(U32 & U32, res) 162 File.close(f) 163 return W64.of_pair(sz) 164 165 # A file's size in bytes, at any size 166 def size(+path: String) -> IO(W64): 167 do IO: 168 f : File <- IO.try(File, File.open(path, "r")) 169 rr : File & Result<&1, &1, U32 & String, U32 & U32> <- R.size64(f) 170 size_done(rr) 171 172 def weights_done(-P: Type, +n: U32, r: T.PA

) -> IO(P): 173 match r: 174 case T.PA{p, xs, +off}: 175 do IO

: 176 require(U32.is_eq(off, n), "Layer.count differs from the loading cursor") 177 return p 178 179 # All parameters in Layer.dump order. The source array is released after the model's arrays are filled. 180 def weights(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: T.Layer, +path: String, p: P) -> IO(P): 181 do IO

: 182 +n : U32 = T.Layer.count(P, X, Y, S, m) 183 a : Array <- read(path, n) 184 weights_done(P, n, T.Layer.load(P, X, Y, S, m, p, a, 0)) 185 186 # The checkpoint layout is determined by the already initialized optimizer state, without a header: 187 # SGD: parameters; momentum: parameters, buffer; Adam/AdamW: parameters, m, v, 1 - b1^t, 1 - b2^t; AdamW with packed 188 # BF16 moments: parameters, the moments' tree (main.bend's pt_store), 1 - b1^t, 1 - b2^t, the whole file read. 189 def train_read(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: T.Layer, +path: String, +base: U32, +copies: U32, +tail: U32, tr: T.Train

) -> IO(T.Train

): 190 do IO>: 191 require(U32.is_le(base, U32.div(U32.sub(4294967295, tail), copies)), "checkpoint exceeds 2^32 - 1 numbers") 192 +n : U32 = ((copies * base : U32) + tail : U32) 193 a : Array <- read(path, n) 194 return T.Train.restore(~P, ~X, ~Y, ~S, ~m, tr, a) 195 196 def cells_done(r: File & Result<&1, &1, U32 & String, U32>) -> IO(U32): 197 match r: 198 case Tuple{f, res}: 199 do IO: 200 +size : U32 <- IO.pass(U32, res) 201 File.close(f) 202 return U32.div(size, 4) 203 204 # the number of F32 cells in a file (a packed-moment checkpoint carries its own sizes: its length is the file's) 205 def cells_of(+path: String) -> IO(U32): 206 do IO: 207 f : File <- IO.try(File, File.open(path, "r")) 208 rr : File & Result<&1, &1, U32 & String, U32> <- File.size(f) 209 cells_done(rr) 210 211 def train(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: T.Layer, +path: String, tr: T.Train

) -> IO(T.Train

): 212 match tr: 213 case T.Train{p, g, T.NoMom{}, p1, p2, loss}: 214 train_read(~P, ~X, ~Y, ~S, ~m, path, T.Layer.count(P, X, Y, S, m), 1, 0, T.Train{p, g, T.NoMom{}, p1, p2, loss}) 215 case T.Train{p, g, T.OneMom{mm}, p1, p2, loss}: 216 train_read(~P, ~X, ~Y, ~S, ~m, path, T.Layer.count(P, X, Y, S, m), 2, 0, T.Train{p, g, T.OneMom{mm}, p1, p2, loss}) 217 case T.Train{p, g, T.TwoMom{mm, vv}, p1, p2, loss}: 218 train_read(~P, ~X, ~Y, ~S, ~m, path, T.Layer.count(P, X, Y, S, m), 3, 2, T.Train{p, g, T.TwoMom{mm, vv}, p1, p2, loss}) 219 case T.Train{p, g, T.Packed{t}, p1, p2, loss}: 220 do IO>: 221 +n : U32 <- cells_of(path) 222 train_read(~P, ~X, ~Y, ~S, ~m, path, n, 1, 0, T.Train{p, g, T.Packed{t}, p1, p2, loss}) 223 224 # ===================================================================================================== 225 # Safetensors (Hugging Face): 8 bytes holding the header's length (little-endian), a JSON header naming each tensor's 226 # dtype, shape and byte range (data_offsets, counted from the end of the header), then the data. The header is read as 227 # bytes and scanned in plain Bend: a lexer into tokens, then one pass that keeps what a tensor's entry holds. It reads 228 # what the format writes; it is not a JSON validator: an escaped character is kept as itself (\uXXXX is not decoded), 229 # and only strings, whole numbers, objects and lists are understood (anything else marks the header malformed). 230 # ===================================================================================================== 231 232 # one tensor the model expects, in Layer.dump order: its name in the file and its shape (rows, columns, ...) 233 type Entry is Data: 234 Entry{name: String, shape: List<&2, U32>} 235 236 def St.entry_name(e: Entry) -> String: 237 match e: 238 case Entry{n, sh}: n 239 240 type Tok is Data: 241 TStr{s: String} 242 TNum{v: W64} 243 TPun{c: U32} 244 TBad{c: U32} 245 246 # the lexer's state between bytes: outside a token, inside a string (its characters reversed), just after a backslash 247 # in a string, inside a number 248 type Lex is Data: 249 LOut{} 250 LStr{acc: String} 251 LEsc{acc: String} 252 LNum{v: W64} 253 254 # a state and the tokens one byte completed, in order 255 type LexR is Data: 256 LexR{st: Lex, toks: List<&2, Tok>} 257 258 def St.digit(+b: U32) -> Bool: 259 U32.is_lt(U32.sub(b, 48), 10) 260 261 def St.space(+b: U32) -> Bool: 262 Bool.or(Bool.or(U32.is_eq(b, 32), U32.is_eq(b, 10)), Bool.or(U32.is_eq(b, 13), U32.is_eq(b, 9))) 263 264 # { } [ ] : , 265 def St.pun(+b: U32) -> Bool: 266 Bool.or(Bool.or(Bool.or(U32.is_eq(b, 123), U32.is_eq(b, 125)), Bool.or(U32.is_eq(b, 91), U32.is_eq(b, 93))), Bool.or(U32.is_eq(b, 58), U32.is_eq(b, 44))) 267 268 def St.out_go(quote: Bool, digit: Bool, space: Bool, pun: Bool, +b: U32) -> LexR: 269 match quote digit space pun: 270 case True{} _ _ _: LexR{LStr{""}, Nil{}} 271 case _ True{} _ _: LexR{LNum{W64.of(U32.sub(b, 48))}, Nil{}} 272 case _ _ True{} _: LexR{LOut{}, Nil{}} 273 case _ _ _ True{}: LexR{LOut{}, Con{TPun{b}, Nil{}}} 274 case _ _ _ _: LexR{LOut{}, Con{TBad{b}, Nil{}}} 275 276 def St.out(+b: U32) -> LexR: 277 St.out_go(U32.is_eq(b, 34), St.digit(b), St.space(b), St.pun(b), b) 278 279 def St.str_go(quote: Bool, slash: Bool, +b: U32, acc: String) -> LexR: 280 match quote slash: 281 case True{} _: LexR{LOut{}, Con{TStr{String.reverse(acc)}, Nil{}}} 282 case _ True{}: LexR{LEsc{acc}, Nil{}} 283 case _ _: LexR{LStr{SCon{Chr{b}, acc}}, Nil{}} 284 285 # a number ends at the first byte that is not a digit, and that byte is then read as outside a token 286 def St.num_end(+v: W64, r: LexR) -> LexR: 287 match r: 288 case LexR{st, ts}: LexR{st, Con{TNum{v}, ts}} 289 290 def St.num_go(digit: Bool, +b: U32, +v: W64) -> LexR: 291 match digit: 292 case True{}: LexR{LNum{W64.add(W64.mul_small(v, 10), W64.of(U32.sub(b, 48)))}, Nil{}} 293 case False{}: St.num_end(v, St.out(b)) 294 295 def St.step(st: Lex, +b: U32) -> LexR: 296 match st: 297 case LOut{}: St.out(b) 298 case LStr{acc}: St.str_go(U32.is_eq(b, 34), U32.is_eq(b, 92), b, acc) 299 case LEsc{acc}: LexR{LStr{SCon{Chr{b}, acc}}, Nil{}} 300 case LNum{+v}: St.num_go(St.digit(b), b, v) 301 302 # ts onto out, which holds the tokens so far in reverse 303 def St.push(ts: List<&2, Tok>, out: List<&2, Tok>) -> List<&2, Tok>: 304 match ts: 305 case Nil{}: out 306 case Con{t, r}: St.push(r, Con{t, out}) 307 308 # the end of the header: a number ends there, an unfinished string is malformed 309 def St.fin(st: Lex, out: List<&2, Tok>) -> List<&2, Tok>: 310 match st: 311 case LOut{}: out 312 case LNum{v}: Con{TNum{v}, out} 313 case LStr{acc}: Con{TBad{0}, out} 314 case LEsc{acc}: Con{TBad{0}, out} 315 316 # the header's tokens, in order; r is what the previous byte left 317 def St.lex(bs: List<&2, U32>, r: LexR, out: List<&2, Tok>) -> List<&2, Tok>: 318 match bs r: 319 case Nil{} LexR{st, ts}: List.reverse(&2, Tok, St.fin(st, St.push(ts, out))) 320 case Con{+b, t} LexR{st, ts}: St.lex(t, St.step(st, b), St.push(ts, out)) 321 322 # one top-level entry of the header: a tensor (or __metadata__, which no layout names) 323 type Rec is Data: 324 Rec{name: String, dtype: String, shape: List<&2, W64>, offs: List<&2, W64>} 325 326 # the pass over the tokens: the depth ({ and [ open a level), whether the last token was a colon, the last string, 327 # the current entry's name and field, what the entry holds so far, the numbers of the open list, the entries read 328 # (last first), and the count of tokens the header may not hold 329 type Ps is Data: 330 Ps{depth: U32, colon: Bool, pend: String, name: String, field: String, dtype: String, shape: List<&2, W64>, offs: List<&2, W64>, nums: List<&2, W64>, recs: List<&2, Rec>, bad: U32} 331 332 # a string is a value after a colon (the dtype, at depth 2) or else a key, kept until its colon 333 def Sp.str(colon: Bool, is_dt: Bool, ps: Ps, s: String) -> Ps: 334 match colon is_dt ps: 335 case True{} True{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, False{}, p, n, f, s, sh, of, ns, rs, bad} 336 case True{} False{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, False{}, p, n, f, dt, sh, of, ns, rs, bad} 337 case False{} _ Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, False{}, s, n, f, dt, sh, of, ns, rs, bad} 338 339 # a colon makes the last string the entry's name (depth 1) or the field's (depth 2) 340 def Sp.colon(d1: Bool, d2: Bool, ps: Ps) -> Ps: 341 match d1 d2 ps: 342 case True{} _ Ps{d, c, +p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, True{}, p, p, f, dt, sh, of, ns, rs, bad} 343 case _ True{} Ps{d, c, +p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, True{}, p, n, p, dt, sh, of, ns, rs, bad} 344 case _ _ Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, True{}, p, n, f, dt, sh, of, ns, rs, bad} 345 346 # { opens a level; at depth 2 an entry starts empty 347 def Sp.open(fresh: Bool, ps: Ps) -> Ps: 348 match fresh ps: 349 case True{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{(d + 1 : U32), False{}, p, n, f, "", Nil{}, Nil{}, ns, rs, bad} 350 case False{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{(d + 1 : U32), False{}, p, n, f, dt, sh, of, ns, rs, bad} 351 352 # } closes a level; closing depth 2 ends an entry 353 def Sp.close(emit: Bool, ps: Ps) -> Ps: 354 match emit ps: 355 case True{} Ps{d, c, p, +n, f, dt, sh, of, ns, rs, bad}: Ps{U32.sub(d, 1), False{}, p, n, f, "", Nil{}, Nil{}, ns, Con{Rec{n, dt, sh, of}, rs}, bad} 356 case False{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{U32.sub(d, 1), False{}, p, n, f, dt, sh, of, ns, rs, bad} 357 358 def Sp.lopen(ps: Ps) -> Ps: 359 match ps: 360 case Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{(d + 1 : U32), False{}, p, n, f, dt, sh, of, Nil{}, rs, bad} 361 362 # ] closes a list; at depth 3 it is the entry's shape or data_offsets 363 def Sp.lclose(is_sh: Bool, is_of: Bool, ps: Ps) -> Ps: 364 match is_sh is_of ps: 365 case True{} _ Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{U32.sub(d, 1), False{}, p, n, f, dt, List.reverse(&2, W64, ns), of, Nil{}, rs, bad} 366 case _ True{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{U32.sub(d, 1), False{}, p, n, f, dt, sh, List.reverse(&2, W64, ns), Nil{}, rs, bad} 367 case _ _ Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{U32.sub(d, 1), False{}, p, n, f, dt, sh, of, Nil{}, rs, bad} 368 369 def Sp.comma(ps: Ps) -> Ps: 370 match ps: 371 case Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, False{}, p, n, f, dt, sh, of, ns, rs, bad} 372 373 # a number counts inside a list at depth 3; elsewhere it is not part of a tensor's entry 374 def Sp.num(in_list: Bool, v: W64, ps: Ps) -> Ps: 375 match in_list ps: 376 case True{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, False{}, p, n, f, dt, sh, of, Con{v, ns}, rs, bad} 377 case False{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, False{}, p, n, f, dt, sh, of, ns, rs, bad} 378 379 def Sp.pun(colon: Bool, ob: Bool, cb: Bool, ol: Bool, cl: Bool, +d: U32, +f: String, ps: Ps) -> Ps: 380 match colon ob cb ol cl: 381 case True{} _ _ _ _: Sp.colon(U32.is_eq(d, 1), U32.is_eq(d, 2), ps) 382 case _ True{} _ _ _: Sp.open(U32.is_eq(d, 1), ps) 383 case _ _ True{} _ _: Sp.close(U32.is_eq(d, 2), ps) 384 case _ _ _ True{} _: Sp.lopen(ps) 385 case _ _ _ _ True{}: Sp.lclose(Bool.and(U32.is_eq(d, 3), String.eq(f, "shape")), Bool.and(U32.is_eq(d, 3), String.eq(f, "data_offsets")), ps) 386 case _ _ _ _ _: Sp.comma(ps) 387 388 def Sp.step(t: Tok, ps: Ps) -> Ps: 389 match t ps: 390 case TStr{s} Ps{+d, +c, p, n, +f, dt, sh, of, ns, rs, bad}: 391 Sp.str(c, Bool.and(U32.is_eq(d, 2), String.eq(f, "dtype")), Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}, s) 392 case TNum{v} Ps{+d, c, p, n, f, dt, sh, of, ns, rs, bad}: 393 Sp.num(U32.is_eq(d, 3), v, Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}) 394 case TPun{+k} Ps{+d, c, p, n, +f, dt, sh, of, ns, rs, bad}: 395 Sp.pun(U32.is_eq(k, 58), U32.is_eq(k, 123), U32.is_eq(k, 125), U32.is_eq(k, 91), U32.is_eq(k, 93), d, f, Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}) 396 case TBad{k} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: 397 Ps{d, c, p, n, f, dt, sh, of, ns, rs, (bad + 1 : U32)} 398 399 def Sp.parse(ts: List<&2, Tok>, ps: Ps) -> Ps: 400 match ts: 401 case Nil{}: ps 402 case Con{t, r}: Sp.parse(r, Sp.step(t, ps)) 403 404 def St.rec_name(r: Rec) -> String: 405 match r: 406 case Rec{n, dt, sh, of}: n 407 408 def St.find_put(r: Rec, rest: Maybe<&2, Rec>, hit: Bool) -> Maybe<&2, Rec>: 409 match hit: 410 case True{}: Some{r} 411 case False{}: rest 412 413 # the entry named `name` 414 def St.find(rs: List<&2, Rec>, +name: String) -> Maybe<&2, Rec>: 415 match rs: 416 case Nil{}: None{} 417 case Con{+r, t}: St.find_put(r, St.find(t, name), String.eq(St.rec_name(r), name)) 418 419 def St.prod(xs: List<&2, U32>) -> U32: 420 match xs: 421 case Nil{}: 1 422 case Con{x, t}: (x * St.prod(t) : U32) 423 424 def St.total(es: List<&2, Entry>) -> U32: 425 match es: 426 case Nil{}: 0 427 case Con{Entry{n, sh}, t}: (St.prod(sh) + St.total(t) : U32) 428 429 def St.shape_eq(a: List<&2, W64>, b: List<&2, U32>) -> Bool: 430 match a b: 431 case Nil{} Nil{}: True{} 432 case Con{x, s} Con{y, t}: Bool.and(W64.eq(x, W64.of(y)), St.shape_eq(s, t)) 433 case _ _: False{} 434 435 def St.show_u(xs: List<&2, U32>, sep: String) -> String: 436 match xs: 437 case Nil{}: "" 438 case Con{x, t}: sep ++ U32.show(x) ++ St.show_u(t, ", ") 439 440 def St.show_w(xs: List<&2, W64>, sep: String) -> String: 441 match xs: 442 case Nil{}: "" 443 case Con{x, t}: sep ++ W64.show(x) ++ St.show_w(t, ", ") 444 445 # where a tensor is and how to read it, or why it cannot be read 446 type Plan is Data: 447 PlanOk{name: String, bf16: Bool, off: W64, n: U32} 448 PlanBad{msg: String} 449 450 def St.plan_span(ok: Bool, +name: String, bf16: Bool, off: W64, +n: U32) -> Plan: 451 match ok: 452 case True{}: PlanOk{name, bf16, off, n} 453 case False{}: PlanBad{"tensor " ++ name ++ ": data_offsets do not span its " ++ U32.show(n) ++ " numbers"} 454 455 # data_offsets must be [begin, end] with end - begin the tensor's bytes 456 def St.plan_of(+name: String, bf16: Bool, +start: W64, +n: U32, +width: U32, of: List<&2, W64>) -> Plan: 457 match of: 458 case Con{+b, Con{e, Nil{}}}: 459 St.plan_span(W64.eq(W64.add(b, W64.mul_small(W64.of(n), width)), e), name, bf16, W64.add(start, b), n) 460 case _: 461 PlanBad{"tensor " ++ name ++ ": data_offsets is not [begin, end]"} 462 463 def St.plan_shape(ok: Bool, +name: String, f32: Bool, +start: W64, +sh: List<&2, W64>, +want: List<&2, U32>, of: List<&2, W64>) -> Plan: 464 match ok f32: 465 case True{} True{}: St.plan_of(name, False{}, start, St.prod(want), 4, of) 466 case True{} False{}: St.plan_of(name, True{}, start, St.prod(want), 2, of) 467 case False{} _: PlanBad{"tensor " ++ name ++ ": shape [" ++ St.show_w(sh, "") ++ "] in the file, [" ++ St.show_u(want, "") ++ "] expected"} 468 469 def St.plan_dtype(f32: Bool, bf16: Bool, +name: String, +dt: String, +start: W64, +sh: List<&2, W64>, +want: List<&2, U32>, of: List<&2, W64>) -> Plan: 470 match f32 bf16: 471 case True{} _: St.plan_shape(St.shape_eq(sh, want), name, True{}, start, sh, want, of) 472 case _ True{}: St.plan_shape(St.shape_eq(sh, want), name, False{}, start, sh, want, of) 473 case _ _: PlanBad{"tensor " ++ name ++ ": dtype " ++ dt ++ " (only F32 and BF16 are read)"} 474 475 def St.plan(+start: W64, e: Entry, mr: Maybe<&2, Rec>) -> Plan: 476 match e mr: 477 case Entry{+name, want} None{}: PlanBad{"tensor " ++ name ++ ": not in the file"} 478 case Entry{+name, want} Some{Rec{rn, +dt, sh, of}}: St.plan_dtype(String.eq(dt, "F32"), String.eq(dt, "BF16"), name, dt, start, sh, want, of) 479 480 def St.read(bf16: Bool, f: File, off: W64, +n: U32, a: Array, +at: U32) -> IO(File & (Array & Result<&1, &1, U32 & String, U32>)): 481 match bf16 off: 482 case True{} W64{hi, lo}: R.read_bf16_into(f, hi, lo, n, a, at) 483 case False{} W64{hi, lo}: R.read_f32_into(f, hi, lo, n, a, at) 484 485 def St.got(+path: String, +name: String, +n: U32, r: File & (Array & Result<&1, &1, U32 & String, U32>), k: File -> Array -> U32 -> IO(Array)) -> IO(Array): 486 match r: 487 case Tuple{f, Tuple{a, res}}: 488 do IO>: 489 +got : U32 <- IO.pass(U32, res) 490 require(U32.is_eq(got, n), path ++ ": tensor " ++ name ++ ": short read (" ++ U32.show(got) ++ " of " ++ U32.show(n) ++ " numbers)") 491 k(f, a, n) 492 493 # reads one planned tensor into cells at .., then goes on with k 494 def St.then(+path: String, p: Plan, f: File, a: Array, +at: U32, k: File -> Array -> U32 -> IO(Array)) -> IO(Array): 495 match p: 496 case PlanBad{msg}: IO.die(Array, 1, path ++ ": " ++ msg) 497 case PlanOk{+name, bf16, off, +n}: 498 IO.bind(File & (Array & Result<&1, &1, U32 & String, U32>), Array, St.read(bf16, f, off, n, a, at), r => St.got(path, name, n, r, k)) 499 500 # every tensor of es, in order, into consecutive cells from at 501 def St.fill(es: List<&2, Entry>, +path: String, +start: W64, +recs: List<&2, Rec>, f: File, a: Array, +at: U32) -> IO(Array): 502 match es: 503 case Nil{}: 504 do IO>: 505 File.close(f) 506 return a 507 case Con{+e, t}: 508 St.then(path, St.plan(start, e, St.find(recs, St.entry_name(e))), f, a, at, f2 => a2 => +n => St.fill(t, path, start, recs, f2, a2, (at + n : U32))) 509 510 def St.le32(+a: U32, +b: U32, +c: U32, +d: U32) -> U32: 511 U32.or(U32.or(a, U32.shln(b, 8n)), U32.or(U32.shln(c, 16n), U32.shln(d, 24n))) 512 513 # the header's length: the file's first 8 bytes, little-endian 514 def St.le64(bs: List<&2, U32>) -> Maybe<&2, W64>: 515 match bs: 516 case Con{b0, Con{b1, Con{b2, Con{b3, Con{b4, Con{b5, Con{b6, Con{b7, t}}}}}}}}: 517 Some{W64{St.le32(b4, b5, b6, b7), St.le32(b0, b1, b2, b3)}} 518 case _: 519 None{} 520 521 def St.count(bs: List<&2, U32>, +acc: U32) -> U32: 522 match bs: 523 case Nil{}: acc 524 case Con{b, t}: St.count(t, (acc + 1 : U32)) 525 526 # A parsed safetensors header: where the data starts and each tensor's record. Reading it once lets a model be filled 527 # one piece at a time (`piece`), so that no array of the whole model exists. One array for all of SmolLM2-1.7B's 528 # 1.71e9 parameters takes a 2^31-cell block (8 GiB) whose memory the training step's smaller blocks never reuse. 529 type Header is Data: 530 Header{start: W64, recs: List<&2, Rec>} 531 532 def St.parsed(+path: String, +start: W64, ps: Ps) -> IO(Header): 533 match ps: 534 case Ps{d, c, p, nm, fl, dt, sh, of, ns, rs, bad}: 535 do IO

: 536 require(Bool.and(U32.is_zero(d), U32.is_zero(bad)), path ++ ": malformed safetensors header") 537 return Header{start, rs} 538 539 def St.header(+path: String, +len: U32, r: File & Result<&1, &1, U32 & String, List<&2, U32>>) -> IO(Header): 540 match r: 541 case Tuple{f, res}: 542 do IO
: 543 +bs : List<&2, U32> <- IO.pass(List<&2, U32>, res) 544 File.close(f) 545 require(U32.is_eq(St.count(bs, 0), len), path ++ ": the header is cut short") 546 St.parsed(path, W64.add(W64.of(8), W64.of(len)), Sp.parse(St.lex(bs, LexR{LOut{}, Nil{}}, Nil{}), Ps{0, False{}, "", "", "", "", Nil{}, Nil{}, Nil{}, Nil{}, 0})) 547 548 # the format caps the header at 100 MB 549 def St.length(f: File, +path: String, h: Maybe<&2, W64>) -> IO(Header): 550 match h: 551 case None{}: IO.die(Header, 1, path ++ ": shorter than a safetensors header") 552 case Some{W64{hi, +lo}}: 553 do IO
: 554 require(Bool.and(U32.is_zero(hi), U32.is_le(lo, 100000000)), path ++ ": a safetensors header past 100 MB") 555 r : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, 8, lo) 556 St.header(path, lo, r) 557 558 def St.first8(+path: String, r: File & Result<&1, &1, U32 & String, List<&2, U32>>) -> IO(Header): 559 match r: 560 case Tuple{f, res}: 561 do IO
: 562 bs : List<&2, U32> <- IO.pass(List<&2, U32>, res) 563 St.length(f, path, St.le64(bs)) 564 565 # A safetensors file's header, parsed; the file is closed after it 566 def header(+path: String) -> IO(Header): 567 do IO
: 568 f : File <- IO.try(File, File.open(path, "r")) 569 r : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, 0, 8) 570 St.first8(path, r) 571 572 # The tensors `layout` names, read from the file whose header is h, in the layout's order, into one array of n numbers 573 def tensors_from(+path: String, h: Header, +n: U32, +layout: List<&2, Entry>) -> IO(Array): 574 match h: 575 case Header{start, recs}: 576 do IO>: 577 require(U32.is_eq(St.total(layout), n), path ++ ": the layout holds " ++ U32.show(St.total(layout)) ++ " numbers, the model " ++ U32.show(n)) 578 f : File <- IO.try(File, File.open(path, "r")) 579 St.fill(layout, path, start, recs, f, Array.new(F32, T.depth(U32.to_nat(n)), 0.0), 0) 580 581 # The tensors `layout` names, read from a safetensors file in its order into one array of n numbers (F32 as they are, 582 # BF16 widened exactly). Each must be in the file with the shape given, as F32 or BF16, and the shapes must hold n 583 # numbers in all; otherwise the program stops with a message naming the tensor. Other tensors in the file are ignored. 584 def tensors(+path: String, +n: U32, +layout: List<&2, Entry>) -> IO(Array): 585 do IO>: 586 require(U32.is_eq(St.total(layout), n), path ++ ": the layout holds " ++ U32.show(St.total(layout)) ++ " numbers, the model " ++ U32.show(n)) 587 h : Header <- header(path) 588 tensors_from(path, h, n, layout) 589 590 # One layer's parameters from the file whose header is h: `layout` names its tensors in Layer.dump order, and 591 # Layer.load fills them; the array they were read into is gone once this answers 592 def piece(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: T.Layer, +path: String, h: Header, +layout: List<&2, Entry>, p: P) -> IO(P): 593 do IO

: 594 +n : U32 = T.Layer.count(P, X, Y, S, m) 595 a : Array <- tensors_from(path, h, n, layout) 596 weights_done(P, n, T.Layer.load(P, X, Y, S, m, p, a, 0)) 597 598 # A pair of parameters filled one after the other: the first piece's source array is released before the second is read 599 def two(-P1: Type, -P2: Type, p: T.Two, k1: P1 -> IO(P1), k2: P2 -> IO(P2)) -> IO(T.Two): 600 match p: 601 case T.Two{a, b}: 602 do IO>: 603 x : P1 <- k1(a) 604 y : P2 <- k2(b) 605 return T.Two{x, y} 606 607 # All parameters from a safetensors file: `layout` names the model's tensors in Layer.dump order (a tied output head 608 # is not in it, as SmolLM2's is not in its file), and Layer.load fills the parameters from them as `weights` does 609 def safetensors(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: T.Layer, +path: String, +layout: List<&2, Entry>, p: P) -> IO(P): 610 do IO

: 611 h : Header <- header(path) 612 piece(~P, ~X, ~Y, ~S, ~m, path, h, layout, p)