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/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 2079a74..0f2e10b 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -381,6 +381,21 @@ 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. + *(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 @@ -648,12 +663,76 @@ 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.)* -- [ ] **OLMoE from HF safetensors** — the port loads GGUF only because HF stores each expert separately + *(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 +- [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/ @@ -664,6 +743,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 @@ -786,6 +875,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 @@ -1351,6 +1457,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`** 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/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..7884bdf 100644 --- a/crates/mummu/src/models/olmoe.rs +++ b/crates/mummu/src/models/olmoe.rs @@ -13,12 +13,19 @@ //! 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. - -use std::path::Path; +//! 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_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}; use burn::module::Module; use burn::nn::{Embedding, EmbeddingConfig, Linear, LinearConfig, RmsNorm, RmsNormConfig}; @@ -26,12 +33,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 +192,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 +322,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 +625,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 +782,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..66ca655 --- /dev/null +++ b/crates/mummu/src/safetensors.rs @@ -0,0 +1,1192 @@ +//! 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. +//! +//! 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 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. + +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 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. +/// 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) +} + +/// 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<(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(); + + 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())); + } + } + + 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 { + // `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_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(); +} 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" + ); +}