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", 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 0f2e10b..1dd7416 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, @@ -719,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. @@ -728,11 +761,84 @@ 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). +- [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 + 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.)* + *(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-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: + 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/ @@ -919,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 @@ -947,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 @@ -1292,6 +1417,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 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 4716ece..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 { @@ -74,6 +90,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)] @@ -281,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 { @@ -375,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 { @@ -394,10 +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(usize::try_from(total_f32_bytes).expect("bounded above")); + let mut start = 0u64; for (index, info) in self.tensors.iter().enumerate() { let (name, shape) = match map(info) { Some(GgufMap::Rename(name)) => { @@ -422,26 +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).expect("string serializes"); - 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. @@ -504,7 +688,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 +735,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 +781,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 +825,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 +886,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 +899,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( @@ -1090,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) } @@ -1430,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/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 7884bdf..cae274d 100644 --- a/crates/mummu/src/models/olmoe.rs +++ b/crates/mummu/src/models/olmoe.rs @@ -23,9 +23,11 @@ //! 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}; +use std::path::Path; use burn::module::Module; use burn::nn::{Embedding, EmbeddingConfig, Linear, LinearConfig, RmsNorm, RmsNormConfig}; @@ -33,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::{ @@ -297,6 +302,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,20 +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)?; - 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, 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_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)?; Ok(LoadedOlmoe { model, @@ -381,48 +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 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, @@ -450,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 { 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");