From 36dd459ce51f410dae94306b5929e90cb3ebe6d9 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 20 Aug 2026 23:01:06 -0500 Subject: [PATCH 1/7] deps: refresh the lockfile; wgpu 30 and burn 0.22 both stay held MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `cargo update` moved four transitive crates (either 1.17->1.18, h2 0.4.17->0.4.18, icu_provider 2.3.0->2.3.1, zerovec-derive 0.11.5->0.11.6). `cargo upgrade --incompatible` offers exactly one major: wgpu 29 -> 30, which stays held for the reason already recorded in P0 — `cargo tree -i wgpu` still shows a single wgpu (29.0.4) reached through `cubecl-wgpu 0.10`, so bumping only our direct handle would put a second, non-Burn wgpu in the graph and the startup adapter probe would stop describing the device Burn actually runs on. It unblocks with the burn bump, and burn's newest published version is still 0.22.0-pre.2 — 0.21.0 remains the latest stable, so the do-not-adopt-a-pre-release rule holds for another run. Green on the refreshed lock: fmt clean, `clippy --all-targets` with no new warnings, 232 lib tests + 6 integration tests passing (verified by test name against the shared cargo target), `cargo build` clean. Co-Authored-By: Claude Opus 5 --- Cargo.lock | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8cb08ea..94a796a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2162,9 +2162,9 @@ checksum = "e1d926b4d407d372f141f93bb444696142c29d32962ccbd3531117cf3aa0bfa9" [[package]] name = "either" -version = "1.17.0" +version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" +checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" [[package]] name = "embassy-futures" @@ -2772,9 +2772,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.17" +version = "0.4.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f877e75f39e9827ec50a572dd592684ac28c029578726c85f1b2aa6ab807449" +checksum = "839c0e8a181239723652be9062bb56ca5bf5f64011f73b623f6f4fc59086a228" dependencies = [ "atomic-waker", "bytes", @@ -3068,9 +3068,9 @@ checksum = "e590f038c1464a96894fd6d10127e90a8be4509f56ff7ecef851b15cee0b7caa" [[package]] name = "icu_provider" -version = "2.3.0" +version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92a7ed671a6aad807a8651a2e1782a6598fda9ce5185dd8158549e95a91c6428" +checksum = "d27bbb9d3abbefac45d55f647c9de1d44aafcd1186eb91879afef17c396c3e73" dependencies = [ "displaydoc", "icu_locale_core", @@ -6657,9 +6657,9 @@ dependencies = [ [[package]] name = "zerovec-derive" -version = "0.11.5" +version = "0.11.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f212a141d820099d57ffafb9569be9617a6f27d3dc881fbee8fb56642f917a9" +checksum = "34df6fc39dbd26ddc9c10e6a2984476e13acce22e64e4487636ef494369225da" dependencies = [ "proc-macro2", "quote", From 74edb2e8bef01947f23a188f224e93497df0762f Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 20 Aug 2026 23:10:50 -0500 Subject: [PATCH 2/7] gguf: no .expect() left on a production path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `safetensors.rs` grew a fallible `to_usize` on 2026-08-20; `gguf.rs` kept the `usize::try_from(..).expect("bounded above")` pattern it was modelled on, so the no-panic-on-production-paths rule was true in the newest import module only. It is now true across the whole import suite. The same helper, same reasoning: every caller has already bounded the value against a `MAX_*` ceiling or a tensor's declared size, so on a 64-bit target the conversion cannot fail — but an import path is exactly where "cannot fail" should be an `Err`, so a 32-bit build degrades to a clean `OverBound` naming the file instead of aborting mid-load. Eight sites, all in production code: - `Reader::{string,value,read_file,tensor_table}` — the four header readers, now via `to_usize`. `tensor_table` converts `count` once instead of twice (it was doing the same fallible cast for the capacity and for the loop bound). - `dequant_to_safetensors` — the payload capacity, plus the tensor-name JSON encode, which now reports through the existing `BadTensor` arm and so names the offending tensor index. - `dequantize` — the two block-geometry casts and the F32 4-byte block cast, all mapped into the `String` error it already returns. No behaviour change by construction: every replaced site was unreachable-on-64-bit. Green — fmt, `clippy --all-targets` no new warnings, 232 lib tests. Real-weights proof that the header reader and the dequant path still agree with the file: the four `real_qwen2_gguf_*` legs pass on the Q4_K_M checkpoint (header parse, dequant-vs-true-weights, tokenizer byte identity, GPU load + decode), 259.3 s. Co-Authored-By: Claude Opus 5 --- crates/mummu/src/gguf.rs | 65 ++++++++++++++++++++++++++++++++-------- 1 file changed, 53 insertions(+), 12 deletions(-) diff --git a/crates/mummu/src/gguf.rs b/crates/mummu/src/gguf.rs index 4716ece..2cd1e33 100644 --- a/crates/mummu/src/gguf.rs +++ b/crates/mummu/src/gguf.rs @@ -74,6 +74,28 @@ pub enum GgufError { }, } +/// `u64` -> `usize` as an error rather than a panic. +/// +/// Every caller has already bounded `value` against one of the `MAX_*` +/// ceilings or a tensor's own declared size, 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 instead of aborting mid-load. (Same helper, same reasoning, +/// as `safetensors::to_usize`.) +fn to_usize(value: u64, path: &str, what: &'static str) -> Result { + debug_assert!( + value <= usize::MAX as u64, + "{what} fits usize on this target" + ); + debug_assert!(!path.is_empty(), "an error needs a file to name"); + usize::try_from(value).map_err(|_| GgufError::OverBound { + path: path.to_string(), + what, + count: value, + bound: usize::MAX as u64, + }) +} + /// One typed metadata value. Arrays are homogeneous per the spec; nested /// arrays are legal but bounded to one level of nesting in practice. #[derive(Debug, Clone, PartialEq)] @@ -396,8 +418,11 @@ impl GgufFile { let mut header = String::from("{"); let mut names = std::collections::HashSet::with_capacity(self.tensors.len()); - let mut data: Vec = - Vec::with_capacity(usize::try_from(total_f32_bytes).expect("bounded above")); + let mut data: Vec = Vec::with_capacity(to_usize( + total_f32_bytes, + &self.path.display().to_string(), + "dequantized f32 payload bytes", + )?); for (index, info) in self.tensors.iter().enumerate() { let (name, shape) = match map(info) { Some(GgufMap::Rename(name)) => { @@ -425,7 +450,8 @@ impl GgufFile { let start = data.len(); let values = self.read_tensor_f32(&info.name)?; data.extend(values.iter().flat_map(|v| v.to_le_bytes())); - let json_name = serde_json::to_string(&name).expect("string serializes"); + let json_name = serde_json::to_string(&name) + .map_err(|e| bad(index, format!("name is not encodable as JSON: {e}")))?; if index > 0 { header.push(','); } @@ -504,7 +530,7 @@ impl Reader { bound: MAX_STRING_BYTES, }); } - let mut buf = vec![0u8; usize::try_from(len).expect("bounded above")]; + let mut buf = vec![0u8; to_usize(len, &self.path, what)?]; self.inner .read_exact(&mut buf) .map_err(|e| self.io_err(e))?; @@ -551,7 +577,7 @@ impl Reader { bound: MAX_ARRAY_LEN, }); } - let mut items = Vec::with_capacity(usize::try_from(len).expect("bounded above")); + let mut items = Vec::with_capacity(to_usize(len, &self.path, "metadata array")?); for _ in 0..len { items.push(self.value(key, elem_type, depth + 1)?); } @@ -597,7 +623,8 @@ impl Reader { } } - let mut metadata = Vec::with_capacity(usize::try_from(kv_count).expect("bounded above")); + let mut metadata = + Vec::with_capacity(to_usize(kv_count, &self.path, "metadata key-values")?); for _ in 0..kv_count { let key = self.string("metadata key")?; let type_id = self.u32()?; @@ -640,9 +667,9 @@ impl Reader { ) -> Result, GgufError> { assert!(count <= MAX_TENSORS, "caller bounded the count"); assert!(alignment.is_power_of_two(), "caller validated alignment"); - let mut tensors: Vec = - Vec::with_capacity(usize::try_from(count).expect("bounded above")); - for index in 0..usize::try_from(count).expect("bounded above") { + let count_usize = to_usize(count, &self.path, "tensor count")?; + let mut tensors: Vec = Vec::with_capacity(count_usize); + for index in 0..count_usize { let bad = |reason: String, path: &str| GgufError::BadTensor { path: path.to_string(), index, @@ -701,7 +728,12 @@ impl Reader { /// Dequantize a whole tensor payload to f32. `bytes` must be whole blocks of /// `dtype` (guaranteed for payload slices sized by [`GgufTensorInfo::byte_len`]). pub fn dequantize(dtype: GgmlType, bytes: &[u8]) -> Result, String> { - let bpb = usize::try_from(dtype.bytes_per_block()).expect("small"); + let bpb = usize::try_from(dtype.bytes_per_block()).map_err(|_| { + format!( + "{dtype:?} block is {} bytes, wider than usize", + dtype.bytes_per_block() + ) + })?; if bytes.is_empty() || !bytes.len().is_multiple_of(bpb) { return Err(format!( "{} bytes is not whole {dtype:?} blocks of {bpb}", @@ -709,11 +741,20 @@ pub fn dequantize(dtype: GgmlType, bytes: &[u8]) -> Result, String> { )); } let blocks = bytes.len() / bpb; - let block_elems = usize::try_from(dtype.block_size()).expect("small"); + let block_elems = usize::try_from(dtype.block_size()).map_err(|_| { + format!( + "{dtype:?} block holds {} elements, more than usize", + dtype.block_size() + ) + })?; let mut out = Vec::with_capacity(blocks * block_elems); for block in bytes.chunks_exact(bpb) { match dtype { - GgmlType::F32 => out.push(f32::from_le_bytes(block.try_into().expect("4 bytes"))), + GgmlType::F32 => { + out.push(f32::from_le_bytes(block.try_into().map_err(|_| { + format!("F32 block is {} bytes, not 4", block.len()) + })?)) + } GgmlType::F16 => out.push(f16_to_f32(u16::from_le_bytes([block[0], block[1]]))), GgmlType::BF16 => { out.push(f32::from_bits( From db36ccb642484e8ec784ebada28b355e5f61f97e Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 20 Aug 2026 23:39:17 -0500 Subject: [PATCH 3/7] docs: fold this run's research - a MoE route that survives a cache miss Three folds, each with a number or a correction attached. P0 (burn 0.22): two of the migration note's claims did not survive reading the published crate. burn-store 0.21 has no `FloatCastAdapter` - its adapters are PyTorchToBurn / BurnToPyTorch / HalfPrecision / Chain / Identity, and HalfPrecision is f32<->f16 only - so replacing our `CastFloatAdapter` is a 0.22 task, not a missed opportunity. And wgpu 30's f16 change is wider than recorded: SHADER_F16 works in WGSL and GLSL now, not only SPIR-V passthrough, which means today's f16-is-Vulkan-only shape ends at the bump. Cooperative matrix load/store also lands, but at 8x8 f32 only. P2 (MoE decode): llama.cpp discussion #24528 is the first route that explains our own 2026-08-03 regression instead of ignoring it. It caches hot experts in VRAM and runs HYBRID hit/miss - the cached rows go to the GPU as one batched matvec while the miss rows compute on the CPU exactly as before - so a miss costs nothing extra and the path degrades to the dense baseline rather than paying an unconditional gather. That is the difference from what we measured and reverted. It also needs no whole-model-in-VRAM, so it unblocks route (b) without waiting on P9. P5 (speculative decoding): the MTP route has a merged implementation (llama.cpp #22673) and a target - ~75% acceptance at k=3 for >2x, with k=2 often beating k=3 as acceptance falls off with the draft window. Co-Authored-By: Claude Opus 5 --- ROADMAP.md | 45 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/ROADMAP.md b/ROADMAP.md index 0f2e10b..4013832 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -396,6 +396,21 @@ a benchmark holds/improves its budget; README perf claims link an artifact. {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 + *(2026-08-20, second run)* Two corrections and one addition, all from reading what is actually + published rather than the migration note's summary. (1) **`burn-store` 0.21 does NOT ship a + `FloatCastAdapter`** — its adapter set is `PyTorchToBurnAdapter` / `BurnToPyTorchAdapter` / + `HalfPrecisionAdapter` / `ChainAdapter` / `IdentityAdapter`, and `HalfPrecisionAdapter` is + f32<->f16 only, not the arbitrary target-dtype cast `CastFloatAdapter` does. So "evaluate + replacing our `CastFloatAdapter`" is a 0.22 task, not something already available and skipped. + (2) **wgpu 30's f16 story is bigger than "f16 beyond SPIR-V"**: `SHADER_F16` now works in WGSL + and GLSL as well as SPIR-V passthrough (add `enable f16;` at the top of the shader), which + matters because it means the f16 win stops being Vulkan-only — the `vulkan` feature's SPIR-V + path is currently the only way we get f16 kernels, so a DX12/Metal consumer gets f32 today and + would not after the bump. (3) wgpu 30 also lands **cooperative matrix load/store** (WGSL in; + SPIR-V/Metal/WGSL out), gated on `vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR` and + currently **8x8 f32 only** — too narrow to be the decode lever yet, but it is the seam a + tensor-core matmul would eventually ride, so watch the supported configurations grow. + — https://docs.rs/burn-store/latest/burn_store/ · https://github.com/gfx-rs/wgpu/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 @@ -682,6 +697,24 @@ a benchmark holds/improves its budget; README perf claims link an artifact. before treating the gather as refuted. Both are gated on `bench/BASELINE.md` like the rest — https://github.com/ggml-org/llama.cpp/pull/25294 · https://huggingface.co/blog/Doctor-Shotgun/llamacpp-moe-offload-guide + *(2026-08-20, second run)* A fourth route, and the one that answers our specific regression: + llama.cpp discussion #24528 proposes a **VRAM cache of hot experts with hybrid hit/miss + execution** — the `MUL_MAT_ID` stays on the CPU, but thread 0 dispatches ONE batched matvec + over the cached (hit) expert rows to the GPU while the remaining threads compute the miss rows + exactly as they do today. The property that matters is that **a miss costs nothing extra**: + performance degrades gracefully to the vanilla dense-CPU path instead of paying a gather that + may not pay for itself, which is precisely how our 2026-08-03 A/B lost (1.58 s vs 1.15 s — the + gather was unconditional). Eviction is LRU with an admission threshold, plus a "soft mode" that + trades ~5-7 % of prefill to avoid evicting during decode; the cache **engages on decode only**. + Reported: +25 % on GLM-5.1 754B, +7 % on Qwen3.5 397B, +28-46 % decode on 2x RTX 3090 — but a + GTX 1080 Ti **regressed**, i.e. the win is a function of how much faster the GPU is than the CPU + at the cached slice, which is exactly the hardware-dependence our own measurement found. + Status: RFC, CUDA-only, unmerged. For us the shape translates directly — a bounded per-layer + VRAM slab cache on the `Gpu` backend with the CPU dense path as the miss fallback — and it is + the first route that does NOT require the whole MoE to fit VRAM first, so it unblocks route (b) + without waiting on P9. Gate it on `bench/BASELINE.md` and the OLMoE 0.76 s/token warm number. + — https://github.com/ggml-org/llama.cpp/discussions/24528 · + https://github.com/ggml-org/llama.cpp/issues/20757 - [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, @@ -1292,6 +1325,18 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari worth re-checking after graph capture (P0) moves the dispatch baseline. — https://github.com/thc1006/qwen3.6-speculative-decoding-rtx3090 · https://inventivehq.com/blog/llama-cpp-speculative-decoding-consumer-gpu + *(2026-08-20, second run)* Route (a) now has a merged reference implementation to read and a + number to aim at: llama.cpp PR **#22673** landed MTP-head support (tested on Qwen3.6-27B and + Qwen3.6-35B-A3B, but written for any MTP model), driven by `--spec-type mtp` plus + `--spec-draft-n-max `. Reported steady-state **acceptance ~75 % at k=3 for >2x end-to-end**, + rising past 80 % on code/math/reasoning, and the community finding that **k=2 often beats k=3** + because acceptance falls off as the draft window widens — so `n_draft` is a tunable to measure, + not a constant to pick. Two things to carry into our design: MTP needs no second model in VRAM + (the heads ride the target checkpoint), which is what makes it the route worth building on a + 16 GB card; and the flag rename noted above means a reference leg should probe BOTH + `--spec-type mtp` and `--spec-type draft-mtp` rather than assume either. — + https://github.com/ggml-org/llama.cpp/pull/22673 · + https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md - [ ] **Grammar-constrained decoding** *(colibri parity)* — colibri forces structured output via `.gbnf` grammars (llama.cpp's GBNF convention) and even uses grammar-forced *drafts* to speed structured generation. Mummu's tool-calling currently *trusts* the model to emit parseable `` JSON From 1398138797e91a48df01b5d6ff531bb3f4621e9e Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Thu, 20 Aug 2026 23:58:28 -0500 Subject: [PATCH 4/7] gguf: stream the dequant instead of buffering it twice `dequant_to_safetensors` filled a `data` Vec of the whole f32 payload, then copied it into a second `blob` Vec of `8 + header + data` - peak 2x the payload, before the model was even built. For OLMoE-1B-7B that is ~28 GB dequantized and ~56 GB of commit charged, which is most of why that model reads as "CPU-only, and only on a big box". `safetensors::fuse_into` was the template and this follows it: plan the whole output first, then stream header + payload into an `impl Write`. `dequant_to_safetensors` is now a thin wrapper over the shared `dequant_into`, joined by `dequant_to_safetensors_file`. The double buffer is gone from BOTH forms - the in-memory one writes header then payload into a single Vec (1x, not 2x), and the file one holds nothing but the current tensor's f32 values plus a fixed 1 MiB staging buffer. Planning first also made the path stricter. `plan_dequant` decides every name, shape and offset before the first payload byte is read, so an unmapped name, a reshape that changes the element count, or a rename collision fails in milliseconds rather than after N tensors have been dequantized. A unit test proves the ordering by pointing a tensor's payload past the end of the file and checking the MAP error still surfaces. `olmoe::load_from_gguf` takes the file variant, reusing the `FusedTemp` guard the HF path already uses; that guard's name now carries a counter as well as the pid, so two concurrent loads in one process cannot pick the same scratch file. REAL-WEIGHTS proof - OLMoE-1B-7B Q4_K_M against llama.cpp on the identical file: top-5 ids exact in order, the 24-token greedy sequence byte-identical, max |dlogprob| 3.687691131310693e-1 against a 7.5e-1 tolerance, 130.5 s. Sampled peak while it ran: 26.5 GB private commit - the model alone - against a 51.7 GB working set, i.e. ~25 GB of the load is file-backed mmap rather than charged to commit, which is the binding memory limit on this box. Nothing else moved. qwen2 GGUF parity 2.6614442413586614e-1 and qwen3 4.015608155114805e-1, the latter bit-identical to the recorded value, both greedy sequences byte-identical; the four `real_qwen2_gguf_*` legs pass in 44.1 s. 234 lib tests (+2), fmt clean, clippy --all-targets with no new warnings. The three OLMoE test gates said "~60 GB free RAM"; they now say ~30 GB free commit plus ~28 GB of scratch disk, which is what they actually need. Co-Authored-By: Claude Opus 5 --- README.md | 5 +- ROADMAP.md | 56 ++++- crates/mummu-bench/tests/budget_moe.rs | 3 +- crates/mummu/src/gguf.rs | 295 ++++++++++++++++++++++--- crates/mummu/src/models/olmoe.rs | 42 +++- crates/mummu/tests/parity_gguf.rs | 3 +- crates/mummu/tests/real_olmoe.rs | 3 +- 7 files changed, 362 insertions(+), 45 deletions(-) diff --git a/README.md b/README.md index 2be8682..10a7bcd 100644 --- a/README.md +++ b/README.md @@ -92,7 +92,10 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys 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. + a second copy of itself. The **GGUF** path streams the same way: its dequant plans the whole output + before reading a payload byte, then writes header + f32 tensors straight to a temp file that + `burn-store` mmaps back, so loading the 1B-7B costs **26.5 GB of measured private commit — the model + alone** — where the old in-RAM dequant paid for the payload twice on top of it. - **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 4013832..f02a92d 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -752,7 +752,7 @@ a benchmark holds/improves its budget; README perf claims link an artifact. 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 +- [x] **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. @@ -761,11 +761,63 @@ a benchmark holds/improves its budget; README perf claims link an artifact. + 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 + *(2026-08-20, second run) Shipped.* `dequant_to_safetensors` is now a thin wrapper over a shared + `dequant_into`, joined by `dequant_to_safetensors_file`. The double-buffer is gone from + BOTH forms: the in-memory one writes header-then-payload into a single `Vec` (peak 1x the payload, + down from 2x), and the file one holds nothing but the current tensor's f32 values plus a fixed + 1 MiB staging buffer. What made streaming possible also made the path stricter — a new + `plan_dequant` decides every name, shape and offset **before the first payload byte is read**, so + an unmapped name, a reshape that changes the element count, or a rename collision now fails in + milliseconds instead of after N tensors have been dequantized (a unit test proves this by + pointing a tensor's payload past the end of the file and checking the MAP error still surfaces). + `olmoe::load_from_gguf` takes the file variant, reusing the same `FusedTemp` guard the HF path + uses, so the scratch file is deleted on every exit path. + **REAL-WEIGHTS proof**, the OLMoE-1B-7B Q4_K_M leg against llama.cpp on the identical file: + top-5 ids exact in order, the 24-token greedy sequence **byte-identical**, max |Δlogprob| + 3.687691131310693e-1 (tolerance 7.5e-1), 130.5 s. Measured peak while it ran: **26.5 GB private + commit** — the model alone — against a 51.7 GB working set, i.e. ~25 GB of the load is now + file-backed mmap rather than charged to commit, which is the binding limit on this box. The old + shape was ~28 GB blob + ~28 GB model ≈ 56 GB of commit. Nothing else moved: qwen2 GGUF parity + 2.6614442413586614e-1 and qwen3 **4.015608155114805e-1, bit-identical to the recorded value**, + both greedy sequences byte-identical; the four `real_qwen2_gguf_*` legs pass (44.1 s). 234 unit + tests (+2). The stale "~60 GB free RAM" gates on the three OLMoE test legs are corrected to + ~30 GB free commit + ~28 GB of scratch disk. +- [x] **`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. + *(2026-08-20, second run) Done — eight sites, not four.* `gguf::to_usize` mirrors + `safetensors::to_usize` exactly. The four header readers (`Reader::{string,value,read_file, + tensor_table}`) go through it — `tensor_table` now converts its count once instead of twice — + and so does `dequant_to_safetensors`' payload capacity. Three more were hiding outside the + `try_from` pattern: the tensor-name JSON encode (now reports through `BadTensor`, so it names + the offending index), and `dequantize`'s two block-geometry casts plus its F32 4-byte block + cast, all mapped into the `String` error that function already returns. No behaviour change by + construction — every replaced site was unreachable on a 64-bit target — proven by the four + `real_qwen2_gguf_*` legs on the Q4_K_M checkpoint (header parse, dequant-vs-true-weights, + tokenizer byte identity, GPU load + decode). +- [ ] **Decide the dequant sink per model, not per code path** — `olmoe::load_from_gguf` takes the new + streaming `dequant_to_safetensors_file`; qwen2 / qwen3 / lfm2 still take the in-memory + `dequant_to_safetensors` (which is now 1x the payload rather than 2x, so they already got half the + win for free). That split is a guess, not a measurement: the file variant trades a spike in commit + for an f32 write plus mmap-back, which should LOSE on a 1-2 GB model and win on anything that is a + meaningful fraction of RAM. Measure GGUF load wall-clock both ways on Qwen2.5-1.5B (~2.2 GB f32) + and Qwen3-0.6B, then either pick per-model or — better — pick automatically from + `total_f32_bytes` against the device inventory's free-RAM figure (P6 already probes it), so a + consumer never has to know. *(2026-08-20, discovered shipping the streaming sink.)* +- [ ] **`FusedTemp` is import-suite machinery living in one model file** — `models/olmoe.rs` owns the + scratch-file guard (create beside the weights, unique per process + counter, `Drop`-delete on + every exit path), and it now serves BOTH import paths. Any other model large enough to want a + streaming sink has to reach into `olmoe` or copy it. Promote it to `import.rs` when a second + model needs it — not before, since a one-caller abstraction moved early is just churn. + *(2026-08-20.)* +- [ ] **The f32 scratch file is a symptom, not the design** — both streaming sinks exist because + `SafetensorsStore` is the only checked-load pipeline we have, so every import must first become + f32 safetensors on disk. A GGUF **keep-quantized** load (P9) would skip the temp file entirely: + no ~28 GB write, no 4x dtype widening, and the 1B-7B would stop being CPU-only. Worth stating + here so the scratch-file machinery is understood as a bridge with a known end, not a fixture. + *(2026-08-20.)* - [ ] **Qwen3.5 hybrid (`qwen35`) architecture port** — split from the Qwen3.5-tier item when the 2026-07-30 header probe showed Qwen3.5-4B/9B are a hybrid **linear-attention/SSM + periodic full-attention** arch (`qwen35.ssm.*` metadata: conv_kernel/state_size/group_count/time_step_rank/ diff --git a/crates/mummu-bench/tests/budget_moe.rs b/crates/mummu-bench/tests/budget_moe.rs index cf35e48..7611b99 100644 --- a/crates/mummu-bench/tests/budget_moe.rs +++ b/crates/mummu-bench/tests/budget_moe.rs @@ -34,7 +34,8 @@ const DECODE_STEPS: usize = 4; const LOAD_BUDGET_SECS: f64 = 300.0; #[test] -#[ignore = "needs the OLMoE Q4_K_M GGUF (MUMMU_OLMOE_GGUF_PATH) and ~40 GB free RAM"] +#[ignore = "needs the OLMoE Q4_K_M GGUF (MUMMU_OLMOE_GGUF_PATH), ~30 GB free COMMIT \ + and ~28 GB of scratch disk beside the gguf"] fn olmoe_moe_cpu_decode_stays_inside_its_budget() { let Some(path) = std::env::var_os("MUMMU_OLMOE_GGUF_PATH").map(PathBuf::from) else { panic!("set MUMMU_OLMOE_GGUF_PATH to the OLMoE-1B-7B q4_k_m gguf"); diff --git a/crates/mummu/src/gguf.rs b/crates/mummu/src/gguf.rs index 2cd1e33..9856551 100644 --- a/crates/mummu/src/gguf.rs +++ b/crates/mummu/src/gguf.rs @@ -9,7 +9,7 @@ //! loaded here; dequantizing them into Burn tensors is the next slice. use std::fs::File; -use std::io::{BufReader, Read, Seek}; +use std::io::{BufReader, Read, Seek, Write}; use std::path::Path; /// GGUF file magic, little-endian `"GGUF"`. @@ -37,10 +37,26 @@ const MAX_ARRAY_LEN: u64 = 1 << 22; /// GGML allows at most 4 tensor dimensions. const MAX_DIMS: u32 = 4; -/// Largest dequantized-to-f32 payload [`GgufFile::dequant_to_safetensors`] -/// will build in RAM (a ~12B-param model; the reference machine has 128 GB). +/// Largest dequantized-to-f32 payload either dequant will produce (a ~12B-param +/// model; the reference machine has 128 GB). Only +/// [`GgufFile::dequant_to_safetensors`] holds this much at once; +/// [`GgufFile::dequant_to_safetensors_file`] streams it and never buffers more +/// than one tensor. const MAX_DEQUANT_BYTES: u64 = 48 << 30; +/// Largest SINGLE dequantized tensor the streaming dequant will hold. The +/// widest real one is a vocabulary embedding — Qwen2.5-1.5B's is 933 MB at +/// f32, a 35B's would be ~3.1 GB — so 8 GiB is a corrupt header claiming an +/// absurd tensor. It is twice `safetensors::MAX_PART_BYTES` on purpose: that +/// path keeps the source dtype, this one widens everything to f32. +const MAX_TENSOR_F32_BYTES: u64 = 8 << 30; + +/// Bytes staged between the f32 values and the sink. Fixed and small: the +/// point of the streaming dequant is that peak allocation tracks the widest +/// TENSOR, never the payload, and this buffer must not reintroduce a second +/// copy of either. +const DEQUANT_STAGE_BYTES: usize = 1 << 20; + /// What went wrong reading a GGUF header. #[derive(Debug, thiserror::Error)] pub enum GgufError { @@ -303,6 +319,20 @@ pub enum GgufMap { Reshape(String, Vec), } +/// Where one tensor lands in the output blob, decided before any payload +/// byte is read ([`GgufFile::plan_dequant`]). +#[derive(Debug)] +struct PlannedTensor { + /// Index into [`GgufFile::tensors`] — the payload to dequantize. + source: usize, + name: String, + shape: Vec, + /// Byte offset within the payload; contiguous and ascending. + start: u64, + /// Dequantized f32 byte length. + len: u64, +} + /// A parsed GGUF header: typed metadata + the located tensor table. #[derive(Debug)] pub struct GgufFile { @@ -397,10 +427,142 @@ impl GgufFile { /// [`GgufMap::Reshape`] overrides that for tensors whose checkpoint /// shape differs by more than dim order (e.g. llama.cpp squeezes the /// middle 1 out of depthwise-conv kernels). + /// + /// This form needs the whole f32 payload resident. For a model whose + /// dequantized size is a meaningful fraction of RAM, use + /// [`Self::dequant_to_safetensors_file`] — the two produce byte-identical + /// output. pub fn dequant_to_safetensors( &self, map: &dyn Fn(&GgufTensorInfo) -> Option, ) -> Result, GgufError> { + let mut blob = Vec::new(); + self.dequant_into(map, &mut blob)?; + assert!(blob.len() > 8, "the header length prefix is always written"); + Ok(blob) + } + + /// [`Self::dequant_to_safetensors`] straight to a file, never holding the + /// payload in RAM. Returns the payload bytes written (header excluded). + /// + /// This is the variant a real quantized checkpoint wants. The in-memory + /// form needs the whole f32 payload resident — ~28 GB for OLMoE-1B-7B — + /// *on top of* the model the load then builds from it, 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. + /// (`safetensors::fuse_checkpoint_to_file` is the same trade on the + /// unquantized path.) + pub fn dequant_to_safetensors_file( + &self, + map: &dyn Fn(&GgufTensorInfo) -> Option, + out: &Path, + ) -> Result { + assert!(!out.as_os_str().is_empty(), "the sink file must be named"); + let io = |source: std::io::Error| GgufError::Io { + path: out.display().to_string(), + source, + }; + let file = File::create(out).map_err(io)?; + let mut sink = std::io::BufWriter::with_capacity(DEQUANT_STAGE_BYTES, file); + let written = self.dequant_into(map, &mut sink)?; + sink.flush().map_err(io)?; + sink.into_inner() + .map_err(|e| GgufError::Io { + path: out.display().to_string(), + source: e.into_error(), + })? + .sync_all() + .map_err(io)?; + Ok(written) + } + + /// The shared dequant: plan, then stream header + payload into `sink` in + /// output order, one tensor at a time. + /// + /// Planning first is not only what makes streaming possible — it moves + /// every *claim* check (unmapped name, reshape that changes the element + /// count, rename collision) ahead of the first payload byte, so a bad map + /// fails in milliseconds instead of after N tensors have been + /// dequantized. + fn dequant_into( + &self, + map: &dyn Fn(&GgufTensorInfo) -> Option, + sink: &mut W, + ) -> Result { + let path = self.path.display().to_string(); + let plan = self.plan_dequant(map)?; + let total: u64 = plan.iter().map(|p| p.len).sum(); + + // 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| GgufError::BadTensor { + path: path.clone(), + index: p.source, + reason: format!("name is not encodable as JSON: {e}"), + })?; + header.push_str(&format!( + "{json_name}:{{\"dtype\":\"F32\",\"shape\":{:?},\"data_offsets\":[{},{}]}}", + p.shape, + p.start, + p.start + p.len, + )); + } + header.push('}'); + + let io = |source: std::io::Error| GgufError::Io { + path: path.clone(), + source, + }; + sink.write_all(&(header.len() as u64).to_le_bytes()) + .map_err(io)?; + sink.write_all(header.as_bytes()).map_err(io)?; + + // One small staging buffer, reused for every tensor: peak allocation + // is the widest tensor's f32 values (from `read_tensor_f32`) plus this + // 1 MiB, never a second copy of the payload. + let mut stage: Vec = Vec::with_capacity(DEQUANT_STAGE_BYTES); + let mut written = 0u64; + for p in &plan { + debug_assert_eq!(written, p.start, "tensors are written in output order"); + let values = self.read_tensor_f32(&self.tensors[p.source].name)?; + debug_assert_eq!( + values.len() as u64 * 4, + p.len, + "the plan sized this tensor from the same element count" + ); + for chunk in values.chunks(DEQUANT_STAGE_BYTES / 4) { + stage.clear(); + stage.extend(chunk.iter().flat_map(|v| v.to_le_bytes())); + sink.write_all(&stage).map_err(io)?; + written += stage.len() as u64; + } + } + assert_eq!( + written, total, + "every planned byte was written exactly once" + ); + assert!( + stage.capacity() <= DEQUANT_STAGE_BYTES, + "the staging buffer never grew past its bound" + ); + Ok(written) + } + + /// First pass: decide what every tensor is called, what shape it claims, + /// and where its bytes land — reading no payload at all. + /// + /// Offsets are contiguous and ascending in tensor-table order, which is + /// what lets the copy pass be a single forward stream. + fn plan_dequant( + &self, + map: &dyn Fn(&GgufTensorInfo) -> Option, + ) -> Result, GgufError> { let total_f32_bytes: u64 = self.tensors.iter().map(|t| t.element_count() * 4).sum(); if total_f32_bytes > MAX_DEQUANT_BYTES { return Err(GgufError::OverBound { @@ -416,13 +578,9 @@ impl GgufFile { reason, }; - let mut header = String::from("{"); + let mut plan: Vec = Vec::with_capacity(self.tensors.len()); let mut names = std::collections::HashSet::with_capacity(self.tensors.len()); - let mut data: Vec = Vec::with_capacity(to_usize( - total_f32_bytes, - &self.path.display().to_string(), - "dequantized f32 payload bytes", - )?); + let mut start = 0u64; for (index, info) in self.tensors.iter().enumerate() { let (name, shape) = match map(info) { Some(GgufMap::Rename(name)) => { @@ -447,27 +605,27 @@ impl GgufFile { if !names.insert(name.clone()) { return Err(bad(index, format!("rename collision on '{name}'"))); } - let start = data.len(); - let values = self.read_tensor_f32(&info.name)?; - data.extend(values.iter().flat_map(|v| v.to_le_bytes())); - let json_name = serde_json::to_string(&name) - .map_err(|e| bad(index, format!("name is not encodable as JSON: {e}")))?; - if index > 0 { - header.push(','); + let len = info.element_count() * 4; + if len > MAX_TENSOR_F32_BYTES { + return Err(GgufError::OverBound { + path: self.path.display().to_string(), + what: "single dequantized tensor bytes", + count: len, + bound: MAX_TENSOR_F32_BYTES, + }); } - header.push_str(&format!( - "{json_name}:{{\"dtype\":\"F32\",\"shape\":{shape:?},\"data_offsets\":[{start},{end}]}}", - end = data.len(), - )); + plan.push(PlannedTensor { + source: index, + name, + shape, + start, + len, + }); + start += len; } - header.push('}'); - - assert_eq!(data.len() as u64, total_f32_bytes, "every element written"); - let mut blob = Vec::with_capacity(8 + header.len() + data.len()); - blob.extend_from_slice(&(header.len() as u64).to_le_bytes()); - blob.extend_from_slice(header.as_bytes()); - blob.extend_from_slice(&data); - Ok(blob) + assert_eq!(start, total_f32_bytes, "the plan covers the whole payload"); + assert_eq!(plan.len(), self.tensors.len(), "every tensor is planned"); + Ok(plan) } /// Look up a metadata value by exact key. @@ -1131,6 +1289,19 @@ mod tests { result } + /// A unique scratch path for tests that write their own output file. + fn scratch_path(tag: &str) -> std::path::PathBuf { + use std::sync::atomic::{AtomicU64, Ordering}; + static NEXT: AtomicU64 = AtomicU64::new(0); + let dir = std::env::temp_dir().join("mummu-gguf-tests"); + std::fs::create_dir_all(&dir).expect("temp dir"); + dir.join(format!( + "{tag}-{}-{}.safetensors", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )) + } + fn open_bytes(bytes: &[u8]) -> Result { with_gguf_bytes(bytes, |r| r) } @@ -1471,6 +1642,74 @@ mod tests { }); } + #[test] + fn dequanting_to_a_file_is_byte_identical_to_dequanting_in_memory() { + // Mixed dtypes and widths, so the pin covers the quantized path and + // the tensor-to-tensor offset arithmetic, not just one F32 copy. + // Q8_0 blocks are 34 B for 32 elements: two blocks = 68 B at offset 0, + // then a 6-element F32 tensor at the next 32-aligned offset. + let mut payload = vec![0u8; 96]; + for (i, b) in payload.iter_mut().enumerate().take(68) { + *b = u8::try_from(i % 251).expect("bounded by the modulus"); + } + payload.extend((1..=6u8).flat_map(|v| f32::from(v).to_le_bytes())); + let bytes = TestGguf::new() + .kv_str("general.architecture", "qwen2") + .tensor("blk.0.attn_q.weight", &[32, 2], 8, 0) // Q8_0 + .tensor("token_embd.weight", &[2, 3], 0, 96) // F32 + .build_with_payload(&payload); + + with_gguf_bytes(&bytes, |f| { + let f = f.expect("parses"); + let map = |i: &GgufTensorInfo| Some(GgufMap::Rename(format!("model.{}", i.name))); + let blob = f.dequant_to_safetensors(&map).expect("serializes"); + + let out = scratch_path("dequant-pin"); + let written = f + .dequant_to_safetensors_file(&map, &out) + .expect("streams to a file"); + let from_file = std::fs::read(&out).expect("reads back"); + let _ = std::fs::remove_file(&out); + + assert_eq!( + blob, from_file, + "the in-memory and to-file dequants must stay ONE importer" + ); + // The returned count is the payload, header excluded — so the + // caller can size the model without re-statting the file. + let header_len = u64::from_le_bytes(blob[0..8].try_into().expect("8 bytes")); + assert_eq!(written, blob.len() as u64 - 8 - header_len); + // 64 Q8_0 elements + 6 F32 elements, all at f32. + assert_eq!(written, (64 + 6) * 4); + }); + } + + #[test] + fn a_bad_map_is_rejected_before_any_payload_is_read() { + // The tensor table claims a payload far past the end of the file, so + // ANY read of it is an io error. A map error must still surface as the + // map error: planning happens first, by construction. + let bytes = TestGguf::new() + .tensor("a.weight", &[8], 0, 0) + .tensor("b.weight", &[1 << 20], 0, 32) + .build_with_payload(&[0u8; 64]); + with_gguf_bytes(&bytes, |f| { + let f = f.expect("parses"); + assert!(matches!( + f.dequant_to_safetensors( + &|i| (i.name != "b.weight").then(|| GgufMap::Rename(i.name.clone())) + ), + Err(GgufError::BadTensor { index: 1, .. }) + )); + // And the collision check too — it is the second claim planning + // makes, on a tensor whose payload is unreadable. + assert!(matches!( + f.dequant_to_safetensors(&|_| Some(GgufMap::Rename("same".into()))), + Err(GgufError::BadTensor { index: 1, .. }) + )); + }); + } + #[test] fn misaligned_offsets_partial_blocks_and_duplicates_are_rejected() { let misaligned = TestGguf::new().tensor("t", &[32], 8, 7).build(); diff --git a/crates/mummu/src/models/olmoe.rs b/crates/mummu/src/models/olmoe.rs index 7884bdf..2da4751 100644 --- a/crates/mummu/src/models/olmoe.rs +++ b/crates/mummu/src/models/olmoe.rs @@ -23,7 +23,9 @@ //! 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. +//! then builds. [`load_from_gguf`] makes the same trade for the same reason: +//! its dequant streams to a temp file (~28 GB of f32) that `burn-store` then +//! mmaps back, so the payload is file-backed rather than charged to commit. use std::path::{Path, PathBuf}; @@ -297,6 +299,12 @@ fn olmoe_gguf_name(name: &str) -> Option { /// same checked-load pipeline every other port uses. Budget note: the 1B-7B's /// ~7B params dequantize to ~28 GB of f32 — size the target device (the /// reference machine runs it on the 128 GB CPU backend). +/// +/// The dequant goes to a TEMP FILE, not to RAM, for the same reason +/// `load_from_dir`'s fuse does: holding ~28 GB of f32 payload *and* building +/// the ~28 GB model from it is the sum a 128 GB box with other tenants +/// actually fails to satisfy. Streaming keeps the peak at the model plus one +/// tensor. The file is this process's to delete, on success or failure. pub fn load_from_gguf( path: &Path, device: &B::Device, @@ -307,17 +315,20 @@ pub fn load_from_gguf( }; let f = GgufFile::open(path).map_err(|e| parse(e.to_string()))?; let config = OlmoeConfig::from_gguf(&f).map_err(parse)?; - let blob = f - .dequant_to_safetensors(&gguf_tensor_to_hf) + // Beside the GGUF, so the scratch write lands on the volume the weights + // already live on. + let scratch = FusedTemp::new(path.parent().unwrap_or(Path::new(".")))?; + let bytes = f + .dequant_to_safetensors_file(&gguf_tensor_to_hf, scratch.path()) .map_err(|e| parse(e.to_string()))?; - assert!(blob.len() > 8, "a parsed GGUF yields a non-empty blob"); + assert!(bytes > 0, "a parsed GGUF yields a non-empty payload"); let mut model = build::(&config, device); // The backend's float dtype, taken from the TYPE (`B::FloatElem`), never // from a probe tensor (per-device default policy hazard). let target_float = ::dtype(); let mut store = install_remaps( - SafetensorsStore::from_bytes(Some(blob)) + SafetensorsStore::from_file(scratch.path().to_path_buf()) .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) .allow_partial(true), ); @@ -381,11 +392,13 @@ fn fused_expert_target(name: &str, num_experts: usize) -> Option { /// /// 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`. +/// Owns the scratch safetensors file for the life of one load — the fused +/// checkpoint on the HF path, the dequantized payload on the GGUF path. /// -/// 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 `?`. +/// It is as large as the weights it carries (13.8 GB fusing the 1B-7B, ~28 GB +/// dequantizing its Q4_K_M), 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, } @@ -393,11 +406,18 @@ struct FusedTemp { 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. + /// + /// The name carries a process-unique counter as well as the pid: two + /// concurrent loads in one process must not choose the same scratch file + /// and interleave their writes into it. fn new(dir: &Path) -> Result { + use std::sync::atomic::{AtomicU64, Ordering}; + static NEXT: AtomicU64 = AtomicU64::new(0); assert!(!dir.as_os_str().is_empty(), "checkpoint dir must be named"); let path = dir.join(format!( - "mummu-fused-{}.safetensors.tmp", - std::process::id() + "mummu-fused-{}-{}.safetensors.tmp", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) )); // A leftover from a killed process must never be mistaken for ours. if path.exists() { diff --git a/crates/mummu/tests/parity_gguf.rs b/crates/mummu/tests/parity_gguf.rs index bdfafe8..effb463 100644 --- a/crates/mummu/tests/parity_gguf.rs +++ b/crates/mummu/tests/parity_gguf.rs @@ -110,7 +110,8 @@ fn lfm2_q4_gguf_matches_llama_cpp_on_the_same_file() { #[test] #[ignore = "needs the OLMoE Q4_K_M GGUF (MUMMU_OLMOE_GGUF_PATH), llama-server \ - (MUMMU_LLAMA_SERVER), and ~60 GB free RAM (28 GB f32 build, CPU backend)"] + (MUMMU_LLAMA_SERVER), ~30 GB free COMMIT and ~28 GB of scratch disk \ + beside the gguf (28 GB f32 build, CPU backend)"] fn olmoe_q4_gguf_matches_llama_cpp_on_the_same_file() { // The zoo's first MoE leg. Our side runs on the CPU backend — the ~28 GB // f32 resident-everything build does not fit a 16 GB card; parity is diff --git a/crates/mummu/tests/real_olmoe.rs b/crates/mummu/tests/real_olmoe.rs index 9f80816..8508ef8 100644 --- a/crates/mummu/tests/real_olmoe.rs +++ b/crates/mummu/tests/real_olmoe.rs @@ -87,7 +87,8 @@ fn olmoe_gguf_tokenizer_matches_the_hf_tokenizer() { /// mapped) and greedy-decodes a correct answer through the MoE stack on the /// CPU backend. #[test] -#[ignore = "needs network (MUMMU_HUB_DEST; ~4.2 GB) and ~60 GB free RAM (CPU backend)"] +#[ignore = "needs network (MUMMU_HUB_DEST; ~4.2 GB), ~30 GB free COMMIT and ~28 GB \ + of scratch disk beside the gguf (CPU backend)"] fn olmoe_gguf_loads_and_decodes_on_cpu() { let path = fetch_olmoe(); let f = GgufFile::open(&path).expect("valid GGUF"); From 624dce98061d5b5cf44c9c1951000daf2e587fb4 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Fri, 21 Aug 2026 00:57:15 -0500 Subject: [PATCH 5/7] import: one gguf_store for all four ports, and measure the sink choice Last commit gave OLMoE's GGUF load the streaming sink and left the other three on the in-memory one, with a roadmap note admitting that split was a guess. Measured, it was the wrong guess. A/B on Qwen2.5-1.5B Q4_K_M (~6 GB of f32), the same load test each way: in-memory 33.90 / 35.93 / 42.15 s, scratch file 36.58 / 37.14 s. The scratch numbers sit inside the in-memory run's own spread, so at any payload of real size the disk round-trip is free and the doubled peak was being paid for nothing. (n=2 on the scratch side: a concurrent routine took the shared cargo target's lock and the third run never got to start.) So the choice becomes automatic. `import::DequantSink` is Memory / Scratch / Auto, and Auto keeps a payload in RAM only up to MAX_IN_MEMORY_DEQUANT_BYTES (512 MiB, i.e. a ~1 GiB peak - free anywhere). Everything else streams. All four GGUF ports now pass DequantSink::Auto, so no consumer has to know the knob exists. The refactor that made this expressible is worth as much as the policy: `import::gguf_store` is the one place the four ports shared - the dequant, the adapter chain, and the comment explaining why the target float dtype comes from the TYPE (`B::FloatElem`) and never from a probe tensor. That reasoning existed in four copies, phrased three different ways. Each port now adds only its own key remaps. `FusedTemp` moves out of models/olmoe.rs and becomes `import::ScratchFile` - the second caller the previous commit predicted arrived immediately, and two copies of a delete-on-every-exit-path guard is one too many. Proof that the new default changed nothing: all three GGUF parity legs against llama.cpp on the identical file, with qwen2 and qwen3 now taking the scratch path they did not take before - qwen2 2.6614442413586614e-1 and qwen3 4.015608155114805e-1, both bit-identical to the values recorded before this change, OLMoE 3.687691131310693e-1, and all three 24-token greedy sequences byte-identical to llama.cpp. The four `real_qwen2_gguf_*` legs pass (40.1 s), the four `real_inference` safetensors legs pass (34.3 s), and the model dirs hold no leftover scratch files afterwards. 236 lib tests (+2: the Auto boundary in both directions, and that two guards never name the same file and both delete on drop). Co-Authored-By: Claude Opus 5 --- crates/mummu/src/import.rs | 212 ++++++++++++++++++++++++++++++- crates/mummu/src/models/lfm2.rs | 19 +-- crates/mummu/src/models/olmoe.rs | 80 ++---------- crates/mummu/src/models/qwen2.rs | 19 +-- crates/mummu/src/models/qwen3.rs | 22 +--- 5 files changed, 243 insertions(+), 109 deletions(-) diff --git a/crates/mummu/src/import.rs b/crates/mummu/src/import.rs index 3e9628d..e7cc58b 100644 --- a/crates/mummu/src/import.rs +++ b/crates/mummu/src/import.rs @@ -11,9 +11,14 @@ use std::path::{Path, PathBuf}; use std::rc::Rc; use burn::module::Module; -use burn::store::{ModuleAdapter, ModuleSnapshot, ModuleStore, TensorSnapshot}; +use burn::store::{ + ModuleAdapter, ModuleSnapshot, ModuleStore, PyTorchToBurnAdapter, SafetensorsStore, + TensorSnapshot, +}; use burn::tensor::DType; +use crate::gguf::{GgufFile, GgufMap, GgufTensorInfo}; + /// Everything that can go wrong turning files on disk into a loaded model. #[derive(Debug, thiserror::Error)] pub enum ImportError { @@ -221,6 +226,163 @@ where Ok(()) } +/// Largest f32 payload [`DequantSink::Auto`] will let a dequant hold in RAM. +/// +/// The in-memory sink's peak is ~2x the payload (the blob, plus the model +/// being built from it), so this ceiling is really a ~1 GiB peak — small +/// enough to be free on any machine that could run the model at all. Above +/// it, [`DequantSink::Scratch`] costs one f32 write plus an mmap read-back +/// and halves the peak; measured on Qwen2.5-1.5B (6.2 GB of f32), that trade +/// is inside the run-to-run noise of the load it is part of, so there is no +/// reason to keep paying the doubled peak for anything of real size. +pub const MAX_IN_MEMORY_DEQUANT_BYTES: u64 = 512 << 20; + +/// Where a GGUF dequant's f32 payload lands on its way into a +/// [`SafetensorsStore`]. +/// +/// The two sinks are byte-identical by construction — `gguf::dequant_into` is +/// ONE function behind both, pinned by a unit test — so this is a resource +/// trade, never a correctness one. [`Self::Memory`] peaks at ~2x the payload +/// and needs nothing; [`Self::Scratch`] peaks at ~1x (the payload becomes +/// file-backed mmap) and needs payload-sized free disk beside the weights. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DequantSink { + /// Hold the whole f32 payload in RAM and hand it to the store as bytes. + Memory, + /// Stream it to a scratch file beside the weights; `burn-store` mmaps + /// that back and materializes tensors lazily, so the payload is never + /// charged to commit. + Scratch, + /// Pick by payload size: [`Self::Memory`] only up to + /// [`MAX_IN_MEMORY_DEQUANT_BYTES`], [`Self::Scratch`] above it. + Auto, +} + +impl DequantSink { + /// Resolve [`Self::Auto`] against a payload size; the other two answer + /// for themselves. Never returns `Auto`. + #[must_use] + pub fn resolve(self, total_f32_bytes: u64) -> Self { + let picked = match self { + Self::Auto if total_f32_bytes > MAX_IN_MEMORY_DEQUANT_BYTES => Self::Scratch, + Self::Auto => Self::Memory, + explicit => explicit, + }; + debug_assert!(picked != Self::Auto, "resolve must decide"); + debug_assert!( + picked != Self::Memory + || total_f32_bytes <= MAX_IN_MEMORY_DEQUANT_BYTES + || self == Self::Memory, + "Auto never picks Memory above the ceiling" + ); + picked + } +} + +/// Owns a scratch file for the life of one load. +/// +/// The file is as large as the weights it carries (~28 GB dequantizing +/// OLMoE-1B-7B's Q4_K_M), 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 `?` — and because the store reads it lazily, the guard +/// must outlive `load_checked`, which holding it as a local does. +#[derive(Debug)] +pub struct ScratchFile { + path: PathBuf, +} + +impl ScratchFile { + /// Placed beside the weights, so the scratch write lands on the same + /// volume as the model rather than on a small system temp drive. + /// + /// The name carries a process-unique counter as well as the pid: two + /// concurrent loads in one process must not choose the same file and + /// interleave their writes into it. + pub fn new(dir: &Path) -> Result { + use std::sync::atomic::{AtomicU64, Ordering}; + static NEXT: AtomicU64 = AtomicU64::new(0); + assert!(!dir.as_os_str().is_empty(), "scratch dir must be named"); + let path = dir.join(format!( + "mummu-scratch-{}-{}.safetensors.tmp", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + // 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 scratch file: {e}"), + })?; + } + Ok(Self { path }) + } + + #[must_use] + pub fn path(&self) -> &Path { + debug_assert!(!self.path.as_os_str().is_empty(), "scratch path is named"); + &self.path + } +} + +impl Drop for ScratchFile { + 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); + } +} + +/// Dequantize `f` to f32 safetensors and return the store every GGUF loader +/// then adds its own key remaps to, plus the scratch guard (if any) that must +/// stay alive until the load finishes. +/// +/// This is the one place the four GGUF ports share: the sink choice, the +/// adapter chain, and — the part worth centralizing — taking the target float +/// dtype from the TYPE (`B::FloatElem`) rather than from a probe tensor. +/// Unspecified-dtype tensor creation follows the per-DEVICE default policy, +/// which another backend alias sharing the device (`Gpu` vs `GpuF16`) may have +/// flipped in this process. +pub fn gguf_store( + f: &GgufFile, + map: &dyn Fn(&GgufTensorInfo) -> Option, + sink: DequantSink, +) -> Result<(SafetensorsStore, Option), ImportError> { + let parse = |reason: String| ImportError::Parse { + file: f.path.clone(), + reason, + }; + let total_f32_bytes: u64 = f.tensors.iter().map(|t| t.element_count() * 4).sum(); + let target_float = ::dtype(); + let adapter = PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float)); + + match sink.resolve(total_f32_bytes) { + DequantSink::Memory => { + let blob = f + .dequant_to_safetensors(map) + .map_err(|e| parse(e.to_string()))?; + assert!(blob.len() > 8, "a parsed GGUF yields a non-empty blob"); + let store = SafetensorsStore::from_bytes(Some(blob)) + .with_from_adapter(adapter) + .allow_partial(true); + Ok((store, None)) + } + DequantSink::Scratch => { + // Beside the gguf, so the scratch write lands on the volume the + // weights already live on. + let scratch = ScratchFile::new(f.path.parent().unwrap_or(Path::new(".")))?; + let bytes = f + .dequant_to_safetensors_file(map, scratch.path()) + .map_err(|e| parse(e.to_string()))?; + assert!(bytes > 0, "a parsed GGUF yields a non-empty payload"); + let store = SafetensorsStore::from_file(scratch.path().to_path_buf()) + .with_from_adapter(adapter) + .allow_partial(true); + Ok((store, Some(scratch))) + } + DequantSink::Auto => unreachable!("resolve never returns Auto"), + } +} + /// `dir/file`, or [`ImportError::MissingFile`] if absent. pub fn required_file(dir: &Path, file: &str) -> Result { assert!(!file.is_empty(), "required_file: empty file name"); @@ -258,6 +420,54 @@ pub fn weights_file(dir: &Path) -> Result { #[cfg(test)] mod tests { + use super::{DequantSink, MAX_IN_MEMORY_DEQUANT_BYTES, ScratchFile}; + + #[test] + fn auto_picks_the_scratch_sink_exactly_above_the_ceiling() { + // The boundary is the whole point: `Auto` must never leave a payload + // in RAM whose 2x peak is what this ceiling exists to cap. + assert_eq!(DequantSink::Auto.resolve(0), DequantSink::Memory); + assert_eq!( + DequantSink::Auto.resolve(MAX_IN_MEMORY_DEQUANT_BYTES), + DequantSink::Memory + ); + assert_eq!( + DequantSink::Auto.resolve(MAX_IN_MEMORY_DEQUANT_BYTES + 1), + DequantSink::Scratch + ); + // An explicit choice is honoured in BOTH directions, including the + // one `Auto` would never make — the A/B that set the ceiling needs it. + assert_eq!(DequantSink::Memory.resolve(u64::MAX), DequantSink::Memory); + assert_eq!(DequantSink::Scratch.resolve(0), DequantSink::Scratch); + } + + #[test] + fn scratch_files_are_unique_per_instance_and_deleted_on_drop() { + let dir = std::env::temp_dir().join("mummu-scratch-tests"); + std::fs::create_dir_all(&dir).expect("temp dir"); + + let a = ScratchFile::new(&dir).expect("names a scratch file"); + let b = ScratchFile::new(&dir).expect("names a second one"); + assert_ne!( + a.path(), + b.path(), + "two live loads must not share one scratch file" + ); + + let path = a.path().to_path_buf(); + std::fs::write(&path, b"payload").expect("writes"); + assert!(path.is_file()); + drop(a); + assert!(!path.exists(), "Drop removes the scratch file"); + + // And a stale file at the same name is cleared, not adopted: write + // one at b's path, then re-create a guard for it. + let stale = b.path().to_path_buf(); + std::fs::write(&stale, b"stale").expect("writes"); + drop(b); + assert!(!stale.exists()); + } + use super::*; #[test] diff --git a/crates/mummu/src/models/lfm2.rs b/crates/mummu/src/models/lfm2.rs index 68020fb..8eb03a9 100644 --- a/crates/mummu/src/models/lfm2.rs +++ b/crates/mummu/src/models/lfm2.rs @@ -16,7 +16,9 @@ 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, required_file}; +use crate::import::{ + CastFloatAdapter, DequantSink, ImportError, gguf_store, load_checked, required_file, +}; use crate::models::CausalLm; use crate::models::qwen2::EosIds; use crate::nn::{ @@ -451,19 +453,12 @@ pub fn load_from_gguf( }; let f = GgufFile::open(path).map_err(|e| parse(e.to_string()))?; let config = Lfm2Config::from_gguf(&f).map_err(parse)?; - let blob = f - .dequant_to_safetensors(&gguf_tensor_to_hf) - .map_err(|e| parse(e.to_string()))?; - assert!(blob.len() > 8, "a parsed GGUF yields a non-empty blob"); + // The scratch guard (Some only when the payload went to disk) must + // outlive `load_checked`: the store reads that file lazily. + let (base, _scratch) = gguf_store::(&f, &gguf_tensor_to_hf, DequantSink::Auto)?; let mut model = build::(&config, device); - // Type-level float dtype (`B::FloatElem`) — a probe tensor would follow - // the per-DEVICE default policy, which another backend alias sharing the - // device (Gpu vs GpuF16) may have flipped in this process. - let target_float = ::dtype(); - let mut store = SafetensorsStore::from_bytes(Some(blob)) - .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) - .allow_partial(true) + let mut store = base .with_key_remapping(r"^model\.", "") .with_key_remapping(r"(self_attn)\.out_proj\.", "$1.o_proj.") .with_key_remapping(r"(self_attn)\.q_layernorm\.weight$", "$1.q_norm.gamma") diff --git a/crates/mummu/src/models/olmoe.rs b/crates/mummu/src/models/olmoe.rs index 2da4751..cae274d 100644 --- a/crates/mummu/src/models/olmoe.rs +++ b/crates/mummu/src/models/olmoe.rs @@ -27,7 +27,7 @@ //! its dequant streams to a temp file (~28 GB of f32) that `burn-store` then //! mmaps back, so the payload is file-backed rather than charged to commit. -use std::path::{Path, PathBuf}; +use std::path::Path; use burn::module::Module; use burn::nn::{Embedding, EmbeddingConfig, Linear, LinearConfig, RmsNorm, RmsNormConfig}; @@ -35,7 +35,10 @@ 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, required_file}; +use crate::import::{ + CastFloatAdapter, DequantSink, ImportError, ScratchFile, gguf_store, load_checked, + required_file, +}; use crate::models::CausalLm; use crate::models::qwen2::{EosIds, gguf_f32, gguf_usize}; use crate::nn::{ @@ -315,23 +318,13 @@ pub fn load_from_gguf( }; let f = GgufFile::open(path).map_err(|e| parse(e.to_string()))?; let config = OlmoeConfig::from_gguf(&f).map_err(parse)?; - // Beside the GGUF, so the scratch write lands on the volume the weights - // already live on. - let scratch = FusedTemp::new(path.parent().unwrap_or(Path::new(".")))?; - let bytes = f - .dequant_to_safetensors_file(&gguf_tensor_to_hf, scratch.path()) - .map_err(|e| parse(e.to_string()))?; - assert!(bytes > 0, "a parsed GGUF yields a non-empty payload"); + // The scratch guard (Some only when the payload went to disk, which at + // OLMoE's size it always does) must outlive `load_checked`: the store + // reads that file lazily. + let (base, _scratch) = gguf_store::(&f, &gguf_tensor_to_hf, DequantSink::Auto)?; let mut model = build::(&config, device); - // The backend's float dtype, taken from the TYPE (`B::FloatElem`), never - // from a probe tensor (per-device default policy hazard). - let target_float = ::dtype(); - let mut store = install_remaps( - SafetensorsStore::from_file(scratch.path().to_path_buf()) - .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) - .allow_partial(true), - ); + let mut store = install_remaps(base); load_checked(&mut model, &mut store, path)?; Ok(LoadedOlmoe { model, @@ -392,57 +385,6 @@ fn fused_expert_target(name: &str, num_experts: usize) -> Option { /// /// 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 scratch safetensors file for the life of one load — the fused -/// checkpoint on the HF path, the dequantized payload on the GGUF path. -/// -/// It is as large as the weights it carries (13.8 GB fusing the 1B-7B, ~28 GB -/// dequantizing its Q4_K_M), 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. - /// - /// The name carries a process-unique counter as well as the pid: two - /// concurrent loads in one process must not choose the same scratch file - /// and interleave their writes into it. - fn new(dir: &Path) -> Result { - use std::sync::atomic::{AtomicU64, Ordering}; - static NEXT: AtomicU64 = AtomicU64::new(0); - assert!(!dir.as_os_str().is_empty(), "checkpoint dir must be named"); - let path = dir.join(format!( - "mummu-fused-{}-{}.safetensors.tmp", - std::process::id(), - NEXT.fetch_add(1, Ordering::Relaxed) - )); - // 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, @@ -470,7 +412,7 @@ pub fn load_from_dir( // 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 fused = ScratchFile::new(dir)?; let bytes = fuse_checkpoint_to_file( dir, &|name| Some(olmoe_hf_fuse(name, num_experts)), diff --git a/crates/mummu/src/models/qwen2.rs b/crates/mummu/src/models/qwen2.rs index 76908d4..9d65146 100644 --- a/crates/mummu/src/models/qwen2.rs +++ b/crates/mummu/src/models/qwen2.rs @@ -16,7 +16,9 @@ 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, required_file}; +use crate::import::{ + CastFloatAdapter, DequantSink, ImportError, gguf_store, load_checked, required_file, +}; use crate::models::CausalLm; use crate::nn::{ GqaAttention, GqaAttentionConfig, LayerKv, SwiGluMlp, SwiGluMlpConfig, causal_mask, rope_tables, @@ -358,19 +360,12 @@ pub fn load_from_gguf( }; let f = GgufFile::open(path).map_err(|e| parse(e.to_string()))?; let config = Qwen2Config::from_gguf(&f).map_err(parse)?; - let blob = f - .dequant_to_safetensors(&gguf_tensor_to_hf) - .map_err(|e| parse(e.to_string()))?; - assert!(blob.len() > 8, "a parsed GGUF yields a non-empty blob"); + // The scratch guard (Some only when the payload went to disk) must + // outlive `load_checked`: the store reads that file lazily. + let (base, _scratch) = gguf_store::(&f, &gguf_tensor_to_hf, DequantSink::Auto)?; let mut model = build::(&config, device); - // Type-level float dtype (`B::FloatElem`) — a probe tensor would follow - // the per-DEVICE default policy, which another backend alias sharing the - // device (Gpu vs GpuF16) may have flipped in this process. - let target_float = ::dtype(); - let mut store = SafetensorsStore::from_bytes(Some(blob)) - .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) - .allow_partial(true) + let mut store = base .with_key_remapping(r"^model\.", "") .with_key_remapping(r"(input_layernorm)\.weight$", "$1.gamma") .with_key_remapping(r"(post_attention_layernorm)\.weight$", "$1.gamma") diff --git a/crates/mummu/src/models/qwen3.rs b/crates/mummu/src/models/qwen3.rs index be273b0..49284a5 100644 --- a/crates/mummu/src/models/qwen3.rs +++ b/crates/mummu/src/models/qwen3.rs @@ -22,7 +22,9 @@ 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, required_file}; +use crate::import::{ + CastFloatAdapter, DequantSink, ImportError, gguf_store, load_checked, required_file, +}; use crate::models::CausalLm; use crate::models::qwen2::{EosIds, gguf_f32, gguf_usize}; use crate::nn::{ @@ -327,22 +329,12 @@ pub fn load_from_gguf( }; let f = GgufFile::open(path).map_err(|e| parse(e.to_string()))?; let config = Qwen3Config::from_gguf(&f).map_err(parse)?; - let blob = f - .dequant_to_safetensors(&gguf_tensor_to_hf) - .map_err(|e| parse(e.to_string()))?; - assert!(blob.len() > 8, "a parsed GGUF yields a non-empty blob"); + // The scratch guard (Some only when the payload went to disk) must + // outlive `load_checked`: the store reads that file lazily. + let (base, _scratch) = gguf_store::(&f, &gguf_tensor_to_hf, DequantSink::Auto)?; let mut model = build::(&config, device); - // The backend's float dtype, taken from the TYPE (`B::FloatElem`), never - // from a probe tensor: unspecified-dtype tensor creation follows the - // per-DEVICE default policy, which another backend alias sharing the - // device (Gpu vs GpuF16) may have flipped in this process. - let target_float = ::dtype(); - let mut store = install_remaps( - SafetensorsStore::from_bytes(Some(blob)) - .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) - .allow_partial(true), - ); + let mut store = install_remaps(base); load_checked(&mut model, &mut store, path)?; // A GGUF is self-contained — no sibling tokenizer_config.json in this path. Ok(LoadedQwen3 { From 5f6c1c2e0a6cd1c0c85b6401cda80de4b6f5aa40 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Fri, 21 Aug 2026 00:57:47 -0500 Subject: [PATCH 6/7] docs: record what the sink A/B settled, and why no planner hook Ticks the two items the last commit closed. The number worth keeping is that streaming to disk is free at ~6 GB of f32 payload (36.58 / 37.14 s vs an in-memory 33.90 / 35.93 / 42.15 s), and the reason it is free - burn-store's `SafetensorsStore::from_file` mmaps and materializes tensors lazily, so the read-back is page faults during a load that was already reading, not a second pass. Also records the design that was deliberately NOT built: this item floated a RAM-aware choice against the device inventory's free-RAM figure. Once streaming is free, a size threshold answers the question, and a planner dependency would be complexity bought for nothing. Co-Authored-By: Claude Opus 5 --- ROADMAP.md | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/ROADMAP.md b/ROADMAP.md index f02a92d..945a08b 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -797,7 +797,7 @@ a benchmark holds/improves its budget; README perf claims link an artifact. construction — every replaced site was unreachable on a 64-bit target — proven by the four `real_qwen2_gguf_*` legs on the Q4_K_M checkpoint (header parse, dequant-vs-true-weights, tokenizer byte identity, GPU load + decode). -- [ ] **Decide the dequant sink per model, not per code path** — `olmoe::load_from_gguf` takes the new +- [x] **Decide the dequant sink per model, not per code path** — `olmoe::load_from_gguf` takes the new streaming `dequant_to_safetensors_file`; qwen2 / qwen3 / lfm2 still take the in-memory `dequant_to_safetensors` (which is now 1x the payload rather than 2x, so they already got half the win for free). That split is a guess, not a measurement: the file variant trades a spike in commit @@ -806,12 +806,33 @@ a benchmark holds/improves its budget; README perf claims link an artifact. and Qwen3-0.6B, then either pick per-model or — better — pick automatically from `total_f32_bytes` against the device inventory's free-RAM figure (P6 already probes it), so a consumer never has to know. *(2026-08-20, discovered shipping the streaming sink.)* -- [ ] **`FusedTemp` is import-suite machinery living in one model file** — `models/olmoe.rs` owns the + *(2026-08-21) Measured, and the guess was wrong.* A/B on Qwen2.5-1.5B Q4_K_M (~6 GB of f32), + the same `real_qwen2_gguf_loads_and_decodes_on_gpu` load each way: **in-memory 33.90 / 35.93 / + 42.15 s vs scratch file 36.58 / 37.14 s** — the scratch numbers sit *inside* the in-memory run's + own spread, so at any payload of real size the disk round-trip is free and the doubled peak was + being paid for nothing. (n=2 on the scratch side: a concurrent routine held the shared cargo + target's lock and the third run never started.) `burn-store` is why it is free — + `SafetensorsStore::from_file` **mmaps** and materializes tensors lazily + (`safetensors_to_snapshots_lazy_file`), so the read-back is page faults during a load that was + already reading, not a second pass. So the choice is automatic, not per model: + `import::DequantSink` is `Memory` / `Scratch` / `Auto`, `Auto` keeps a payload in RAM only up to + `MAX_IN_MEMORY_DEQUANT_BYTES` (512 MiB — a ~1 GiB peak, free anywhere) and streams everything + else, and all four ports pass `Auto` so no consumer sees the knob. Proof the new default changed + nothing: qwen2 and qwen3 now take the scratch path they did not take before and their parity + numbers are **bit-identical** to the recorded ones (2.6614442413586614e-1 / 4.015608155114805e-1), + greedy sequences byte-identical, no scratch files left behind. Note the RAM-aware variant this + item floated (compare against the device inventory's free-RAM figure) is deliberately NOT built: + once streaming is free, a size threshold answers the question and a planner dependency would be + complexity bought for nothing. +- [x] **`FusedTemp` is import-suite machinery living in one model file** — `models/olmoe.rs` owns the scratch-file guard (create beside the weights, unique per process + counter, `Drop`-delete on every exit path), and it now serves BOTH import paths. Any other model large enough to want a streaming sink has to reach into `olmoe` or copy it. Promote it to `import.rs` when a second model needs it — not before, since a one-caller abstraction moved early is just churn. - *(2026-08-20.)* + *(2026-08-20.)* *(2026-08-21) The second caller arrived the same night* — `import::gguf_store` + needs the guard for every port, not just OLMoE's — so it is now `import::ScratchFile`, with a + unit test pinning the two properties that matter (two live guards never name the same file; + `Drop` deletes on every exit path). - [ ] **The f32 scratch file is a symptom, not the design** — both streaming sinks exist because `SafetensorsStore` is the only checked-load pipeline we have, so every import must first become f32 safetensors on disk. A GGUF **keep-quantized** load (P9) would skip the temp file entirely: From a77804ba795e562f6c60e01a4bdb7ca356ee2907 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Fri, 21 Aug 2026 01:04:08 -0500 Subject: [PATCH 7/7] docs: close the tokenizer-import item on evidence, not on bookkeeping The parent item listed two remaining pieces - SentencePiece `tokenizer.model` import, and calling the consistency validators from `load_from_dir` - and both have been `[x]` in split sub-items for a while. Rather than tick it on that, the whole surface was re-verified on real files this run: config-to-tokenizer id agreement on qwen3-0.6b, both SentencePiece legs (Unigram and BPE-type) byte-matching their own tokenizer.json, the 10-case template BYTE gate across qwen2/qwen3/lfm2, and the jinja fallback renderer byte-identical to the family renderers. Also files one thing the pass surfaced: `imported_render` reported FAILED because MUMMU_LFM2_DIR was unset, while both other legs were byte-identical - the same trap already recorded for `parity_gguf`. The item states the tension rather than assuming the fix: a gate that silently skips is a gate that does not exist, so the answer is a summarized "N legs skipped for missing fixtures" outcome, not a quiet skip. Co-Authored-By: Claude Opus 5 --- ROADMAP.md | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/ROADMAP.md b/ROADMAP.md index 945a08b..1dd7416 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -1025,7 +1025,7 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari rope-theta, vocab, tie-word-embeddings, …) so a model is **config-driven**, not hardcoded per checkpoint. *(2026-07-09) Per-architecture serde configs with validation (`Qwen2Config`, `Lfm2Config` incl. `layer_types` + auto-adjusted ff_dim); both real checkpoints parse and drive the build.* -- [ ] **Tokenizer + chat-template import** — HF `tokenizer.json` (fast), SentencePiece `tokenizer.model`, BPE +- [x] **Tokenizer + chat-template import** — HF `tokenizer.json` (fast), SentencePiece `tokenizer.model`, BPE merges/vocab; special-tokens map + the chat template from `tokenizer_config.json`. *(2026-07-18)* **`tokenizer_config.json` import shipped** (`mummu::tok_config::TokenizerConfig`, no new deps): parses the *conventions* HF keeps beside `tokenizer.json` — `add_bos_token`/`add_eos_token`, the BOS/EOS/PAD/UNK @@ -1053,6 +1053,25 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari qwen3-0.6b: all 26 ids pass `check_ids_against`, and `config.json` eos 151645 agrees with the resolved `<|im_end|>`. Remaining on this item: SentencePiece `tokenizer.model` import, and *calling* these validators from `load_from_dir` (config-driven EOS + template-vs-renderer consistency) — split below. + *(2026-08-21) Closed — every remaining piece it named is `[x]` below, and the whole surface was + re-verified on real files this run rather than taken on the sub-items' word:* + `real_tokenizer_config` (qwen3-0.6b config↔tokenizer ids agree), `real_spm` **both** legs + (Unigram via flan-t5-small and BPE-type via tinyllama, ids byte-matching each checkpoint's own + `tokenizer.json`), the **10-case template BYTE gate** (qwen2 / qwen3 / lfm2, plain + history + + tools), and `imported_render` under `--features jinja-template` (the fallback renderer's output + byte-identical to the family renderer at 142 / 748 / 324 B for qwen3, 157 / 379 B for lfm2, plus + the no-family-renderer fallback at 57 B). Anything further on tokenizers is its own item, not + this one. +- [ ] **The real-file gates panic on a missing env var instead of reporting a skip** — already recorded + operationally for `parity_gguf` (its lfm2 leg has no local GGUF, so the whole binary reports + FAILED even when every other leg passes); `imported_render` did the same thing this run, failing + on an unset `MUMMU_LFM2_DIR` while both other legs were byte-identical. The tension is real and + the current behaviour is the safer half of it — a gate that silently skips is a gate that does + not exist, which is exactly how a parity suite rots. So the fix is NOT "skip quietly": make the + missing-fixture case a distinct, *summarized* outcome — collect the unrunnable legs and print + one "N legs skipped for missing fixtures: …" line, so the run summary distinguishes "no fixture" + from "wrong answer" without ever letting the second hide inside the first. *(2026-08-21, + discovered re-verifying the tokenizer gates.)* - [x] **SentencePiece `tokenizer.model` import** — the `.model` proto tokenizer (Llama/Gemma/T5 family) that HF ships instead of a `tokenizer.json`; build the equivalent HF `tokenizers` pipeline (or convert), and byte-verify ids against a `tokenizer.json` of the same checkpoint where one exists. *(2026-07-18, split