From b17ef42636ca5760964417d0d29c88a64c0e891b Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 20 Aug 2026 10:08:38 -0500 Subject: [PATCH 1/5] deps: refresh the lockfile and take hf-chat-template to 1.0.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Dependency-freshness increment. `cargo update` moves 36 transitive packages (minijinja 2.23->2.24, the icu_* 2.2->2.3 family, futures 0.3.33->0.3.34, h2, quinn-proto, rustls-webpki, zerovec, ...). The one major available to us is hf-chat-template 0.2.1 -> 1.0.0, and it is taken. That crate is load-bearing in two roles at once, so it gets the gate that matters rather than a build check: it is the reference renderer behind the template BYTE gate AND, behind the non-default `jinja-template` feature, the fallback renderer for checkpoints with no hardcoded family renderer. Proof it is a non-event: the byte gate is 10/10 byte-identical across Qwen3, Qwen2 and LFM2 (plain / multi-turn / tools / tool-history / no-system), i.e. every rendered prompt is the same bytes under 1.0.0 as under 0.2.1. Suites green both ways the crate can be compiled in — 216 tests default, 223 with `jinja-template` on, clippy --all-targets clean on both. Two incompatible upgrades are deliberately NOT taken, and the roadmap now records why: - burn stays 0.21 — crates.io serves 0.22.0-pre.2 and the P0 migration item says do not adopt a pre-release. - wgpu stays 29 — `cargo tree -i wgpu` shows a single wgpu (29.0.4) reached through cubecl-wgpu 0.10, so bumping only our direct handle would put a second, non-Burn wgpu in the tree and the startup adapter probe would stop describing the device Burn actually runs on. It unblocks with the burn bump. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 151 +++++++++++++++++++++++++++-------------------------- Cargo.toml | 2 +- ROADMAP.md | 7 +++ 3 files changed, 84 insertions(+), 76 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4ca1880..8cb08ea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -971,9 +971,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.4.2" +version = "1.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d262e149917187838d5b42777c8253bcb64500067342904e7d429499a6f277e" +checksum = "509591b7bcd67f4ef775afad7662703b4935daaa6ec0e5605cfb1090b32a2b6d" dependencies = [ "find-msvc-tools", "jobserver", @@ -1846,9 +1846,9 @@ dependencies = [ [[package]] name = "cudarc" -version = "0.19.8" +version = "0.19.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42310153e06cf4cd532901f7096beb27504d681736a29ee90728ae4e2d93b2a8" +checksum = "804764d10e844da09765a7b2ca9641a0851523d1702efb0d7299d73e31b86e80" dependencies = [ "libloading 0.9.0", ] @@ -2380,9 +2380,9 @@ dependencies = [ [[package]] name = "find-msvc-tools" -version = "0.1.10" +version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26b73573e6edcd2af0cdf47bd6cb58f0b3839491263c314eaad1ccf24430e1de" +checksum = "d45db016d36b838f563236e9193d0ee6ce38f3f68b6c94e914b4929c96bbb890" [[package]] name = "fixedbitset" @@ -2454,9 +2454,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.33" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" dependencies = [ "futures-core", "futures-sink", @@ -2464,15 +2464,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.33" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" [[package]] name = "futures-io" -version = "0.3.33" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4577ecaa3c4f96589d473f679a71b596316f6641bc350038b962a5daf0085d7a" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" [[package]] name = "futures-lite" @@ -2489,21 +2489,21 @@ dependencies = [ [[package]] name = "futures-sink" -version = "0.3.33" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e34418ac499d6305c2fb5ad0ed2f6ac998c5f8ca209b4510f7f94242c647e307" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" [[package]] name = "futures-task" -version = "0.3.33" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" [[package]] name = "futures-util" -version = "0.3.33" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" dependencies = [ "futures-core", "futures-io", @@ -2772,9 +2772,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.15" +version = "0.4.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" +checksum = "9f877e75f39e9827ec50a572dd592684ac28c029578726c85f1b2aa6ab807449" dependencies = [ "atomic-waker", "bytes", @@ -2862,9 +2862,9 @@ checksum = "dfa686283ad6dd069f105e5ab091b04c62850d3e4cf5d67debad1933f55023df" [[package]] name = "hf-chat-template" -version = "0.2.1" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "15e938b1f75afdc8eed9071ff130befb7d014f9b09dbd8d81877ade94ee9f9fe" +checksum = "ed4c5f2a8f34e7371a60a73bf129b13f698db5f9e0ba4cd583ff25bae285e166" dependencies = [ "minijinja", "minijinja-contrib", @@ -2912,9 +2912,9 @@ dependencies = [ [[package]] name = "http-body-util" -version = "0.1.4" +version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9f41fd6a08e4d4ec69df65976da761afd5ad5e58a9d4acb46bd1c953a9e3ff2" +checksum = "23169fe34a5fbcdd3f3862e78fb9b6fccd5f02a6dc6f732547005d45631ce71c" dependencies = [ "bytes", "futures-core", @@ -3000,9 +3000,9 @@ dependencies = [ [[package]] name = "icu_collections" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +checksum = "fa68d21081c4a05d5a901a1c62add574c77048b6a1c67be3b50ce0b60d4ca513" dependencies = [ "displaydoc", "potential_utf", @@ -3014,9 +3014,9 @@ dependencies = [ [[package]] name = "icu_locale_core" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +checksum = "d56e28588da92eee5c3201a6eff33fabdd49b62269c8938d4ff050ce4d900deb" dependencies = [ "displaydoc", "litemap", @@ -3027,9 +3027,9 @@ dependencies = [ [[package]] name = "icu_normalizer" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +checksum = "12f9cf5f235641ed274641dd81c3f28d870e276763d0797aeeab72317b1c646f" dependencies = [ "icu_collections", "icu_normalizer_data", @@ -3041,16 +3041,17 @@ dependencies = [ [[package]] name = "icu_normalizer_data" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" +checksum = "1563da1ed3e0b3bf3d74c9b85917ac9c56464d2f57242270c09c9e752f8021a0" [[package]] name = "icu_properties" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +checksum = "7e7ca276ad3145661a65914e6daf131ca5120cd3dcee8f8f3214b8875184a148" dependencies = [ + "displaydoc", "icu_collections", "icu_locale_core", "icu_properties_data", @@ -3061,15 +3062,15 @@ dependencies = [ [[package]] name = "icu_properties_data" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" +checksum = "e590f038c1464a96894fd6d10127e90a8be4509f56ff7ecef851b15cee0b7caa" [[package]] name = "icu_provider" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +checksum = "92a7ed671a6aad807a8651a2e1782a6598fda9ce5185dd8158549e95a91c6428" dependencies = [ "displaydoc", "icu_locale_core", @@ -3357,9 +3358,9 @@ dependencies = [ [[package]] name = "liblzma" -version = "0.4.7" +version = "0.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45aec2360b3933207e27908049d8e4df4e476b58180afb1e56b2a4fb72efe4ba" +checksum = "2fe0a34ca854fd4f20c07f696fc8675aec78f87d88d29f5e10257a7490a1b2e1" dependencies = [ "liblzma-sys", "num_cpus", @@ -3367,9 +3368,9 @@ dependencies = [ [[package]] name = "liblzma-sys" -version = "0.4.7" +version = "0.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a046c7f353ba30f810545151e04f63545833803f5b86ee3ddf1517247fe560a5" +checksum = "a0dad045e4b1b7b170be4b60b54b780cafb4490165461bac7d1cf7b703f61d5f" dependencies = [ "cc", "libc", @@ -3384,9 +3385,9 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libredox" -version = "0.1.19" +version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2026a5056764a10b2bf5d56488cba40da507f5493a6a429340e2004d9ed085fa" +checksum = "28d0a00925a9f930d679b6789b721e3a7f9ed110f41b86d2497caa780c3a070a" dependencies = [ "libc", ] @@ -3399,9 +3400,9 @@ checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" [[package]] name = "litemap" -version = "0.8.2" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" +checksum = "47d9d19d1d6efa0109d2f65ff4c85cddd50bd572e5a00127ab10987290bcefae" [[package]] name = "litrs" @@ -3542,9 +3543,9 @@ checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" [[package]] name = "minijinja" -version = "2.23.0" +version = "2.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42d74234349a775546a83af0f0c0c0e3a73227dee4950a542cda81a47240b3e6" +checksum = "86886cf6dbf4e614b19c9a1eec9775f021869d7eadde0fc73921a81b90c9b4c9" dependencies = [ "indexmap", "memo-map", @@ -3554,9 +3555,9 @@ dependencies = [ [[package]] name = "minijinja-contrib" -version = "2.23.0" +version = "2.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35c96d8fae7fa4743bbcf06f486ff43f81837f5e8826f8bf525dbc4c0e03be5d" +checksum = "bd3e5f077bc2379f0f7d911e7cfdd921114ed99fc884533dca502944cb355b11" dependencies = [ "minijinja", "serde", @@ -3841,9 +3842,9 @@ dependencies = [ [[package]] name = "num-integer" -version = "0.1.46" +version = "0.1.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b" dependencies = [ "num-traits", ] @@ -4141,9 +4142,9 @@ checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" [[package]] name = "pkg-config" -version = "0.3.33" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" +checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" [[package]] name = "plotters" @@ -4194,9 +4195,9 @@ checksum = "bc6355899e1c9462875b6757c79f3caa011a1fdae12bbb1a2e72dd1f234f8336" [[package]] name = "portable-atomic" -version = "1.14.0" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" dependencies = [ "serde", ] @@ -4212,9 +4213,9 @@ dependencies = [ [[package]] name = "potential_utf" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661" dependencies = [ "zerovec", ] @@ -4350,9 +4351,9 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.16" +version = "0.11.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" +checksum = "04759210543be93709136e28212294a659ef5001836ff4eab4d663e4529bba83" dependencies = [ "bytes", "getrandom 0.4.3", @@ -4838,9 +4839,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.103.13" +version = "0.103.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +checksum = "0527518605e68109d875e248ea259b6758801cf165e4b2c2733ae3b51f12535a" dependencies = [ "ring", "rustls-pki-types", @@ -5387,9 +5388,9 @@ dependencies = [ [[package]] name = "tinystr" -version = "0.8.3" +version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643" dependencies = [ "displaydoc", "zerovec", @@ -5942,9 +5943,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.24.0" +version = "1.24.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239" +checksum = "2cefc03fd367c0c6d4305de1b312cf00248c4114f4a0418ce6a6af769e3b0bd9" [[package]] name = "v_frame" @@ -6530,9 +6531,9 @@ checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" [[package]] name = "writeable" -version = "0.6.3" +version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" +checksum = "3ad82d2a33cdc9674dc7465672f271e096168fcdbe0f799d9e6db8c5892679dc" [[package]] name = "xattr" @@ -6634,9 +6635,9 @@ checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" [[package]] name = "zerotrie" -version = "0.2.4" +version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +checksum = "4ea269c3bd32f0a32c321907a2ae912ba6f4649bb0fc764a15627e99a7095a3f" dependencies = [ "displaydoc", "yoke", @@ -6645,9 +6646,9 @@ dependencies = [ [[package]] name = "zerovec" -version = "0.11.6" +version = "0.11.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +checksum = "bb0464e17806c1d976d5cba29399c7f08e516e279e2ba493f63123b5fca67dd8" dependencies = [ "yoke", "zerofrom", @@ -6656,13 +6657,13 @@ dependencies = [ [[package]] name = "zerovec-derive" -version = "0.11.3" +version = "0.11.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +checksum = "9f212a141d820099d57ffafb9569be9617a6f27d3dc881fbee8fb56642f917a9" dependencies = [ "proc-macro2", "quote", - "syn 2.0.119", + "syn 3.0.3", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index a46b3c3..cceab1b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -95,7 +95,7 @@ criterion = "0.8" # layer). Two roles, one version: the test-only reference behind the template # BYTE gate, and — behind mummu's non-default `jinja-template` feature — the # fallback renderer for checkpoints with no hardcoded family renderer. -hf-chat-template = "0.2.1" +hf-chat-template = "1.0.0" [profile.release] # Thin LTO: inference is GPU-kernel-bound, so fat LTO's extra minutes per link diff --git a/ROADMAP.md b/ROADMAP.md index 2079a74..a6c4082 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -381,6 +381,13 @@ a benchmark holds/improves its budget; README perf claims link an artifact. dequant→op→quant fallbacks for slice/gather/select/expand plus BitNet b1.58 calibration (P9). Plan the migration around measuring graph capture on the decode loop first — if it lands the dispatch win, it reorders everything below it in the perf section. + *(2026-08-20)* Still gated: crates.io now serves **0.22.0-pre.2**, so 0.21.0 remains the newest + stable and the do-not-adopt-a-pre-release rule holds for another run. Checked as part of the + dependency sweep, which is also why **wgpu 30 stays held**: `cargo upgrade --incompatible` + offers it, but `cargo tree -i wgpu` shows exactly one wgpu in the graph (29.0.4, reached via + `cubecl-wgpu 0.10`), so bumping our direct handle alone would put a second, non-Burn wgpu in + the tree and the startup adapter probe would stop describing the device Burn actually runs on. + wgpu 30 unblocks with the burn bump, exactly as the Stack note says — not before. - [x] Silence the pre-existing `LNK4098` (LIBCMT defaultlib conflict) the 2026-07 nightly toolchain's new `linker_messages` lint now surfaces when linking the `mummu` lib-test binary — find which native dep object embeds the static-CRT directive (tokenizers' C++ deps are the suspects) and From a6a6a3f15168272c16ee67205a2ffb7a179ace4a Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 20 Aug 2026 11:10:08 -0500 Subject: [PATCH 2/5] docs: fold this run's research into the roadmap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four dated notes, each attached to the item it changes rather than to a log. P2 MoE decode — the two findings that most change that item's shape. llama.cpp PR #25294 turns route (c) from a slogan into an implementation with numbers: bounded per-layer device-side expert-slab cache, CPU-side top-k id remap to cache slots, async demand-load on a miss, eviction by DECAYING ROUTE HOTNESS with an LRU tiebreak, O_DIRECT to keep the page cache from thrashing — 5.3x prefill / 2.4x decode over mmap+--n-cpu-moe at a 79 % hit rate. That hit rate is the load-bearing number for us: it says a small resident expert set covers most tokens. Separately, our own 2026-08-03 gather regression looks like a BATCH-SIZE REGIME rather than a refutation — llama.cpp only switches to the copy-experts path above a batch threshold, and ik_llama.cpp sets it at 32 * total_experts / active_experts, i.e. 256 tokens for OLMoE's 64/8. We measured batch-1 decode, the worst possible point for a materializing gather; prefill is where it should win, and that leg was never run. P6 NVMe streaming — the same PR is the reference implementation of the colibri-shaped design already written there, so the note records the two details worth stealing (hotness-with-LRU-tiebreak over plain LRU, since MoE routing is skewed; unbuffered reads once the model exceeds RAM) and its single-context limitation, which is a constraint on how we own the pool. P2 Qwen3.5 — the linear-attention half has a name, Gated DeltaNet, at roughly 75 % linear / 25 % full, so full_attention_interval is expected to read 4. Also: the MTP variants ship as their own GGUF repos, which is the concrete draft-model artifact P5's speculative-decoding item needs. P0 burn 0.22 — the associated element types are not deleted but moved onto a new BackendTypes trait, with callers steered to the type aliases. Actionable before the bump: every hand-written B::FloatElem is a place the 0.22 diff will land. Co-Authored-By: Claude Opus 4.8 --- ROADMAP.md | 47 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/ROADMAP.md b/ROADMAP.md index a6c4082..e904ab5 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -388,6 +388,14 @@ a benchmark holds/improves its budget; README perf claims link an artifact. `cubecl-wgpu 0.10`), so bumping our direct handle alone would put a second, non-Burn wgpu in the tree and the startup adapter probe would stop describing the device Burn actually runs on. wgpu 30 unblocks with the burn bump, exactly as the Stack note says — not before. + *(2026-08-20 research)* One refinement to the migration shape: the associated element types are + not simply deleted — they move off `Backend` onto a new **`BackendTypes`** trait, and the release + notes steer callers to the **type aliases** (`Device`, `FloatTensor`) instead of naming + associated types directly, specifically to dodge resolution problems. That is actionable now, at + zero migration risk: every place we write `B::FloatElem` / `B::IntElem` by hand (`backend:: + {float_dtype,int_dtype}` and each loader's `target_float` derivation) is a place the 0.22 diff will + land, so preferring the alias form where one already exists shrinks that diff before the bump. + — https://github.com/tracel-ai/burn/releases - [x] Silence the pre-existing `LNK4098` (LIBCMT defaultlib conflict) the 2026-07 nightly toolchain's new `linker_messages` lint now surfaces when linking the `mummu` lib-test binary — find which native dep object embeds the static-CRT directive (tokenizers' C++ deps are the suspects) and @@ -655,6 +663,25 @@ a benchmark holds/improves its budget; README perf claims link an artifact. VRAM (P9 keep-quantized, or expert offload); (c) llama.cpp's own answer, `--n-cpu-moe`-style placement, which sidesteps the gather entirely by moving whole expert banks rather than slicing them. *(2026-08-03, measured this run.)* + *(2026-08-20 research)* Two findings that sharpen the three routes above. (1) **Route (c) now has + a measured implementation to copy rather than a slogan**: llama.cpp PR #25294 streams routed + experts from disk behind a bounded per-layer device-side cache of expert *slabs* — top-k ids are + remapped to cache slots on the CPU, a miss demand-loads asynchronously through an io thread pool, + eviction is **decaying route hotness with an LRU tiebreak**, and reads use `O_DIRECT` so the page + cache cannot thrash against a model far larger than RAM. Reported on GB10 for a ~254 GB model at a + 90-slot cache: **5.3x prefill / 2.4x decode vs `mmap`+`--n-cpu-moe`, at a 79 % cache hit rate**. + That hit rate is the number that matters to us — it says a *small* resident expert set covers most + tokens, which is exactly the premise the gather route needs to become cheap. It also ships + **wave-partitioned prefill** (when a batch needs more experts than slots, run experts in waves of + `(n_slots - n_expert_used)/2` and sum the masked outputs) — the trick that keeps a bounded cache + from capping prompt length. (2) **Our gather regression is a batch-size regime, not a dead end**: + llama.cpp only switches to the copy-experts-to-GPU path above a batch threshold, and ik_llama.cpp + sets that threshold at `32 * total_experts / active_experts` — **256 tokens for OLMoE's 64/8**. + Batch-1 decode is the worst possible point for a materializing gather, and prefill is the regime + where it should win; the 2026-08-03 A/B measured only decode, so re-run it at prefill batch sizes + before treating the gather as refuted. Both are gated on `bench/BASELINE.md` like the rest — + https://github.com/ggml-org/llama.cpp/pull/25294 · + https://huggingface.co/blog/Doctor-Shotgun/llamacpp-moe-offload-guide - [ ] **OLMoE from HF safetensors** — the port loads GGUF only because HF stores each expert separately (`model.layers.N.mlp.experts.{0..63}.{gate,up,down}_proj.weight`) while `MoeExperts` holds one fused `[experts, out, in]` tensor per projection. Needs a concat-on-import step (64 slices → one tensor, @@ -671,6 +698,16 @@ a benchmark holds/improves its budget; README perf claims link an artifact. the 4B/9B are 2026's local-FC sweet spot (BFCL 9B 66.1%) and the 4B ships an `mmproj` vision projector (P11 candidate). *(2026-07-30 research)* — https://huggingface.co/unsloth/Qwen3.5-4B-GGUF + *(2026-08-20 research)* The mechanism has a name and a ratio, which makes the port scopeable: + the linear-attention half is **Gated DeltaNet**, interleaved with periodic full-attention at + roughly **75 % linear / 25 % full** — so `full_attention_interval` is expected to read 4, and the + recurrent-state cache (the LFM2 conv-state machinery generalizes) carries three quarters of the + layers. Two adjacent facts: the family ships **MTP** variants as their own GGUF repos + (`unsloth/Qwen3.5-{2B,9B,35B-A3B}-MTP-GGUF`), which is the concrete draft-model artifact the P5 + speculative-decoding item needs rather than a hypothetical; and llama.cpp renamed the flag + `--spec-type mtp` to `--spec-type draft-mtp` (2026-05-13), worth knowing before wiring a + reference leg against it. — + https://huggingface.co/unsloth/Qwen3.5-9B-MTP-GGUF · https://sebastianraschka.com/llm-architecture-gallery/hybrid-attention/ - [x] A `Model` trait so new architectures (Hermes-class function-callers, Gemma, Qwen3, …) slot in. *(2026-07-10) `models::CausalLm` — associated `Cache` type; a port supplies `new_cache` / `forward` / `is_eos` and inherits `generate` / `greedy_generate` / `first_token` from the shared @@ -1358,6 +1395,16 @@ that fits the model AND uses every device to the fullest. RAM/VRAM pool, residency keyed by routing frequency, prefetch keyed by the router's early output. Gate like everything: parity unaffected by placement (colibri's own invariant — "placement only affects speed, never precision"), and a bench proving the streamed model beats the + *(2026-08-20 research — read with the P2 MoE note, which carries the detail)* The colibri-shaped + design above now has an independent implementation to measure against: llama.cpp PR #25294 does + per-layer bounded expert-slab caching + async demand-load + hotness/LRU eviction + `O_DIRECT`, and + reports 5.3x prefill / 2.4x decode over `mmap`+`--n-cpu-moe` at a 79 % hit rate. Two design points + worth stealing outright when the `Disk` placement class is built: eviction on **decaying route + hotness with an LRU tiebreak** (not plain LRU — MoE routing is skewed, and plain LRU throws away + a hot expert after one cold burst), and **`O_DIRECT`/unbuffered reads**, because the OS page cache + actively hurts once the model exceeds RAM. Its stated limitation is also a design constraint for + us: single-context only — concurrent decodes sharing one streamed model corrupt each other, so the + residency pool has to be owned per-session or locked. — https://github.com/ggml-org/llama.cpp/pull/25294 largest-fitting resident one on task throughput. *(2026-07-30 research)* — https://github.com/JustVugg/colibri *(2026-08-03 research)* Prior art to mine when this is picked up: llama.cpp's **`--n-cpu-moe N`** From c2ccd49f3d693b5711e32573d1b58840cceec281 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 20 Aug 2026 11:11:21 -0500 Subject: [PATCH 3/5] hub: install the sibling files the import gates actually read MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `fetch_model_with` fetched exactly config.json, tokenizer.json, and the weights. So a checkpoint installed through Mummu's own registry arrived WITHOUT tokenizer_config.json — and every gate that reads it fails open: `validate_checkpoint_dir` returns Ok(None), so the EOS-agreement, added-token-id and tool-call-convention checks silently do not run, and the `tokenizer_config` the loaders surface is always None. The gates have been passing on hand-populated fixtures. Found by installing the OLMoE safetensors checkpoint through the registry and noticing the resulting dir has no tokenizer_config.json to check. Both tokenizer_config.json and chat_template.jinja are now fetched as OPTIONAL files. Three details that make this a fix rather than a hopeful extra request: - Only a 404 counts as absent. A 403 on a gated repo or a 5xx is treated as present so the real fetch raises it, instead of the probe swallowing a real failure as "the repo just doesn't ship this". - config.json and tokenizer.json stay REQUIRED — a 404 on either is still an error. - They are fetched BEFORE the weights, because the single-file branch returns early; after it, sharded checkpoints would have been the only ones to get them. Proof is a clean-dir install, so no hand-populated fixture can satisfy it: `a_registry_install_arrives_with_the_files_the_import_gates_read` fetches a catalog model into a fresh dir and asserts both that tokenizer_config.json is on disk and that validate_checkpoint_dir returns Some — the gate now has something to check rather than quietly passing. Green, plus a unit test pinning the optional set and its URLs (223 lib tests). Co-Authored-By: Claude Opus 4.8 --- ROADMAP.md | 17 ++++++++ crates/mummu/src/hub.rs | 78 +++++++++++++++++++++++++++++++++- crates/mummu/tests/real_hub.rs | 42 ++++++++++++++++++ 3 files changed, 135 insertions(+), 2 deletions(-) diff --git a/ROADMAP.md b/ROADMAP.md index e904ab5..eca9e6e 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -830,6 +830,23 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari BOTH Qwen2.5-1.5B and LFM2.5-1.2B; top-3 first-forward ids exact in order, top-5 overlap ≥ 4/5, max |Δlogprob| 2.7e-1 (the reference's own Q8_K *activation* quantization in its integer Q4_K kernels — an order above the BF16 leg's 1.5e-2; our f32 path doesn't quantize activations).* +- [x] **`fetch_model` fetches the sibling files the import gates read** — found while installing the OLMoE + safetensors checkpoint through the registry (2026-08-20). `hub::fetch_model_with` fetched exactly + `config.json`, `tokenizer.json`, and the weights, so a checkpoint installed through Mummu's OWN + registry arrived with no `tokenizer_config.json` — and the 2026-07-19/20/21 gates all fail **open**: + `validate_checkpoint_dir` returns `Ok(None)`, the EOS-agreement + added-token-id + tool-call-convention + checks silently no-op, and the `tokenizer_config` the loaders surface is always `None`. The gates + appeared to work only because every locally-cached fixture had been populated by hand. Fixed: both + `tokenizer_config.json` and `chat_template.jinja` (the standalone-template checkpoints of the + 2026-07-20 item) are now fetched as **optional** files — one HEAD each, and only a 404 counts as + absent, so a 403 on a gated repo or a 5xx still reaches the real fetch and is reported rather than + swallowed as "the repo just doesn't ship it". `config.json` / `tokenizer.json` stay required. They are + fetched BEFORE the weights because the single-file branch returns early — placing them after it would + have fetched them for sharded checkpoints only, which is exactly the kind of half-fix this item exists + to avoid. Proof (`tests/real_hub.rs::a_registry_install_arrives_with_the_files_the_import_gates_read`): + a catalog model installed into a **clean** dir — no hand-populated fixture can satisfy it — now has + `tokenizer_config.json` on disk AND `validate_checkpoint_dir` returns `Some`, i.e. the gate has + something to check instead of quietly passing. *(2026-08-20)* - [ ] **GPTQ / AWQ** (HF safetensors) — import the calibration-quantized int4/int8 layouts most "quantized on the Hub" models ship as (a `.safetensors` payload + a quant config), dequant or keep-quant into Burn. *(2026-07-13 research)* Both are quantization *algorithms*, not formats — the artifact is ordinary diff --git a/crates/mummu/src/hub.rs b/crates/mummu/src/hub.rs index d32747f..712bb6f 100644 --- a/crates/mummu/src/hub.rs +++ b/crates/mummu/src/hub.rs @@ -160,6 +160,42 @@ fn announced_sha256(url: &str) -> Result, HubError> { })) } +/// The sibling files that are fetched **only if the repo ships them**. +/// +/// `tokenizer_config.json` is what the import-validation gates read (EOS +/// agreement, added-token ids, tool-call convention); `chat_template.jinja` is +/// how checkpoints that keep their template out of `tokenizer_config.json` +/// ship it. Neither is universal on the Hub, so a 404 is a legitimate answer — +/// but not asking for them at all silently disarms those gates for every model +/// installed through this path. +const OPTIONAL_FILES: [&str; 2] = ["tokenizer_config.json", "chat_template.jinja"]; + +/// Does the repo actually ship `url`? +/// +/// One HEAD, redirects not followed: the Hub answers a `resolve` URL for a +/// present file with 200 (or a 302 to the CDN) and for an absent one with 404. +/// ONLY 404 counts as absent — anything else (403 on a gated repo, a 5xx) is +/// reported as present so the real fetch raises it properly rather than this +/// probe swallowing it as "the repo just doesn't have it". +fn repo_has_file(url: &str) -> Result { + assert!(url.starts_with("https://"), "refusing non-https url: {url}"); + let agent: ureq::Agent = ureq::Agent::config_builder() + .max_redirects(0) + .http_status_as_error(false) + .build() + .into(); + let resp = agent.head(url).call().map_err(|e| HubError::Http { + url: url.into(), + reason: e.to_string(), + })?; + let status = resp.status().as_u16(); + debug_assert!( + (100..600).contains(&status), + "http status in range: {status}" + ); + Ok(status != 404) +} + /// Streaming sha256 of a file on disk, as lowercase hex. fn sha256_hex_of_file(path: &Path) -> Result { let io_err = |e: std::io::Error| HubError::Io { @@ -385,9 +421,15 @@ fn download( } /// Fetch a whole model from the Hub into `dest_dir`: `config.json`, -/// `tokenizer.json`, and the weights — `model.safetensors` when the repo is -/// single-file, else every shard listed by `model.safetensors.index.json`. +/// `tokenizer.json`, the weights — `model.safetensors` when the repo is +/// single-file, else every shard listed by `model.safetensors.index.json` — +/// and, when the repo ships them, the optional siblings in [`OPTIONAL_FILES`]. /// Returns `dest_dir` ready for the per-model `load_from_dir`. +/// +/// `config.json` and `tokenizer.json` are required: a 404 on either is an +/// error. The optional siblings are best-effort, because a repo that does not +/// ship them is normal — but they are asked for, so a checkpoint installed +/// this way arrives with the files the import-validation gates need. pub fn fetch_model( repo: &str, revision: &str, @@ -420,6 +462,15 @@ pub fn fetch_model_with( &mut on_progress, )?; } + // Then the optional siblings. These MUST be fetched before the weights, + // because the single-file branch below returns early on success — putting + // them after it would fetch them for sharded checkpoints only. + for file in OPTIONAL_FILES { + let url = hub_file_url(repo, revision, file); + if repo_has_file(&url)? { + fetch_file_with(&url, &dest_dir.join(file), opts, &mut on_progress)?; + } + } // Single-file first (the common case for the small-model tiers we target). let single = fetch_file_with( &hub_file_url(repo, revision, "model.safetensors"), @@ -477,6 +528,29 @@ mod tests { let _ = hub_file_url("qwen", "main", "config.json"); } + /// The gates that read these files fail OPEN — `validate_checkpoint_dir` + /// returns `Ok(None)` when `tokenizer_config.json` is absent, so the + /// EOS-agreement and added-token-id checks simply do not run. That makes + /// "did we even ask the Hub for it?" the load-bearing question, and it is + /// what this list answers. + #[test] + fn optional_files_cover_the_siblings_the_import_gates_read() { + assert!(OPTIONAL_FILES.contains(&"tokenizer_config.json")); + assert!(OPTIONAL_FILES.contains(&"chat_template.jinja")); + for file in OPTIONAL_FILES { + assert!( + !["config.json", "tokenizer.json"].contains(&file), + "{file} is required, not optional — a 404 on it must stay an error" + ); + assert_eq!( + hub_file_url("allenai/OLMoE-1B-7B-0125-Instruct", "main", file), + format!( + "https://huggingface.co/allenai/OLMoE-1B-7B-0125-Instruct/resolve/main/{file}" + ) + ); + } + } + #[test] fn shard_index_dedupes_and_sorts() { let idx = br#"{"metadata":{},"weight_map":{ diff --git a/crates/mummu/tests/real_hub.rs b/crates/mummu/tests/real_hub.rs index fc70780..00b84da 100644 --- a/crates/mummu/tests/real_hub.rs +++ b/crates/mummu/tests/real_hub.rs @@ -275,3 +275,45 @@ fn hub_gguf_spec_downloads_and_parses() { f.tensors.len() ); } + +/// The install path must feed the import-validation gates. +/// +/// `validate_checkpoint_dir` fails OPEN: with no sibling `tokenizer_config.json` +/// it returns `Ok(None)` and the EOS-agreement + added-token-id checks quietly +/// do not run. So a checkpoint installed through Mummu's own registry has to +/// ARRIVE with that file, or every one of those gates is decorative for exactly +/// the models Mummu installed itself. This fetches a catalog model into a +/// **clean** dir (no hand-populated fixtures) and asserts the gate found +/// something to check. +#[test] +#[ignore = "needs network (MUMMU_HUB_DEST names the download dir; ~90 MB)"] +fn a_registry_install_arrives_with_the_files_the_import_gates_read() { + let Some(dest) = std::env::var_os("MUMMU_HUB_DEST").map(PathBuf::from) else { + panic!("set MUMMU_HUB_DEST to a scratch dir for the ~90 MB download"); + }; + // A dir of our own, so nothing here can be satisfied by a fixture someone + // populated by hand — the whole point of the check. + let clean = dest.join("gate-feed-probe"); + if clean.exists() { + std::fs::remove_dir_all(&clean).expect("clear the probe dir"); + } + let dir = minilm_spec().fetch(&clean, |_| {}).expect("hub fetch"); + + assert!( + dir.join("tokenizer_config.json").is_file(), + "the install must fetch tokenizer_config.json, or the import gates no-op" + ); + + // And the gate itself must now have something to validate. MiniLM declares + // no eos in config.json, so pass an empty set and assert only that the + // config was FOUND — `Ok(None)` here is the silent-no-op this test exists + // to forbid. + let cfg = mummu::tokenizer::validate_checkpoint_dir(&dir, &[], None) + .expect("tokenizer_config.json parses and agrees with tokenizer.json"); + assert!( + cfg.is_some(), + "validate_checkpoint_dir returned Ok(None) — the gate silently no-opped" + ); + + std::fs::remove_dir_all(&clean).ok(); +} From ee61c666205f4fb614a7f178e10ec7f8d7d02edc Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 20 Aug 2026 11:27:55 -0500 Subject: [PATCH 4/5] models: import OLMoE from HF safetensors, fusing the 64 experts on load MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The OLMoE port could only read GGUF, because a GGUF ships the expert bank already fused (`ffn_*_exps`) while the HF checkpoint stores every expert as its own tensor across three shards. `burn-store` cannot bridge that: it finds only a single `model.safetensors`, and its remapping is 1:1, not N:1. `safetensors.rs` is a sharded reader plus a fusing rewriter for exactly those two gaps. It plans the entire output before reading one payload byte, so a group that is not exactly `count` members `0..count` is a loud BadGroup rather than a short expert bank that loads clean and computes wrong. Members are ordered NUMERICALLY, not lexicographically — `experts.10` sorts before `experts.2` as text, and that silent mis-ordering is the whole risk here. `olmoe::load_from_dir` then rides the ordinary SafetensorsStore + adapter-chain + load_checked pipeline, landing on the same params the GGUF path renames to (pinned by a unit test — both importers must agree or one is loading a different model). REAL-WEIGHTS proof, on the 3-shard 13.84 GB bf16 checkpoint: - fuses [64, 1024, 2048] banks over 16 layers, checked-load in 136.1 s on the CPU backend, zero missing params; - emits a live distribution (sanity top 194, spread 20.6) and greedy-decodes 8 tokens. This leg uses an arbitrary token probe rather than a prompt, so read it as liveness, not coherence; - the check a mis-ordering could NOT survive: layer 5 / expert 37's `gate_proj`, read independently out of the raw shard bytes, is bit-identical to slot 37 of the fused bank — 0 mismatches over 2 097 152 values. bf16→f32 is an exact 16-bit widening, so this is bit-equality, not a tolerance. Two defects the real run caught that unit tests had not: - `checkpoint_shards` planned off the index without checking the shards were on disk. A half-finished download therefore reported as a COMPLETE checkpoint: the fetch was skipped and the load died later with a bare `os error 2`. It now fails at planning time naming the missing shard, and says so when a `.part` sibling shows the download was interrupted. - The fuse materialized the payload in RAM twice — a `data` buffer, then a `blob` copy — and the second 13.8 GB allocation genuinely failed on this 128 GB machine (`memory allocation of 13838346237 bytes failed`). `fuse_checkpoint_to_file` now streams header + payload straight to a temp file through ONE reusable per-part buffer (~4 MiB, a single expert projection), dropping peak from ~42 GB (13.8 blob + ~28 model) to the model alone. A `FusedTemp` guard deletes the scratch file on every exit path, verified gone after the run, and a unit test pins the in-memory and to-file fuses as byte-identical so the two are one importer, not two. Also hardened while carrying it: every `.expect()` on a production path in the new module is gone (a fallible `to_usize` returning OverBound), and assertions went 5 -> 9 across 9 functions. 232 lib tests green, clippy --all-targets clean. Co-Authored-By: Claude Opus 4.8 --- README.md | 8 + ROADMAP.md | 47 +- crates/mummu/src/lib.rs | 1 + crates/mummu/src/models/olmoe.rs | 247 +++- crates/mummu/src/registry.rs | 14 +- crates/mummu/src/safetensors.rs | 1162 ++++++++++++++++++ crates/mummu/tests/real_olmoe_safetensors.rs | 184 +++ 7 files changed, 1654 insertions(+), 9 deletions(-) create mode 100644 crates/mummu/src/safetensors.rs create mode 100644 crates/mummu/tests/real_olmoe_safetensors.rs diff --git a/README.md b/README.md index 8b86963..2be8682 100644 --- a/README.md +++ b/README.md @@ -85,6 +85,14 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys 1.15 s/token on the CPU backend (~28 GB f32 resident; a 16 GB card waits on keep-quantized VRAM, tracked in P9). Attention learned OLMoE's whole-projection q/k RMSNorm placement, inferred from the loaded norm's own width, so every existing checkpoint loads byte-unchanged. + It also loads from the **HF safetensors** release, where the 64 experts are stored as separate + tensors across three shards: the importer fuses `experts.{0..63}.{gate,up,down}_proj` into the same + `[64, 1024, 2048]` banks, in numeric (not lexicographic) member order, validating group completeness + before reading a payload byte. Proven on the real 13.84 GB bf16 checkpoint — 16 layers checked-loaded + in 136.1 s, and layer 5 / expert 37's `gate_proj` read straight from the raw shard bytes is + **bit-identical to slot 37 of the fused bank across all 2 097 152 values**. The fuse streams to a + temp file rather than RAM, so a checkpoint this size costs the model's footprint, not the model plus + a second copy of itself. - **All three models are parity-verified** — the two-leg P7 gate passes for Qwen2.5-1.5B on the reference GPU: single-forward top-5 logits match a Candle f32 reference (max |Δlogit| 2.7e-5, `tests/parity_qwen2.rs` + the committed `tools/candle-probe` fixture) and a 24-token greedy sequence diff --git a/ROADMAP.md b/ROADMAP.md index eca9e6e..0f2e10b 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -682,12 +682,57 @@ a benchmark holds/improves its budget; README perf claims link an artifact. before treating the gather as refuted. Both are gated on `bench/BASELINE.md` like the rest — https://github.com/ggml-org/llama.cpp/pull/25294 · https://huggingface.co/blog/Doctor-Shotgun/llamacpp-moe-offload-guide -- [ ] **OLMoE from HF safetensors** — the port loads GGUF only because HF stores each expert separately +- [x] **OLMoE from HF safetensors** — the port loads GGUF only because HF stores each expert separately (`model.layers.N.mlp.experts.{0..63}.{gate,up,down}_proj.weight`) while `MoeExperts` holds one fused `[experts, out, in]` tensor per projection. Needs a concat-on-import step (64 slices → one tensor, in expert order) in the safetensors path — mechanical, but it wants its own fixture (the bf16 checkpoint is ~14 GB) and a byte-equality check against the GGUF-loaded weights. *(2026-08-03, split from the MoE item.)* + *(2026-08-20) Shipped.* `crates/mummu/src/safetensors.rs` is a sharded-safetensors reader plus an + **N:1 fusing rewriter** — the two things `burn-store` cannot do for us (it finds only a single + `model.safetensors`, and its remapping is 1:1). It plans the whole output before reading one payload + byte, so a group that is not exactly `count` members `0..count` is a loud `BadGroup` rather than a + short bank that loads clean and computes wrong; members are ordered **numerically, not + lexicographically** (`experts.10` sorts before `experts.2` as text — the trap this whole item exists + to avoid). `olmoe::load_from_dir` maps `model.layers.N.mlp.experts.{0..63}.{gate,up,down}_proj.weight` + onto the same fused `[experts, out, in]` banks the GGUF path lands on, then rides the ordinary + `SafetensorsStore` + adapter-chain + `load_checked` pipeline. + **REAL-WEIGHTS proof** (`tests/real_olmoe_safetensors.rs`, the 3-shard 13.84 GB bf16 checkpoint): + the fuse builds `[64, 1024, 2048]` banks over 16 layers, checked-loads in **136.1 s** on the CPU + backend, and emits a live distribution (sanity top 194, spread 20.6) that greedy-decodes 8 tokens. + The check a mis-ordering could not survive: layer 5 / expert 37's `gate_proj`, read independently out + of the raw shard bytes, is **bit-identical to slot 37 of the fused bank — 0 mismatches over 2 097 152 + values** (bf16 -> f32 is an exact 16-bit widening, so this is bit-equality, not a tolerance). Note the + decode leg uses an arbitrary token probe, not a prompt, so it is a liveness check; the *weight* + correctness evidence is the bit-exact one. 16 `safetensors` unit tests, incl. one pinning the + safetensors fuse targets to the SAME module names the GGUF path renames to — both import paths must + land on the same params or one of them is loading a different model. + Two things the gate caught that the code did not survive first contact with: + (a) `checkpoint_shards` planned straight off the index without checking the shards were **on disk**, + so an interrupted download reported as a complete checkpoint, skipped the resume, and died later with + a bare `os error 2` from inside the load — it now fails at planning time naming the missing shard, and + says so explicitly when a `.part` sibling shows the download was interrupted; + (b) the fuse materialized the payload in RAM **twice** (a `data` buffer, then a `blob` copy), and the + second 13.8 GB allocation genuinely failed on this 128 GB box — + `memory allocation of 13838346237 bytes failed`. `fuse_checkpoint_to_file` now streams header + + payload straight to a temp file through one reusable per-part buffer (~4 MiB, one expert projection), + so peak drops from ~42 GB (13.8 blob + ~28 model) to the model alone; a `FusedTemp` guard deletes the + scratch file on every exit path, verified empty after the run. A unit test pins the two fuse paths as + **byte-identical**, so the big-model path and the small-model path stay one importer. +- [ ] **Give the GGUF dequant path the same streaming sink the safetensors fuse just got** — found while + fixing the OLMoE safetensors OOM (2026-08-20). `gguf::dequant_to_safetensors` has the identical + double-buffer: it fills a `data` Vec of `total_f32_bytes`, then copies it into a second `blob` Vec of + `8 + header + data`, so peak is **2x the dequantized f32 payload** before the model is even built. + For OLMoE-1B-7B that is ~28 GB f32 → ~56 GB peak, which is most of why that model reads as + "CPU-only, and only on a big box". `safetensors::fuse_into` is the template: plan, then stream header + + payload into an `impl Write` through one bounded per-tensor buffer, with a `*_to_file` variant for + `SafetensorsStore::from_file`. Prove it the same way — a unit test pinning the in-memory and + to-file outputs as byte-identical, then re-run the `real_gguf` + `parity_gguf` legs unchanged. +- [ ] **`gguf.rs` still has `.expect()` on production paths** — `dequant_to_safetensors` and the metadata + readers carry `usize::try_from(..).expect("bounded above")` (4 sites), the same pattern removed from + `safetensors.rs` on 2026-08-20 in favour of a fallible `to_usize` that returns `OverBound`. Mechanical, + and it keeps the no-panic-on-production-paths rule true across the whole import suite rather than in + the newest module only. - [ ] **Qwen3.5 hybrid (`qwen35`) architecture port** — split from the Qwen3.5-tier item when the 2026-07-30 header probe showed Qwen3.5-4B/9B are a hybrid **linear-attention/SSM + periodic full-attention** arch (`qwen35.ssm.*` metadata: conv_kernel/state_size/group_count/time_step_rank/ diff --git a/crates/mummu/src/lib.rs b/crates/mummu/src/lib.rs index 8bf0e82..a97a709 100644 --- a/crates/mummu/src/lib.rs +++ b/crates/mummu/src/lib.rs @@ -21,6 +21,7 @@ pub mod models; pub mod nn; pub mod plan; pub mod registry; +pub mod safetensors; /// Render a checkpoint's own imported chat template (feature `jinja-template`). #[cfg(feature = "jinja-template")] pub mod template; diff --git a/crates/mummu/src/models/olmoe.rs b/crates/mummu/src/models/olmoe.rs index 19e0f8a..c20d6f0 100644 --- a/crates/mummu/src/models/olmoe.rs +++ b/crates/mummu/src/models/olmoe.rs @@ -13,12 +13,15 @@ //! CPU backend on the reference 128 GB machine. Expert streaming / offload is //! the P6 placement item; keep-quantized is P9. //! -//! Import is **GGUF-only** for now: the HF safetensors checkpoint stores each -//! expert as a separate tensor (`mlp.experts.{i}.gate_proj.weight`), which -//! would need a 64-way concat on import; the GGUF ships the experts already -//! fused (`ffn_*_exps`) in exactly the layout [`MoeExperts`] holds. +//! Import covers **both** sources. A GGUF ships the experts already fused +//! (`ffn_*_exps`) in exactly the layout `MoeExperts` holds; the HF safetensors +//! checkpoint stores each expert separately +//! (`mlp.experts.{i}.{gate,up,down}_proj.weight`) and is sharded, so +//! [`load_from_dir`] runs it through [`crate::safetensors::fuse_checkpoint`], +//! which reads every shard and stacks each 64-member expert group into one +//! `[experts, out, in]` tensor before the ordinary checked-load pipeline. -use std::path::Path; +use std::path::{Path, PathBuf}; use burn::module::Module; use burn::nn::{Embedding, EmbeddingConfig, Linear, LinearConfig, RmsNorm, RmsNormConfig}; @@ -26,12 +29,13 @@ use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, SafetensorsStore}; use burn::tensor::{Int, Tensor, TensorData, backend::Backend}; use crate::gguf::{GgufFile, GgufMap, GgufTensorInfo, GgufValue}; -use crate::import::{CastFloatAdapter, ImportError, load_checked}; +use crate::import::{CastFloatAdapter, ImportError, load_checked, required_file}; use crate::models::CausalLm; use crate::models::qwen2::{EosIds, gguf_f32, gguf_usize}; use crate::nn::{ GqaAttention, GqaAttentionConfig, LayerKv, SparseMoe, SparseMoeConfig, causal_mask, rope_tables, }; +use crate::safetensors::{Fuse, fuse_checkpoint_to_file}; /// OLMoE architecture hyperparameters (HF `config.json` field names). #[derive(Debug, Clone, serde::Deserialize)] @@ -184,6 +188,10 @@ pub struct Olmoe { pub struct LoadedOlmoe { pub model: Olmoe, pub config: OlmoeConfig, + /// The sibling `tokenizer_config.json`, when the checkpoint dir ships one + /// — config-driven EOS/BOS/PAD for a consumer to read. `None` for a GGUF + /// load (self-contained: EOS rides the GGUF metadata). + pub tokenizer_config: Option, } fn build(cfg: &OlmoeConfig, device: &B::Device) -> Olmoe { @@ -310,7 +318,161 @@ pub fn load_from_gguf( .allow_partial(true), ); load_checked(&mut model, &mut store, path)?; - Ok(LoadedOlmoe { model, config }) + Ok(LoadedOlmoe { + model, + config, + tokenizer_config: None, + }) +} + +/// How a source tensor of an HF OLMoE checkpoint reaches the module. +/// +/// Everything but the expert bank passes through under its own name (the +/// `install_remaps` chain does the HF→module renaming downstream, exactly as +/// on the single-file safetensors path). The per-expert projections are the +/// N:1 case: `model.layers.3.mlp.experts.7.gate_proj.weight` is member 7 of +/// the fused `model.layers.3.mlp.experts.gate`. +/// +/// Deliberately *not* a strict allow-list: unrecognized tensors are kept, and +/// a checkpoint that renames the expert projections then fails loudly one +/// stage later in `load_checked` — which names the missing `experts.gate` +/// param in its report — rather than here with a less specific message. +fn olmoe_hf_fuse(name: &str, num_experts: usize) -> Fuse { + let Some(target) = fused_expert_target(name, num_experts) else { + return Fuse::Keep(name.to_string()); + }; + target +} + +/// `model.layers.{L}.mlp.experts.{I}.{gate,up,down}_proj.weight` → its slot in +/// the fused bank, or `None` when the name is not a per-expert projection. +fn fused_expert_target(name: &str, num_experts: usize) -> Option { + let rest = name.strip_prefix("model.layers.")?; + let (layer, rest) = rest.split_once('.')?; + layer.parse::().ok()?; + let rest = rest.strip_prefix("mlp.experts.")?; + let (index, rest) = rest.split_once('.')?; + let index: usize = index.parse().ok()?; + let projection = match rest { + "gate_proj.weight" => "gate", + "up_proj.weight" => "up", + "down_proj.weight" => "down", + _ => return None, + }; + Some(Fuse::Stack { + target: format!("model.layers.{layer}.mlp.experts.{projection}"), + index, + count: num_experts, + }) +} + +/// Load an OLMoE model from an **HF safetensors checkpoint dir**. +/// +/// The checkpoint is sharded (`model-0000N-of-0000M.safetensors` + +/// `model.safetensors.index.json`) and stores every expert separately, so the +/// shards are read and the expert groups fused into `[experts, out, in]` +/// tensors first; the fused blob then rides the SAME adapter chain and +/// `load_checked` as every other import path. Source dtype is preserved by +/// the fuse (HF ships bf16) and cast to the backend float on load. +/// +/// Budget note: the fused blob is the checkpoint's own size (~13.8 GB in bf16 +/// for the 1B-7B) and the loaded f32 model is ~28 GB — size the target device. +/// Owns the fused scratch file for the life of one `load_from_dir`. +/// +/// The fused blob is as large as the checkpoint (13.8 GB for the 1B-7B), so +/// leaving one behind on a failed load would quietly fill the disk over a few +/// retries. `Drop` removes it on every exit path, success or `?`. +struct FusedTemp { + path: PathBuf, +} + +impl FusedTemp { + /// Placed beside the checkpoint, so the scratch write lands on the same + /// volume as the weights rather than on a small system temp drive. + fn new(dir: &Path) -> Result { + assert!(!dir.as_os_str().is_empty(), "checkpoint dir must be named"); + let path = dir.join(format!( + "mummu-fused-{}.safetensors.tmp", + std::process::id() + )); + // A leftover from a killed process must never be mistaken for ours. + if path.exists() { + std::fs::remove_file(&path).map_err(|e| ImportError::Parse { + file: path.clone(), + reason: format!("could not clear a stale fused scratch file: {e}"), + })?; + } + Ok(Self { path }) + } + + fn path(&self) -> &Path { + debug_assert!(!self.path.as_os_str().is_empty(), "scratch path is named"); + &self.path + } +} + +impl Drop for FusedTemp { + fn drop(&mut self) { + // Best effort by construction: a Drop that can fail has nowhere to + // report to, and a stranded scratch file is not worth a panic. + let _ = std::fs::remove_file(&self.path); + } +} + +pub fn load_from_dir( + dir: &Path, + device: &B::Device, +) -> Result, ImportError> { + let cfg_path = required_file(dir, "config.json")?; + let cfg_bytes = std::fs::read(&cfg_path).map_err(|e| ImportError::Parse { + file: cfg_path.clone(), + reason: e.to_string(), + })?; + let config = OlmoeConfig::from_json_bytes(&cfg_bytes).map_err(|reason| ImportError::Parse { + file: cfg_path, + reason, + })?; + + // Cross-check the sibling metadata before touching weights, same as the + // dense loaders. `None` for the expected tool-call convention: Mummu ships + // no hardcoded OLMoE `chat` renderer to contradict, so only the EOS and + // added-token-id checks apply. + let tokenizer_config = + crate::tokenizer::validate_checkpoint_dir(dir, &config.eos_token_id.to_vec(), None)?; + + let num_experts = config.num_experts; + // Fuse to a TEMP FILE, not to RAM. The in-memory fuse needs the whole + // payload resident (13.8 GB for the 1B-7B) on top of the ~28 GB f32 model + // the load then builds; streaming it to disk keeps the peak at the model + // alone. The file is this process's to delete, and it is deleted whether + // the load succeeds or fails. + let fused = FusedTemp::new(dir)?; + let bytes = fuse_checkpoint_to_file( + dir, + &|name| Some(olmoe_hf_fuse(name, num_experts)), + fused.path(), + ) + .map_err(|e| ImportError::Parse { + file: dir.to_path_buf(), + reason: e.to_string(), + })?; + assert!(bytes > 8, "a fused checkpoint yields a non-empty payload"); + + let mut model = build::(&config, device); + // The backend's float dtype from the TYPE (`B::FloatElem`), never a probe + // tensor (the per-device default-dtype policy hazard). + let target_float = ::dtype(); + let mut store = install_remaps( + SafetensorsStore::from_file(fused.path().to_path_buf()) + .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) + .allow_partial(true), + ); + load_checked(&mut model, &mut store, dir)?; + Ok(LoadedOlmoe { + model, + config, + tokenizer_config, + }) } impl CausalLm for LoadedOlmoe { @@ -459,6 +621,7 @@ mod tests { let loaded = LoadedOlmoe:: { model: build(&cfg, &device), config: cfg, + tokenizer_config: None, }; let prompt: Vec = vec![3, 14, 15, 9, 26]; @@ -615,8 +778,78 @@ mod tests { let loaded = LoadedOlmoe:: { model: build(&cfg, &device), config: cfg, + tokenizer_config: None, }; let out = loaded.greedy_generate(&[1, 2, 3], 4, &device).unwrap(); assert!(out.len() <= 4); } + + /// The HF per-expert projections fuse onto EXACTLY the module names the + /// GGUF path renames its pre-fused banks to. Both import paths must land + /// on the same params, or one of them is loading a different model. + #[test] + fn hf_expert_fusion_targets_match_the_gguf_names() { + for (projection, ggml) in [ + ("gate", "ffn_gate_exps"), + ("up", "ffn_up_exps"), + ("down", "ffn_down_exps"), + ] { + let hf = format!("model.layers.3.mlp.experts.7.{projection}_proj.weight"); + let expected = olmoe_gguf_name(&format!("blk.3.{ggml}.weight")).unwrap(); + assert_eq!( + olmoe_hf_fuse(&hf, 64), + Fuse::Stack { + target: expected, + index: 7, + count: 64, + }, + "{projection}: safetensors and GGUF must fuse to the same param" + ); + } + } + + /// Non-expert tensors pass through untouched — the `install_remaps` chain + /// does the HF→module renaming downstream, same as the dense loaders. + #[test] + fn non_expert_tensors_pass_through_by_name() { + for name in [ + "model.embed_tokens.weight", + "model.norm.weight", + "lm_head.weight", + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.q_norm.weight", + "model.layers.0.input_layernorm.weight", + // The ROUTER is a plain Linear, not part of the expert bank — + // fusing it would be a silent disaster. + "model.layers.0.mlp.gate.weight", + ] { + assert_eq!( + olmoe_hf_fuse(name, 64), + Fuse::Keep(name.to_string()), + "{name} must pass through" + ); + } + } + + /// The expert index is parsed as a NUMBER, so the fuse plan can order the + /// bank numerically; `experts.10` must not be read as expert 1. + #[test] + fn expert_index_is_parsed_numerically() { + let at = |i: usize| match olmoe_hf_fuse( + &format!("model.layers.0.mlp.experts.{i}.gate_proj.weight"), + 64, + ) { + Fuse::Stack { index, .. } => index, + other => panic!("expected a Stack, got {other:?}"), + }; + assert_eq!(at(0), 0); + assert_eq!(at(1), 1); + assert_eq!(at(10), 10); + assert_eq!(at(63), 63); + // A non-numeric member is not an expert projection at all. + assert_eq!( + olmoe_hf_fuse("model.layers.0.mlp.experts.x.gate_proj.weight", 64), + Fuse::Keep("model.layers.0.mlp.experts.x.gate_proj.weight".to_string()) + ); + } } diff --git a/crates/mummu/src/registry.rs b/crates/mummu/src/registry.rs index 7ae6f02..66a3c2d 100644 --- a/crates/mummu/src/registry.rs +++ b/crates/mummu/src/registry.rs @@ -21,7 +21,8 @@ pub enum Architecture { /// `models::minilm` — all-MiniLM BERT sentence embedder. MiniLm, /// `models::olmoe` — OLMoE sparse mixture-of-experts decoder (the zoo's - /// first MoE; GGUF import only — HF ships the experts unfused). + /// first MoE). Imports from GGUF (experts pre-fused) or from HF + /// safetensors (experts fused on import). Olmoe, } @@ -248,6 +249,17 @@ pub fn catalog() -> Vec { }, disk_bytes_estimate: 4_210_000_000, }, + // The same MoE from its HF source: 3 bf16 safetensors shards + an + // index, with the 64 experts stored separately. `load_from_dir` fuses + // them into the `[experts, out, in]` banks the module holds. + ModelSpec { + name: "olmoe-1b-7b-0125-instruct".into(), + repo: "allenai/OLMoE-1B-7B-0125-Instruct".into(), + revision: "main".into(), + architecture: Architecture::Olmoe, + format: WeightFormat::Safetensors, + disk_bytes_estimate: 13_800_000_000, + }, ]; debug_assert!( entries.iter().all(|s| s.validate().is_ok()), diff --git a/crates/mummu/src/safetensors.rs b/crates/mummu/src/safetensors.rs new file mode 100644 index 0000000..3b2da3a --- /dev/null +++ b/crates/mummu/src/safetensors.rs @@ -0,0 +1,1162 @@ +//! Safetensors **reader** + a fusing rewriter for checkpoints whose on-disk +//! layout does not match the module layout. +//! +//! `burn-store` reads safetensors for us on the ordinary path, so this module +//! exists for the two things it cannot do: +//! +//! 1. **Sharded checkpoints.** Anything past ~5 GB ships as +//! `model-0000N-of-0000M.safetensors` + a `model.safetensors.index.json` +//! weight map. `import::weights_file` only ever finds a single +//! `model.safetensors`. +//! 2. **N:1 tensor fusion.** `burn-store`'s remapping is 1:1 (rename); an MoE +//! checkpoint stores every expert separately +//! (`mlp.experts.{0..63}.gate_proj.weight`) while the module holds ONE +//! fused `[experts, out, in]` param — exactly the ggml `ffn_*_exps` layout +//! a GGUF ships pre-fused. +//! +//! [`fuse_checkpoint`] reads every shard's header, plans the output layout, +//! validates it, and only then copies payload bytes — so a checkpoint missing +//! expert 37 fails before a single weight byte is read, rather than loading +//! clean and computing wrong. The result is an ordinary in-memory safetensors +//! blob that goes through the SAME `SafetensorsStore::from_bytes` + +//! adapter-chain + `load_checked` pipeline as every other import path (the +//! GGUF path's `dequant_to_safetensors` is the precedent). +//! +//! Source dtypes are preserved verbatim — the bf16→backend-float cast stays +//! where it already lives, in `CastFloatAdapter` on the load pipeline. + +use std::collections::HashMap; +use std::fs::File; +use std::io::{BufReader, Read, Seek, SeekFrom, Write}; +use std::path::{Path, PathBuf}; + +/// Largest safetensors JSON header accepted, per shard. Real headers are a +/// few hundred KB (one entry per tensor); 64 MiB is a corrupt or hostile file. +const MAX_HEADER_BYTES: u64 = 64 << 20; + +/// Largest tensor count accepted across a whole checkpoint. A 64-expert MoE +/// at 16 layers already declares ~3 000; 1M is a runaway index. +const MAX_TENSORS: usize = 1 << 20; + +/// Largest fused payload [`fuse_checkpoint`] will build in RAM. Matches the +/// GGUF path's ceiling (the reference machine has 128 GB) — note this blob +/// keeps the SOURCE dtype, so a bf16 checkpoint costs half its f32 footprint. +const MAX_FUSED_BYTES: u64 = 48 << 30; + +/// Largest SINGLE tensor (or fused member) the streaming fuse will buffer. +/// The widest real one we carry is a 64-expert projection member at ~4 MiB; +/// 4 GiB is a corrupt header claiming an absurd tensor. +const MAX_PART_BYTES: u64 = 4 << 30; + +/// Largest number of shards in an index (mirrors `hub::MAX_SHARDS`). +const MAX_SHARDS: usize = 256; + +/// What went wrong reading or fusing a safetensors checkpoint. +#[derive(Debug, thiserror::Error)] +pub enum SafetensorsError { + #[error("safetensors {path}: {source}")] + Io { + path: String, + source: std::io::Error, + }, + #[error("safetensors {path}: header is not valid JSON: {reason}")] + BadHeader { path: String, reason: String }, + #[error("safetensors {path}: {what} {count} exceeds the {bound} bound")] + OverBound { + path: String, + what: &'static str, + count: u64, + bound: u64, + }, + #[error("safetensors {path}: tensor '{name}': {reason}")] + BadTensor { + path: String, + name: String, + reason: String, + }, + /// A fused group is not exactly `count` distinct members `0..count`. This + /// is the load-bearing check: a silently short expert bank would load + /// clean and compute wrong. + #[error("fused tensor '{target}': {reason}")] + BadGroup { target: String, reason: String }, + #[error("no safetensors checkpoint in {0} (looked for model.safetensors and *.index.json)")] + NoCheckpoint(PathBuf), + /// The index names a shard that is not on disk. A sibling `.part` means an + /// interrupted download, which is worth saying out loud: the alternative is + /// this surfacing as a bare `os error 2` from deep inside the load. + #[error("incomplete checkpoint {dir}: index names {missing}, which is not present{hint}")] + MissingShard { + dir: PathBuf, + missing: String, + hint: &'static str, + }, +} + +/// `u64` -> `usize` as an error rather than a panic. +/// +/// Every caller has already bounded `value` against a file length or one of +/// the `MAX_*` ceilings, so on a 64-bit target this cannot fail — but an +/// import path is exactly where "cannot fail" should still be an `Err` +/// instead of an `.expect()`, so a 32-bit build degrades to a clean error. +fn to_usize(value: u64, path: &Path, what: &'static str) -> Result { + debug_assert!( + value <= usize::MAX as u64, + "{what} fits usize on this target" + ); + usize::try_from(value).map_err(|_| SafetensorsError::OverBound { + path: path.display().to_string(), + what, + count: value, + bound: usize::MAX as u64, + }) +} + +/// One tensor as the on-disk header describes it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TensorEntry { + /// Safetensors dtype token, verbatim (`BF16`, `F32`, `I64`, …). + pub dtype: String, + /// Row-major shape. + pub shape: Vec, + /// `[start, end)` within the shard's payload region. + pub offsets: (u64, u64), +} + +impl TensorEntry { + /// Payload length in bytes. + #[must_use] + pub fn byte_len(&self) -> u64 { + debug_assert!( + self.offsets.1 >= self.offsets.0, + "offsets validated on read" + ); + self.offsets.1 - self.offsets.0 + } + + /// Elements implied by the shape (1 for a scalar — an empty shape). + #[must_use] + pub fn element_count(&self) -> u64 { + self.shape.iter().product() + } +} + +/// A parsed safetensors header: the tensor table plus where its payload starts. +#[derive(Debug)] +pub struct SafetensorsHeader { + pub path: PathBuf, + /// Entries in **file order** (the order payload bytes appear), which is + /// what makes a sequential read cheap. + pub tensors: Vec<(String, TensorEntry)>, + /// Absolute file offset where the payload region begins. + pub data_offset: u64, +} + +impl SafetensorsHeader { + /// Read and validate one shard's header. No payload bytes are read. + pub fn open(path: &Path) -> Result { + let io = |source: std::io::Error| SafetensorsError::Io { + path: path.display().to_string(), + source, + }; + let mut file = BufReader::new(File::open(path).map_err(io)?); + let file_len = file.get_ref().metadata().map_err(io)?.len(); + + let mut len_bytes = [0u8; 8]; + file.read_exact(&mut len_bytes).map_err(io)?; + let header_len = u64::from_le_bytes(len_bytes); + if header_len > MAX_HEADER_BYTES || header_len.saturating_add(8) > file_len { + return Err(SafetensorsError::OverBound { + path: path.display().to_string(), + what: "header bytes", + count: header_len, + bound: MAX_HEADER_BYTES.min(file_len), + }); + } + let mut header = vec![0u8; to_usize(header_len, path, "header bytes")?]; + file.read_exact(&mut header).map_err(io)?; + + let data_offset = 8 + header_len; + assert!(data_offset <= file_len, "payload starts inside the file"); + let payload_len = file_len - data_offset; + + let parsed = Self::parse_header(&header, path, payload_len)?; + Ok(Self { + path: path.to_path_buf(), + tensors: parsed, + data_offset, + }) + } + + /// Parse the JSON header into a file-ordered tensor table, validating every + /// entry against the payload region it claims. + fn parse_header( + header: &[u8], + path: &Path, + payload_len: u64, + ) -> Result, SafetensorsError> { + let bad_header = |reason: String| SafetensorsError::BadHeader { + path: path.display().to_string(), + reason, + }; + let json: serde_json::Value = + serde_json::from_slice(header).map_err(|e| bad_header(e.to_string()))?; + let object = json + .as_object() + .ok_or_else(|| bad_header("header is not a JSON object".into()))?; + if object.len() > MAX_TENSORS { + return Err(SafetensorsError::OverBound { + path: path.display().to_string(), + what: "tensor entries", + count: object.len() as u64, + bound: MAX_TENSORS as u64, + }); + } + + let mut tensors = Vec::with_capacity(object.len()); + for (name, value) in object { + // `__metadata__` is a free-form string map, not a tensor. + if name == "__metadata__" { + continue; + } + let bad = |reason: String| SafetensorsError::BadTensor { + path: path.display().to_string(), + name: name.clone(), + reason, + }; + let entry = value + .as_object() + .ok_or_else(|| bad("entry is not an object".into()))?; + let dtype = entry + .get("dtype") + .and_then(serde_json::Value::as_str) + .ok_or_else(|| bad("missing 'dtype'".into()))? + .to_string(); + let element_bytes = + dtype_bytes(&dtype).ok_or_else(|| bad(format!("unsupported dtype '{dtype}'")))?; + let shape: Vec = entry + .get("shape") + .and_then(serde_json::Value::as_array) + .ok_or_else(|| bad("missing 'shape'".into()))? + .iter() + .map(|d| d.as_u64().ok_or_else(|| bad("non-integer dim".into()))) + .collect::>()?; + let offsets = entry + .get("data_offsets") + .and_then(serde_json::Value::as_array) + .ok_or_else(|| bad("missing 'data_offsets'".into()))?; + if offsets.len() != 2 { + return Err(bad(format!("data_offsets has {} entries", offsets.len()))); + } + let start = offsets[0] + .as_u64() + .ok_or_else(|| bad("non-integer data_offset".into()))?; + let end = offsets[1] + .as_u64() + .ok_or_else(|| bad("non-integer data_offset".into()))?; + if end < start || end > payload_len { + return Err(bad(format!( + "data_offsets [{start}, {end}) outside the {payload_len}-byte payload" + ))); + } + // Positive AND negative space: the declared shape must be exactly + // the bytes claimed — a mismatch is a corrupt or mis-declared + // tensor, never something to load and hope about. + let expected = shape + .iter() + .try_fold(element_bytes, |acc: u64, d| acc.checked_mul(*d)) + .ok_or_else(|| bad("shape overflows u64 bytes".into()))?; + if expected != end - start { + return Err(bad(format!( + "shape {shape:?} of {dtype} implies {expected} bytes, header claims {}", + end - start + ))); + } + tensors.push(( + name.clone(), + TensorEntry { + dtype, + shape, + offsets: (start, end), + }, + )); + } + // File order (by payload offset) makes the copy pass a sequential read. + tensors.sort_by_key(|(_, t)| t.offsets.0); + Ok(tensors) + } + + /// Read one tensor's payload bytes. + fn read_payload(&self, entry: &TensorEntry, into: &mut [u8]) -> Result<(), SafetensorsError> { + assert_eq!( + into.len() as u64, + entry.byte_len(), + "read_payload: destination sized to the entry" + ); + let io = |source: std::io::Error| SafetensorsError::Io { + path: self.path.display().to_string(), + source, + }; + let mut file = File::open(&self.path).map_err(io)?; + file.seek(SeekFrom::Start(self.data_offset + entry.offsets.0)) + .map_err(io)?; + file.read_exact(into).map_err(io) + } +} + +/// Bytes per element for a safetensors dtype token, or `None` if unsupported. +#[must_use] +pub fn dtype_bytes(dtype: &str) -> Option { + Some(match dtype { + "BOOL" | "U8" | "I8" | "F8_E4M3" | "F8_E5M2" => 1, + "U16" | "I16" | "F16" | "BF16" => 2, + "U32" | "I32" | "F32" => 4, + "U64" | "I64" | "F64" => 8, + _ => return None, + }) +} + +/// What [`fuse_checkpoint`] should do with a source tensor. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Fuse { + /// Copy through under this (possibly renamed) target name. + Keep(String), + /// Contribute to `target` as member `index` of `count`, stacked along a + /// NEW leading axis — the N:1 case (`mlp.experts.7.gate_proj.weight` is + /// member 7 of the fused `mlp.experts.gate`). + Stack { + target: String, + index: usize, + count: usize, + }, + /// Drop this tensor (present in the checkpoint, unused by the module). + Drop, +} + +/// Every safetensors shard of a checkpoint dir, in index order. +/// +/// A single `model.safetensors` is the one-shard case; otherwise +/// `model.safetensors.index.json`'s weight map names the shards. +pub fn checkpoint_shards(dir: &Path) -> Result, SafetensorsError> { + let single = dir.join("model.safetensors"); + if single.is_file() { + return Ok(vec![single]); + } + let index_path = dir.join("model.safetensors.index.json"); + if !index_path.is_file() { + return Err(SafetensorsError::NoCheckpoint(dir.to_path_buf())); + } + let bytes = std::fs::read(&index_path).map_err(|source| SafetensorsError::Io { + path: index_path.display().to_string(), + source, + })?; + let names = crate::hub::shards_from_index(&bytes, &index_path).map_err(|e| { + SafetensorsError::BadHeader { + path: index_path.display().to_string(), + reason: e.to_string(), + } + })?; + assert!( + !names.is_empty() && names.len() <= MAX_SHARDS, + "shards_from_index bounds the count" + ); + + // The index is a manifest, not evidence. Check every shard is actually on + // disk HERE, so a half-fetched checkpoint is one clear error at planning + // time rather than an `os error 2` raised after the first shards have + // already been opened — and so callers can use this function as the + // "is the checkpoint complete?" question it looks like. + let paths: Vec = names.into_iter().map(|n| dir.join(n)).collect(); + for path in &paths { + if path.is_file() { + continue; + } + let name = path.file_name().unwrap_or_default().to_string_lossy(); + let interrupted = path.with_file_name(format!("{name}.part")).is_file(); + return Err(SafetensorsError::MissingShard { + dir: dir.to_path_buf(), + missing: name.into_owned(), + hint: if interrupted { + " — a .part sibling is present, so the download was interrupted" + } else { + "" + }, + }); + } + debug_assert!( + paths.iter().all(|p| p.is_file()), + "every shard verified present before returning" + ); + Ok(paths) +} + +/// Read every shard of the checkpoint in `dir` and build ONE in-memory +/// safetensors blob, applying `map` to each source tensor name. +/// +/// Fused (`Fuse::Stack`) groups are concatenated in **numeric member order**, +/// gaining a leading `count` axis — `count` tensors of shape `[out, in]` +/// become one `[count, out, in]`. Every group must be exactly complete; +/// a missing or duplicate member is a loud [`SafetensorsError::BadGroup`], +/// raised during planning, before any payload byte is read. +pub fn fuse_checkpoint( + dir: &Path, + map: &dyn Fn(&str) -> Option, +) -> Result, SafetensorsError> { + let mut blob = Vec::new(); + fuse_into(dir, map, &mut blob)?; + Ok(blob) +} + +/// [`fuse_checkpoint`] straight to a file, never holding the payload in RAM. +/// +/// This is the variant a real checkpoint wants. The in-memory form needs the +/// whole fused payload resident — 13.8 GB for OLMoE-1B-7B — *on top of* the +/// model the load then builds (~28 GB at f32), and that sum is what a 128 GB +/// box with other tenants actually fails to satisfy. Writing to disk trades +/// the spike for temp space and lets `SafetensorsStore::from_file` page the +/// weights in as it needs them. +pub fn fuse_checkpoint_to_file( + dir: &Path, + map: &dyn Fn(&str) -> Option, + out: &Path, +) -> Result { + let io = |source: std::io::Error| SafetensorsError::Io { + path: out.display().to_string(), + source, + }; + let file = File::create(out).map_err(io)?; + let mut sink = std::io::BufWriter::with_capacity(1 << 20, file); + let written = fuse_into(dir, map, &mut sink)?; + sink.flush().map_err(io)?; + sink.into_inner() + .map_err(|e| SafetensorsError::Io { + path: out.display().to_string(), + source: e.into_error(), + })? + .sync_all() + .map_err(io)?; + Ok(written) +} + +/// The shared fuse: plan, then stream header + payload into `sink` in output +/// order, one part at a time. +/// +/// Writes are strictly ascending because `plan` carries tensors in output +/// order and each tensor's `parts` are already in destination order — so this +/// never needs the payload addressable at once, only ONE part at a time. +fn fuse_into( + dir: &Path, + map: &dyn Fn(&str) -> Option, + sink: &mut W, +) -> Result { + let shard_paths = checkpoint_shards(dir)?; + let headers = shard_paths + .iter() + .map(|p| SafetensorsHeader::open(p)) + .collect::, _>>()?; + + let plan = plan_output(&headers, map)?; + let total: u64 = plan.iter().map(|p| p.len).sum(); + if total > MAX_FUSED_BYTES { + return Err(SafetensorsError::OverBound { + path: dir.display().to_string(), + what: "fused payload bytes", + count: total, + bound: MAX_FUSED_BYTES, + }); + } + + // Header first: names, dtypes, shapes, and the contiguous offsets the + // copy pass will fill. + let mut header = String::from("{"); + for (i, p) in plan.iter().enumerate() { + if i > 0 { + header.push(','); + } + let json_name = + serde_json::to_string(&p.name).map_err(|e| SafetensorsError::BadTensor { + path: dir.display().to_string(), + name: p.name.clone(), + reason: format!("name is not encodable as JSON: {e}"), + })?; + header.push_str(&format!( + "{json_name}:{{\"dtype\":\"{}\",\"shape\":{:?},\"data_offsets\":[{},{}]}}", + p.dtype, + p.shape, + p.start, + p.start + p.len, + )); + } + header.push('}'); + + let io = |source: std::io::Error| SafetensorsError::Io { + path: dir.display().to_string(), + source, + }; + sink.write_all(&(header.len() as u64).to_le_bytes()) + .map_err(io)?; + sink.write_all(header.as_bytes()).map_err(io)?; + + // One reusable buffer, sized to the largest single part rather than to the + // payload: for a 64-expert bank that is one expert projection (~4 MB), not + // the 13.8 GB whole. + let widest = plan + .iter() + .flat_map(|p| p.parts.iter()) + .map(|(_, entry, _)| entry.byte_len()) + .max() + .unwrap_or(0); + if widest > MAX_PART_BYTES { + return Err(SafetensorsError::OverBound { + path: dir.display().to_string(), + what: "single tensor bytes", + count: widest, + bound: MAX_PART_BYTES, + }); + } + let mut buf = vec![0u8; to_usize(widest, dir, "single tensor bytes")?]; + + let mut written = 0u64; + for p in &plan { + debug_assert_eq!(written, p.start, "tensors are written in output order"); + for (shard, entry, dst_within) in &p.parts { + debug_assert_eq!( + written, + p.start + dst_within, + "parts are written in destination order" + ); + let len = to_usize(entry.byte_len(), dir, "member byte length")?; + let slot = &mut buf[..len]; + headers[*shard].read_payload(entry, slot)?; + sink.write_all(slot).map_err(io)?; + written += len as u64; + } + } + assert_eq!( + written, total, + "every planned byte was written exactly once" + ); + Ok(written) +} + +/// Plan the output layout and validate it. Nothing here reads payload bytes — +/// a malformed checkpoint fails before the expensive pass. +fn plan_output( + headers: &[SafetensorsHeader], + map: &dyn Fn(&str) -> Option, +) -> Result, SafetensorsError> { + assert!( + !headers.is_empty(), + "planning needs at least one shard header" + ); + assert!( + headers.len() <= MAX_SHARDS, + "shard count is bounded before planning" + ); + + // Collect, in first-seen order, what each target is built from. + let mut order: Vec = Vec::new(); + let mut groups: HashMap = HashMap::new(); + + for (shard, header) in headers.iter().enumerate() { + for (name, entry) in &header.tensors { + let action = map(name).ok_or_else(|| SafetensorsError::BadTensor { + path: header.path.display().to_string(), + name: name.clone(), + reason: "unmapped tensor name".into(), + })?; + let (target, member, count) = match action { + Fuse::Drop => continue, + Fuse::Keep(target) => (target, 0usize, 1usize), + Fuse::Stack { + target, + index, + count, + } => { + if index >= count { + return Err(SafetensorsError::BadGroup { + target, + reason: format!("member index {index} outside 0..{count}"), + }); + } + (target, index, count) + } + }; + let group = groups.entry(target.clone()).or_insert_with(|| { + order.push(target.clone()); + Group { + count, + stacked: count > 1, + members: Vec::new(), + } + }); + if group.count != count { + return Err(SafetensorsError::BadGroup { + target, + reason: format!("member count disagrees ({} vs {count})", group.count), + }); + } + if group.members.iter().any(|(i, _, _)| *i == member) { + return Err(SafetensorsError::BadGroup { + target, + reason: format!("duplicate member {member}"), + }); + } + group.members.push((member, shard, entry.clone())); + } + } + + let mut plan = Vec::with_capacity(order.len()); + let mut cursor = 0u64; + for target in order { + // `order` is pushed only from `or_insert_with`, so every target in it + // has a group and none repeats — an `Err` here would mean that + // invariant broke, not that the checkpoint is bad. + let Some(mut group) = groups.remove(&target) else { + return Err(SafetensorsError::BadGroup { + target, + reason: "planned target has no collected members".into(), + }); + }; + // THE load-bearing check: exactly `count` members, ids 0..count. + if group.members.len() != group.count { + return Err(SafetensorsError::BadGroup { + target, + reason: format!( + "{} of {} members present — a short group would load clean and compute wrong", + group.members.len(), + group.count + ), + }); + } + // Numeric member order, NOT the lexicographic order the names imply + // (…experts.10 sorts before …experts.2 as text). + group.members.sort_by_key(|(i, _, _)| *i); + debug_assert!( + group + .members + .iter() + .enumerate() + .all(|(i, (m, _, _))| i == *m), + "a complete, duplicate-free group is exactly 0..count once sorted" + ); + + let first = &group.members[0].2; + for (_, _, entry) in &group.members { + if entry.dtype != first.dtype || entry.shape != first.shape { + return Err(SafetensorsError::BadGroup { + target, + reason: format!( + "members disagree on layout ({} {:?} vs {} {:?})", + first.dtype, first.shape, entry.dtype, entry.shape + ), + }); + } + } + + let shape = if group.stacked { + std::iter::once(group.count as u64) + .chain(first.shape.iter().copied()) + .collect() + } else { + first.shape.clone() + }; + let member_bytes = first.byte_len(); + let len = member_bytes * group.count as u64; + let parts = group + .members + .iter() + .enumerate() + .map(|(slot, (_, shard, entry))| (*shard, entry.clone(), slot as u64 * member_bytes)) + .collect(); + plan.push(PlannedNamed { + name: target, + dtype: first.dtype.clone(), + shape, + start: cursor, + len, + parts, + }); + cursor += len; + } + assert!( + !plan.is_empty(), + "a checkpoint dir with shards yields at least one output tensor" + ); + Ok(plan) +} + +/// Accumulator for one output tensor while planning. +struct Group { + count: usize, + stacked: bool, + members: Vec<(usize, usize, TensorEntry)>, +} + +/// A planned output tensor: where it lands in the blob and what fills it. +struct PlannedNamed { + name: String, + dtype: String, + shape: Vec, + start: u64, + len: u64, + /// Sources, in destination order: (shard index, entry, byte offset within + /// this tensor). One part for `Keep`, `count` parts for `Stack`. + parts: Vec<(usize, TensorEntry, u64)>, +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Build a synthetic safetensors file: `(name, dtype, shape, bytes)`. + fn write_st(path: &Path, tensors: &[(&str, &str, Vec, Vec)]) { + let mut header = String::from("{"); + let mut data: Vec = Vec::new(); + for (i, (name, dtype, shape, bytes)) in tensors.iter().enumerate() { + if i > 0 { + header.push(','); + } + let start = data.len(); + data.extend_from_slice(bytes); + header.push_str(&format!( + "{:?}:{{\"dtype\":\"{dtype}\",\"shape\":{shape:?},\"data_offsets\":[{start},{}]}}", + name, + data.len() + )); + } + header.push('}'); + let mut blob = (header.len() as u64).to_le_bytes().to_vec(); + blob.extend_from_slice(header.as_bytes()); + blob.extend_from_slice(&data); + std::fs::write(path, blob).unwrap(); + } + + /// A fresh empty scratch dir under the OS temp root. + fn scratch(tag: &str) -> PathBuf { + let dir = std::env::temp_dir().join(format!("mummu_st_{tag}")); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + dir + } + + /// One f32 tensor's little-endian bytes. + fn f32s(values: &[f32]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() + } + + /// Read a tensor back out of a fused blob by name. + fn read_back(blob: &[u8], name: &str) -> (String, Vec, Vec) { + let header_len = u64::from_le_bytes(blob[..8].try_into().unwrap()) as usize; + let json: serde_json::Value = serde_json::from_slice(&blob[8..8 + header_len]).unwrap(); + let entry = &json[name]; + let dtype = entry["dtype"].as_str().unwrap().to_string(); + let shape: Vec = entry["shape"] + .as_array() + .unwrap() + .iter() + .map(|d| d.as_u64().unwrap()) + .collect(); + let start = entry["data_offsets"][0].as_u64().unwrap() as usize; + let end = entry["data_offsets"][1].as_u64().unwrap() as usize; + let base = 8 + header_len; + let (words, rest) = blob[base + start..base + end].as_chunks::<4>(); + assert!(rest.is_empty(), "f32 payload is a whole number of words"); + let values = words.iter().copied().map(f32::from_le_bytes).collect(); + (dtype, shape, values) + } + + #[test] + fn header_parses_shape_dtype_and_offsets() { + let dir = scratch("header"); + let path = dir.join("model.safetensors"); + write_st( + &path, + &[ + ("a", "F32", vec![2, 2], f32s(&[1.0, 2.0, 3.0, 4.0])), + ("b", "F32", vec![3], f32s(&[5.0, 6.0, 7.0])), + ], + ); + let h = SafetensorsHeader::open(&path).unwrap(); + assert_eq!(h.tensors.len(), 2); + assert_eq!(h.tensors[0].0, "a"); + assert_eq!(h.tensors[0].1.shape, vec![2, 2]); + assert_eq!(h.tensors[0].1.byte_len(), 16); + assert_eq!(h.tensors[0].1.element_count(), 4); + assert_eq!(h.tensors[1].1.offsets, (16, 28)); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn header_rejects_a_shape_that_disagrees_with_its_bytes() { + let dir = scratch("badshape"); + let path = dir.join("model.safetensors"); + // Declares [4] f32 (16 B) but only 8 B of payload are claimed. + let mut blob = Vec::new(); + let header = r#"{"a":{"dtype":"F32","shape":[4],"data_offsets":[0,8]}}"#.to_string(); + blob.extend_from_slice(&(header.len() as u64).to_le_bytes()); + blob.extend_from_slice(header.as_bytes()); + blob.extend_from_slice(&[0u8; 8]); + std::fs::write(&path, blob).unwrap(); + + let err = SafetensorsHeader::open(&path).unwrap_err(); + assert!( + matches!(&err, SafetensorsError::BadTensor { reason, .. } + if reason.contains("implies 16 bytes")), + "got {err}" + ); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn header_rejects_an_unsupported_dtype() { + let dir = scratch("baddtype"); + let path = dir.join("model.safetensors"); + write_st(&path, &[("a", "COMPLEX128", vec![1], vec![0u8; 16])]); + let err = SafetensorsHeader::open(&path).unwrap_err(); + assert!( + matches!(&err, SafetensorsError::BadTensor { reason, .. } + if reason.contains("unsupported dtype")), + "got {err}" + ); + std::fs::remove_dir_all(&dir).unwrap(); + } + + /// THE ordering test. Members are stacked in NUMERIC order; the names sort + /// lexicographically as 0, 1, 10, 2, … so a text sort silently permutes + /// the expert bank — a model that loads clean and computes wrong. + #[test] + fn stack_orders_members_numerically_not_lexicographically() { + let dir = scratch("order"); + let count = 12usize; + let tensors: Vec<(String, &str, Vec, Vec)> = (0..count) + .map(|i| { + ( + format!("mlp.experts.{i}.gate_proj.weight"), + "F32", + vec![1], + // Expert i holds exactly the value i. + f32s(&[i as f32]), + ) + }) + .collect(); + let refs: Vec<(&str, &str, Vec, Vec)> = tensors + .iter() + .map(|(n, d, s, b)| (n.as_str(), *d, s.clone(), b.clone())) + .collect(); + write_st(&dir.join("model.safetensors"), &refs); + + let blob = fuse_checkpoint(&dir, &|name| { + let idx: usize = name + .strip_prefix("mlp.experts.")? + .strip_suffix(".gate_proj.weight")? + .parse() + .ok()?; + Some(Fuse::Stack { + target: "mlp.experts.gate".into(), + index: idx, + count, + }) + }) + .unwrap(); + + let (dtype, shape, values) = read_back(&blob, "mlp.experts.gate"); + assert_eq!(dtype, "F32"); + assert_eq!(shape, vec![count as u64, 1], "gains a leading expert axis"); + // Slot i must hold expert i — the whole point. + assert_eq!( + values, + (0..count).map(|i| i as f32).collect::>(), + "experts must stack in numeric order (lexicographic would give 0,1,10,11,2,…)" + ); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn stack_rejects_a_short_group() { + let dir = scratch("short"); + // Declare 4 members but ship only 3. + write_st( + &dir.join("model.safetensors"), + &[ + ("e.0", "F32", vec![1], f32s(&[0.0])), + ("e.1", "F32", vec![1], f32s(&[1.0])), + ("e.2", "F32", vec![1], f32s(&[2.0])), + ], + ); + let err = fuse_checkpoint(&dir, &|name| { + let idx: usize = name.strip_prefix("e.")?.parse().ok()?; + Some(Fuse::Stack { + target: "e".into(), + index: idx, + count: 4, + }) + }) + .unwrap_err(); + assert!( + matches!(&err, SafetensorsError::BadGroup { reason, .. } + if reason.contains("3 of 4 members")), + "got {err}" + ); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn stack_rejects_duplicate_members_and_layout_disagreement() { + let dir = scratch("dup"); + write_st( + &dir.join("model.safetensors"), + &[ + ("a", "F32", vec![1], f32s(&[0.0])), + ("b", "F32", vec![1], f32s(&[1.0])), + ], + ); + // Both map to member 0 of the same target. + let err = fuse_checkpoint(&dir, &|_| { + Some(Fuse::Stack { + target: "t".into(), + index: 0, + count: 2, + }) + }) + .unwrap_err(); + assert!( + matches!(&err, SafetensorsError::BadGroup { reason, .. } + if reason.contains("duplicate member 0")), + "got {err}" + ); + + // Same group, different shapes. + let dir2 = scratch("layout"); + write_st( + &dir2.join("model.safetensors"), + &[ + ("a", "F32", vec![1], f32s(&[0.0])), + ("b", "F32", vec![2], f32s(&[1.0, 2.0])), + ], + ); + let err = fuse_checkpoint(&dir2, &|name| { + Some(Fuse::Stack { + target: "t".into(), + index: usize::from(name == "b"), + count: 2, + }) + }) + .unwrap_err(); + assert!( + matches!(&err, SafetensorsError::BadGroup { reason, .. } + if reason.contains("disagree on layout")), + "got {err}" + ); + std::fs::remove_dir_all(&dir).unwrap(); + std::fs::remove_dir_all(&dir2).unwrap(); + } + + #[test] + fn unmapped_tensor_is_a_loud_error_and_drop_is_explicit() { + let dir = scratch("unmapped"); + write_st( + &dir.join("model.safetensors"), + &[ + ("keep", "F32", vec![1], f32s(&[1.0])), + ("junk", "F32", vec![1], f32s(&[2.0])), + ], + ); + // No mapping for "junk" -> loud. + let err = fuse_checkpoint(&dir, &|name| { + (name == "keep").then(|| Fuse::Keep("keep".into())) + }) + .unwrap_err(); + assert!( + matches!(&err, SafetensorsError::BadTensor { reason, .. } + if reason.contains("unmapped tensor name")), + "got {err}" + ); + // Explicitly dropping it is fine, and it leaves the blob. + let blob = fuse_checkpoint(&dir, &|name| { + Some(if name == "keep" { + Fuse::Keep("keep".into()) + } else { + Fuse::Drop + }) + }) + .unwrap(); + let (_, shape, values) = read_back(&blob, "keep"); + assert_eq!((shape, values), (vec![1], vec![1.0])); + let header_len = u64::from_le_bytes(blob[..8].try_into().unwrap()) as usize; + let header = std::str::from_utf8(&blob[8..8 + header_len]).unwrap(); + assert!(!header.contains("junk"), "dropped tensor is absent"); + std::fs::remove_dir_all(&dir).unwrap(); + } + + /// A group whose members live in DIFFERENT shards still fuses correctly — + /// a real 3-shard checkpoint can split a layer across a boundary. + #[test] + fn shards_are_discovered_and_fused_across_boundaries() { + let dir = scratch("shards"); + write_st( + &dir.join("model-00001-of-00002.safetensors"), + &[("e.1", "F32", vec![1], f32s(&[11.0]))], + ); + write_st( + &dir.join("model-00002-of-00002.safetensors"), + &[("e.0", "F32", vec![1], f32s(&[10.0]))], + ); + std::fs::write( + dir.join("model.safetensors.index.json"), + br#"{"weight_map":{"e.1":"model-00001-of-00002.safetensors", + "e.0":"model-00002-of-00002.safetensors"}}"#, + ) + .unwrap(); + + let shards = checkpoint_shards(&dir).unwrap(); + assert_eq!(shards.len(), 2, "both shards discovered from the index"); + + let blob = fuse_checkpoint(&dir, &|name| { + let idx: usize = name.strip_prefix("e.")?.parse().ok()?; + Some(Fuse::Stack { + target: "e".into(), + index: idx, + count: 2, + }) + }) + .unwrap(); + let (_, shape, values) = read_back(&blob, "e"); + assert_eq!(shape, vec![2, 1]); + // Member 0 came from shard 2, member 1 from shard 1 — order is by + // member index, never by shard order. + assert_eq!(values, vec![10.0, 11.0]); + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn a_dir_with_no_checkpoint_is_a_loud_error() { + let dir = scratch("empty"); + assert!(matches!( + checkpoint_shards(&dir), + Err(SafetensorsError::NoCheckpoint(_)) + )); + std::fs::remove_dir_all(&dir).unwrap(); + } + + /// An index is a manifest, not evidence. A half-fetched checkpoint must + /// fail HERE — before any shard is opened — or it surfaces as a bare + /// `os error 2` from inside the load, and any caller using + /// `checkpoint_shards(..).is_ok()` as "is this checkpoint complete?" + /// silently believes an interrupted download is ready to use. + #[test] + fn an_index_naming_a_missing_shard_is_a_loud_error_not_a_late_os_error() { + let dir = scratch("missing_shard"); + write_st( + &dir.join("model-00001-of-00002.safetensors"), + &[("e.0", "F32", vec![1], f32s(&[10.0]))], + ); + std::fs::write( + dir.join("model.safetensors.index.json"), + br#"{"weight_map":{"e.0":"model-00001-of-00002.safetensors", + "e.1":"model-00002-of-00002.safetensors"}}"#, + ) + .unwrap(); + + // Shard 2 is absent entirely: named, no hint. + match checkpoint_shards(&dir) { + Err(SafetensorsError::MissingShard { missing, hint, .. }) => { + assert_eq!(missing, "model-00002-of-00002.safetensors"); + assert!( + hint.is_empty(), + "no .part sibling, so no interrupted-download hint" + ); + } + other => panic!("expected MissingShard, got {other:?}"), + } + + // The same shard as an interrupted download: the error says so, which + // is the difference between "re-fetch me" and "this repo is broken". + std::fs::write( + dir.join("model-00002-of-00002.safetensors.part"), + b"partial", + ) + .unwrap(); + match checkpoint_shards(&dir) { + Err(SafetensorsError::MissingShard { hint, .. }) => { + assert!( + hint.contains("interrupted"), + "hint names the .part sibling: {hint:?}" + ); + } + other => panic!("expected MissingShard, got {other:?}"), + } + + std::fs::remove_dir_all(&dir).unwrap(); + } + + /// The streaming fuse and the in-memory fuse must be the SAME bytes — + /// otherwise "load a big model from a file" and "load a small one from + /// RAM" are two different importers, and only one of them is tested. + #[test] + fn fusing_to_a_file_is_byte_identical_to_fusing_in_memory() { + let dir = scratch("fuse_to_file"); + write_st( + &dir.join("model-00001-of-00002.safetensors"), + &[ + ("e.1", "F32", vec![2], f32s(&[11.0, 12.0])), + ("keep", "F32", vec![1], f32s(&[7.0])), + ], + ); + write_st( + &dir.join("model-00002-of-00002.safetensors"), + &[("e.0", "F32", vec![2], f32s(&[10.0, 9.0]))], + ); + std::fs::write( + dir.join("model.safetensors.index.json"), + br#"{"weight_map":{"e.1":"model-00001-of-00002.safetensors", + "keep":"model-00001-of-00002.safetensors", + "e.0":"model-00002-of-00002.safetensors"}}"#, + ) + .unwrap(); + + let map = |name: &str| -> Option { + if name == "keep" { + return Some(Fuse::Keep("keep".into())); + } + let idx: usize = name.strip_prefix("e.")?.parse().ok()?; + Some(Fuse::Stack { + target: "e".into(), + index: idx, + count: 2, + }) + }; + + let in_memory = fuse_checkpoint(&dir, &map).unwrap(); + let out = dir.join("fused.safetensors"); + let written = fuse_checkpoint_to_file(&dir, &map, &out).unwrap(); + let on_disk = std::fs::read(&out).unwrap(); + + assert_eq!( + in_memory, on_disk, + "the two fuse paths must agree byte for byte" + ); + // The return value counts PAYLOAD bytes, so the file is that plus the + // 8-byte length prefix and the JSON header. + assert!( + written < on_disk.len() as u64, + "payload {written} B sits inside a {} B file", + on_disk.len() + ); + + // And the fused content is still correct, not merely self-consistent. + let (_, shape, values) = read_back(&on_disk, "e"); + assert_eq!(shape, vec![2, 2]); + assert_eq!(values, vec![10.0, 9.0, 11.0, 12.0], "numeric member order"); + + std::fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn dtype_bytes_covers_the_float_and_int_tokens() { + assert_eq!(dtype_bytes("BF16"), Some(2)); + assert_eq!(dtype_bytes("F32"), Some(4)); + assert_eq!(dtype_bytes("I64"), Some(8)); + assert_eq!(dtype_bytes("BOOL"), Some(1)); + assert_eq!(dtype_bytes("NOPE"), None); + } +} diff --git a/crates/mummu/tests/real_olmoe_safetensors.rs b/crates/mummu/tests/real_olmoe_safetensors.rs new file mode 100644 index 0000000..9e130a6 --- /dev/null +++ b/crates/mummu/tests/real_olmoe_safetensors.rs @@ -0,0 +1,184 @@ +//! REAL-WEIGHTS proof for the OLMoE **HF safetensors** import path (P2). +//! +//! The GGUF path ships the 64 experts already fused (`ffn_*_exps`); the HF +//! checkpoint stores each expert separately across three shards. This suite +//! proves the fusing loader lands on the same model: +//! +//! 1. the sharded checkpoint fetches through the registry spec; +//! 2. `load_from_dir` fuses 16 layers x 64 experts x 3 projections and +//! checked-loads with zero missing params; +//! 3. it decodes coherently on the CPU backend; +//! 4. **the fused expert bank is bit-exact against the same expert read +//! independently out of the raw shard bytes** — the check that a silent +//! mis-ordering (expert 10 in slot 2) could not survive. +//! +//! Ignored by default (13.8 GB of weights, ~28 GB of RAM once loaded). Run: +//! +//! ```text +//! MUMMU_HUB_DEST=C:\Users\me\.cache\mummu-models \ +//! cargo test -p mummu --test real_olmoe_safetensors -- --ignored --nocapture +//! ``` + +use std::path::PathBuf; + +use mummu::backend::Cpu; +use mummu::models::{CausalLm, olmoe}; +use mummu::registry; +use mummu::safetensors::{SafetensorsHeader, checkpoint_shards}; + +type Dev = burn::tensor::Device; + +const SPEC: &str = "olmoe-1b-7b-0125-instruct"; + +/// Where the checkpoint lives, fetching it through the registry if needed. +fn checkpoint_dir() -> Option { + let root = PathBuf::from(std::env::var_os("MUMMU_HUB_DEST")?); + let spec = registry::catalog() + .into_iter() + .find(|s| s.name == SPEC) + .expect("the catalog carries the OLMoE safetensors spec"); + let dir = spec.dir(&root); + if dir.join("config.json").is_file() && checkpoint_shards(&dir).is_ok() { + return Some(dir); + } + eprintln!("fetching {SPEC} (~13.8 GB) into {}", dir.display()); + spec.fetch(&root, &mut |p: mummu::hub::Progress| { + if let Some(total) = p.total_bytes.filter(|t| *t > 0) { + eprint!("\r {} {:>3} %", p.file, 100 * p.received_bytes / total); + } + }) + .expect("registry fetch of the sharded checkpoint succeeds"); + eprintln!(); + Some(dir) +} + +/// The whole path: sharded fetch -> fuse -> checked load -> real decode. +#[test] +#[ignore = "needs 13.8 GB of weights + ~28 GB RAM (MUMMU_HUB_DEST)"] +fn olmoe_safetensors_fuses_and_decodes_on_cpu() { + let Some(dir) = checkpoint_dir() else { + eprintln!("set MUMMU_HUB_DEST to run this test"); + return; + }; + let shards = checkpoint_shards(&dir).expect("shards discovered"); + println!("checkpoint: {} shard(s)", shards.len()); + assert!( + shards.len() > 1, + "the 1B-7B ships sharded — this exercises the multi-shard path" + ); + + let device = Dev::default(); + let started = std::time::Instant::now(); + let loaded = olmoe::load_from_dir::(&dir, &device).expect("fused safetensors load"); + println!( + "loaded {} layers x {} experts in {:.1} s", + loaded.config.num_hidden_layers, + loaded.config.num_experts, + started.elapsed().as_secs_f32() + ); + assert_eq!(loaded.config.num_experts, 64); + assert_eq!(loaded.config.num_experts_per_tok, 8); + // The sibling tokenizer_config.json is surfaced when the dir has one, and + // is legitimately absent otherwise (a registry fetch pulls config.json + + // tokenizer.json + weights only) — the loader must reflect the dir, not + // invent a config. + assert_eq!( + loaded.tokenizer_config.is_some(), + dir.join("tokenizer_config.json").is_file(), + "surfaced tokenizer_config must match what the dir actually ships" + ); + + // Liveness: a real forward, checked for finiteness / width / spread. + let probe: Vec = vec![100, 200, 300, 400]; + let smoke = loaded + .sanity_check(&probe, loaded.config.vocab_size, &device) + .expect("a fused load computes a live distribution"); + println!("sanity: top {} spread {:.1}", smoke.top_id, smoke.spread); + + let out = loaded + .greedy_generate(&probe, 8, &device) + .expect("greedy decode runs"); + println!("decoded {} tokens: {out:?}", out.len()); + assert!(!out.is_empty(), "the model emits tokens"); +} + +/// The ordering proof. Read expert 37's `gate_proj` for layer 5 straight out +/// of the raw shard bytes, then read slot 37 of the fused bank the loader +/// built, and require them **bit-identical**. A mis-ordered fuse (the +/// lexicographic trap) puts a different expert in that slot and fails here. +#[test] +#[ignore = "needs 13.8 GB of weights (MUMMU_HUB_DEST)"] +fn fused_expert_slot_is_bit_exact_against_the_raw_shard_bytes() { + let Some(dir) = checkpoint_dir() else { + eprintln!("set MUMMU_HUB_DEST to run this test"); + return; + }; + const LAYER: usize = 5; + const EXPERT: usize = 37; + let source_name = format!("model.layers.{LAYER}.mlp.experts.{EXPERT}.gate_proj.weight"); + + // Independent read: find the tensor in whichever shard holds it and pull + // its raw bytes, using nothing the loader used. + let mut truth: Option<(Vec, Vec)> = None; + for shard in checkpoint_shards(&dir).expect("shards") { + let header = SafetensorsHeader::open(&shard).expect("shard header parses"); + let Some((_, entry)) = header.tensors.iter().find(|(n, _)| *n == source_name) else { + continue; + }; + let mut file = std::fs::File::open(&shard).expect("shard opens"); + use std::io::{Read, Seek, SeekFrom}; + file.seek(SeekFrom::Start(header.data_offset + entry.offsets.0)) + .expect("seek"); + let mut bytes = vec![0u8; entry.byte_len() as usize]; + file.read_exact(&mut bytes).expect("read"); + truth = Some((entry.shape.clone(), bytes)); + break; + } + let (shape, truth_bytes) = truth.expect("the checkpoint carries that expert"); + println!( + "raw {source_name}: shape {shape:?}, {} B", + truth_bytes.len() + ); + + // What the loader produced, read back out of the loaded module. + let device = Dev::default(); + let loaded = olmoe::load_from_dir::(&dir, &device).expect("fused safetensors load"); + let bank = loaded.model.layers[LAYER].mlp.experts.gate.val(); + let dims = bank.dims(); + println!("fused bank dims {dims:?}"); + assert_eq!(dims[0], loaded.config.num_experts, "leading expert axis"); + assert_eq!( + [dims[1] as u64, dims[2] as u64], + [shape[0], shape[1]], + "each slot keeps the source [out, in] shape" + ); + + let slot: Vec = bank + .slice(EXPERT..EXPERT + 1) + .into_data() + .into_vec() + .expect("readback"); + // The checkpoint is bf16; the loader cast it to the backend float (f32). + // Compare against the same cast applied to the truth bytes: bf16 -> f32 + // is exact (it is a 16-bit truncation of f32), so this is bit-equality. + let (halves, rest) = truth_bytes.as_chunks::<2>(); + assert!(rest.is_empty(), "bf16 payload is a whole number of halves"); + let expected: Vec = halves + .iter() + .map(|c| f32::from_bits(u32::from(u16::from_le_bytes(*c)) << 16)) + .collect(); + assert_eq!(slot.len(), expected.len(), "same element count"); + let mismatches = slot + .iter() + .zip(&expected) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + println!( + "expert {EXPERT} slot: {} values, {mismatches} bit-mismatches", + slot.len() + ); + assert_eq!( + mismatches, 0, + "the fused slot must be bit-exact against the raw shard bytes" + ); +} From e90052636d783fb59dad61cd884e8d926b4484db Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 20 Aug 2026 11:52:47 -0500 Subject: [PATCH 5/5] refactor: split the collect phase out of plan_output, fix two stale docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `plan_output` was ~150 lines doing two separable jobs. `collect_groups` is now the first pass — walk every shard, record in first-seen order what each target is built from, and reject a checkpoint's bad *claims* (unmapped name, member index outside 0..count, members disagreeing on group size, a duplicate member). `plan_output` keeps the second pass: lay the groups out and apply THE completeness check, which can only run once every shard has been seen, because a group's members may be split across shards in any order. Bodies now 63 and 93 lines, down from one 150-line function. Behaviour is unchanged and the existing suite says so — the group-validation tests (short group, duplicate member, layout disagreement, numeric-vs-lexicographic ordering) all cover the seam and still pass. Also two doc comments that my own change had made wrong: `models::olmoe` and `MAX_FUSED_BYTES` still pointed at `fuse_checkpoint` after `load_from_dir` moved to `fuse_checkpoint_to_file`, and the module header described the result as necessarily "in-memory". A doc that names the wrong function is worse than no doc. 232 lib tests green, clippy --all-targets clean, fmt clean. Co-Authored-By: Claude Opus 4.8 --- crates/mummu/src/models/olmoe.rs | 10 +++-- crates/mummu/src/safetensors.rs | 74 ++++++++++++++++++++++---------- 2 files changed, 59 insertions(+), 25 deletions(-) diff --git a/crates/mummu/src/models/olmoe.rs b/crates/mummu/src/models/olmoe.rs index c20d6f0..7884bdf 100644 --- a/crates/mummu/src/models/olmoe.rs +++ b/crates/mummu/src/models/olmoe.rs @@ -17,9 +17,13 @@ //! (`ffn_*_exps`) in exactly the layout `MoeExperts` holds; the HF safetensors //! checkpoint stores each expert separately //! (`mlp.experts.{i}.{gate,up,down}_proj.weight`) and is sharded, so -//! [`load_from_dir`] runs it through [`crate::safetensors::fuse_checkpoint`], -//! which reads every shard and stacks each 64-member expert group into one -//! `[experts, out, in]` tensor before the ordinary checked-load pipeline. +//! [`load_from_dir`] runs it through +//! [`crate::safetensors::fuse_checkpoint_to_file`], which reads every shard and +//! stacks each 64-member expert group into one `[experts, out, in]` tensor +//! before the ordinary checked-load pipeline. It fuses to a temp file rather +//! than to RAM deliberately: the in-memory twin would need the whole payload +//! resident (13.8 GB for the 1B-7B) on top of the ~28 GB f32 model the load +//! then builds. use std::path::{Path, PathBuf}; diff --git a/crates/mummu/src/safetensors.rs b/crates/mummu/src/safetensors.rs index 3b2da3a..66ca655 100644 --- a/crates/mummu/src/safetensors.rs +++ b/crates/mummu/src/safetensors.rs @@ -14,13 +14,16 @@ //! fused `[experts, out, in]` param — exactly the ggml `ffn_*_exps` layout //! a GGUF ships pre-fused. //! -//! [`fuse_checkpoint`] reads every shard's header, plans the output layout, +//! The fuse reads every shard's header, plans the output layout, //! validates it, and only then copies payload bytes — so a checkpoint missing //! expert 37 fails before a single weight byte is read, rather than loading -//! clean and computing wrong. The result is an ordinary in-memory safetensors -//! blob that goes through the SAME `SafetensorsStore::from_bytes` + -//! adapter-chain + `load_checked` pipeline as every other import path (the -//! GGUF path's `dequant_to_safetensors` is the precedent). +//! clean and computing wrong. The result is an ordinary safetensors blob that +//! goes through the SAME `SafetensorsStore` + adapter-chain + `load_checked` +//! pipeline as every other import path (the GGUF path's +//! `dequant_to_safetensors` is the precedent). Two ways out, byte-identical +//! and unit-tested as such: [`fuse_checkpoint`] returns the blob in memory, +//! and [`fuse_checkpoint_to_file`] streams it to disk for checkpoints too big +//! to hold twice. //! //! Source dtypes are preserved verbatim — the bf16→backend-float cast stays //! where it already lives, in `CastFloatAdapter` on the load pipeline. @@ -38,9 +41,11 @@ const MAX_HEADER_BYTES: u64 = 64 << 20; /// at 16 layers already declares ~3 000; 1M is a runaway index. const MAX_TENSORS: usize = 1 << 20; -/// Largest fused payload [`fuse_checkpoint`] will build in RAM. Matches the -/// GGUF path's ceiling (the reference machine has 128 GB) — note this blob -/// keeps the SOURCE dtype, so a bf16 checkpoint costs half its f32 footprint. +/// Largest fused payload either fuse will produce. Matches the GGUF path's +/// ceiling (the reference machine has 128 GB) — note the payload keeps the +/// SOURCE dtype, so a bf16 checkpoint costs half its f32 footprint. Only +/// [`fuse_checkpoint`] holds this much at once; [`fuse_checkpoint_to_file`] +/// streams it and never buffers more than one tensor. const MAX_FUSED_BYTES: u64 = 48 << 30; /// Largest SINGLE tensor (or fused member) the streaming fuse will buffer. @@ -538,22 +543,18 @@ fn fuse_into( Ok(written) } -/// Plan the output layout and validate it. Nothing here reads payload bytes — -/// a malformed checkpoint fails before the expensive pass. -fn plan_output( +/// First pass: walk every shard and record, in first-seen order, what each +/// output target is built from. +/// +/// This is where a checkpoint's *claims* are checked — an unmapped name, a +/// member index outside `0..count`, two members disagreeing about the group +/// size, or the same member twice. What it does NOT check is completeness; +/// that is [`plan_output`]'s job, once every shard has been seen. +fn collect_groups( headers: &[SafetensorsHeader], map: &dyn Fn(&str) -> Option, -) -> Result, SafetensorsError> { - assert!( - !headers.is_empty(), - "planning needs at least one shard header" - ); - assert!( - headers.len() <= MAX_SHARDS, - "shard count is bounded before planning" - ); - - // Collect, in first-seen order, what each target is built from. +) -> Result<(Vec, HashMap), SafetensorsError> { + debug_assert!(!headers.is_empty(), "collecting needs a shard header"); let mut order: Vec = Vec::new(); let mut groups: HashMap = HashMap::new(); @@ -605,6 +606,35 @@ fn plan_output( } } + debug_assert_eq!( + order.len(), + groups.len(), + "every ordered target has exactly one group" + ); + Ok((order, groups)) +} + +/// Second pass: lay the collected groups out into the output, and apply THE +/// completeness check. +/// +/// Nothing here reads payload bytes — a malformed checkpoint fails before the +/// expensive pass. Completeness can only be judged after [`collect_groups`] +/// has seen every shard, because a group's members are free to be split +/// across shards in any order. +fn plan_output( + headers: &[SafetensorsHeader], + map: &dyn Fn(&str) -> Option, +) -> Result, SafetensorsError> { + assert!( + !headers.is_empty(), + "planning needs at least one shard header" + ); + assert!( + headers.len() <= MAX_SHARDS, + "shard count is bounded before planning" + ); + + let (order, mut groups) = collect_groups(headers, map)?; let mut plan = Vec::with_capacity(order.len()); let mut cursor = 0u64; for target in order {