From ca980ec2a8baf4ad6e8eebb462889be520e07601 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Mon, 13 Jul 2026 07:06:57 -0500 Subject: [PATCH 1/6] =?UTF-8?q?chore(deps):=20cargo=20update=20=E2=80=94?= =?UTF-8?q?=20uuid=201.23.5,=20zmij=201.0.22=20(pins=20respected)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit wgpu 30 / tokenizers 0.23 majors offered by cargo-upgrade were reverted: burn 0.21 + wgpu 29 + tokenizers 0.22 is the intentional parity-validated pin (see Cargo.toml). Gates green: 121 unit tests, GPU budget 107.4 ms / 13.3 tok/s, CPU budget 15.1 tok/s. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2e828e6..96e7259 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5840,9 +5840,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.23.4" +version = "1.23.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf80a72845275afea99e7f2b434723d3bc7e38470fcd1c7ed39a599c73319a53" +checksum = "ea5fab0d6c3c01ae70085a09cb03d4c7a1d6314e2b3e075392783396d724ca0a" [[package]] name = "v_frame" @@ -6630,9 +6630,9 @@ checksum = "b142a20ec14a91d5bc708c1dc21b080c550113d8aa77afa29635673a65dd02c5" [[package]] name = "zmij" -version = "1.0.21" +version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +checksum = "bd2f034a4bebf216c9e4b7083603e024cf930873fd67830cfb083c9fa33129d9" [[package]] name = "zopfli" From 03724d104b53e985501282f0baa6bff7df9f235a Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Mon, 13 Jul 2026 07:43:41 -0500 Subject: [PATCH 2/6] =?UTF-8?q?feat(p3):=20GGUF=20->=20running=20model=20?= =?UTF-8?q?=E2=80=94=20full=20dequant=20suite=20+=20qwen2::load=5Ffrom=5Fg?= =?UTF-8?q?guf?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - dequantize every GGUF storage dtype: adds Q4_0/Q4_1/Q5_0/Q5_1 legacy blocks and Q2_K/Q3_K/Q5_K superblocks (exact ggml-quants ports, each with hand-computed block tests); Q8_1/Q8_K stay loud errors (activation formats, never tensor storage) - GgufFile::dequant_to_safetensors: bridge a GGUF onto the exact checked-load pipeline safetensors uses (in-memory blob, dims reversed = row-major HF layout, unmapped tensor names error loudly, 48 GiB RAM bound) - Qwen2Config::from_gguf: hyperparameters from qwen2.* metadata; vocab from the embedding tensor; EOS from tokenizer.ggml.eos_token_id - qwen2::load_from_gguf: one .gguf file -> a running model - Qwen2 gains an optional untied lm_head: llama.cpp GGUFs materialize the tied head as a separate (higher-precision) output.weight; also unlocks untied safetensors checkpoints (Qwen2-7B class) Proof (real weights, 4070 Ti SUPER): Qwen2.5-1.5B Q4_K_M greedy-decodes '2+2 equals 4.'; first-token top-1 identical to the bf16 build, top-5 overlap 4/5, logit cosine 0.977. Parity gate re-passed byte-identically after the lm_head change (max dlogit 2.670e-5; Ollama greedy exact). 133 unit tests; f16, real-inference, and budget gates green (GPU 107.4 ms / 13.6 tok/s, CPU 14.8 tok/s). Co-Authored-By: Claude Opus 4.8 --- README.md | 15 +- ROADMAP.md | 24 +- crates/mummu/src/gguf.rs | 429 ++++++++++++++++++++++++++++++- crates/mummu/src/models/qwen2.rs | 278 +++++++++++++++++++- crates/mummu/tests/real_gguf.rs | 95 +++++++ 5 files changed, 820 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 1e6d1a0..d4b7d9b 100644 --- a/README.md +++ b/README.md @@ -64,12 +64,15 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys - **Model management** — `ModelManager` gives settings UIs the whole lifecycle over a declarative model catalog (`registry::ModelSpec`): install with per-chunk download progress, `is_installed`, per-model disk usage, and traversal-safe removal; model switching rides `ModelSlot`. -- **GGUF import (container + K-quant dequant)** — `mummu::gguf` parses the llama.cpp container (typed, - bounded metadata; fully validated tensor table) and dequantizes F32/F16/BF16, Q8_0, and the - **Q4_K/Q6_K superblocks** to f32 — proven against the model's true weights: on the real Qwen2.5-1.5B - Q4_K_M file the F32 norms come back **bit-exact** vs the bf16 safetensors of the same checkpoint and - Q4_K embedding rows dequantize at cosine 0.9975 (`tests/real_gguf.rs`). Next: the remaining quant - types and GGUF→model load (tracked in P3/P9). +- **GGUF import, end to end** — `mummu::gguf` parses the llama.cpp container (typed, bounded metadata; + fully validated tensor table) and dequantizes **every storage dtype** (F32/F16/BF16, the legacy + Q4_0/Q4_1/Q5_0/Q5_1/Q8_0 blocks, and the Q2_K–Q6_K superblocks) to f32; `qwen2::load_from_gguf` + turns the one file into a running model — hyperparameters from the GGUF metadata, weights bridged + through the same checked-load pipeline as safetensors (tied *and* untied lm-heads). Proven against + the model's true weights (`tests/real_gguf.rs`): F32 norms **bit-exact** vs the bf16 safetensors of + the same checkpoint, Q4_K rows at cosine 0.9975, and the real Qwen2.5-1.5B **Q4_K_M file greedy-decodes + "2+2 equals 4." on the GPU** with first-token top-1 identical to the bf16 build (logit cosine 0.977). + Next: tokenizer-from-GGUF metadata and the LFM2 GGUF map (tracked in P3). - **Hub downloads** — streaming HuggingFace fetches into the model cache: resumable (`.part` + HTTP Range, proven byte-identical after an interrupted transfer), length-verified, shard-index aware, with a per-chunk progress callback; verified end-to-end by downloading all-MiniLM and embedding with it. diff --git a/ROADMAP.md b/ROADMAP.md index 889bef9..9d7cf54 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -162,7 +162,7 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari no adapter chaining; `.bin`-era checkpoints are f32), decoder loaders adopt `weights_file` when a real `.pth` decoder checkpoint exists to verify against, and `hub::fetch_model` learning the `pytorch_model.bin` fallback. -- [ ] **GGUF** (llama.cpp) — parse the GGUF container (metadata KV + tensor table), map tensors to modules, +- [x] **GGUF** (llama.cpp) — parse the GGUF container (metadata KV + tensor table), map tensors to modules, and **dequantize** Q4/Q5/Q8/K-quant blocks into Burn tensors (or hand keep-quantized to P9). GGUF is how most small models are distributed — this makes the whole ecosystem importable. *(2026-07-10 research)* K-quant superblocks are 256 values: Q4_K = fp16 d + fp16 dmin + 12 B of 6-bit sub-scales/mins + 128 B @@ -190,6 +190,28 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari embedding rows hit cosine **0.9975** vs truth (garbage layout ⇒ ≈ 0) — 5 hand-computed-block unit tests (121 total). NEXT slice: remaining dequants (Q4_0/Q5/Q2_K/Q3_K/Q5_K), GGUF→model load (name remap + ggml dim-order transpose), tokenizer-from-GGUF-metadata.* + *(2026-07-13) **GGUF→running model shipped.** Every storage dtype now dequantizes (added + Q4_0/Q4_1/Q5_0/Q5_1 + K-quants Q2_K/Q3_K/Q5_K, exact ggml ports, hand-computed block tests — + 133 unit tests total); `GgufFile::dequant_to_safetensors` bridges a GGUF onto the SAME checked-load + pipeline as safetensors (dims reversed = row-major HF layout, unmapped tensor names are loud + errors); `Qwen2Config::from_gguf` reads hyperparameters from `qwen2.*` metadata (vocab from the + embedding tensor, EOS from tokenizer metadata); `qwen2::load_from_gguf` = one file → running model. + The Qwen2 module gained an optional **untied lm_head** (llama.cpp GGUFs materialize the tied head + as a separate higher-precision `output.weight` — the real Q4_K_M carries it as Q6_K; also unlocks + untied safetensors like Qwen2-7B). REAL-GPU proof (`real_gguf.rs`): the Q4_K_M file alone + greedy-decodes "2+2 equals 4.", first-token top-1 identical to the bf16 build, top-5 overlap 4/5, + logit cosine 0.977 (28 layers of Q4_K drift; a layout bug reads ≈ 0). Parity gate re-passed + byte-identically after the lm_head change (max |Δlogit| 2.670e-5, Ollama greedy exact); f16 + + budget gates green (GPU 107.4 ms / 13.6 tok/s, CPU 14.8 tok/s).* +- [ ] **Tokenizer-from-GGUF metadata** — build the HF `tokenizers` pipeline from `tokenizer.ggml.*` + (tokens, merges, token types, BPE pre-tokenizer regex) so a GGUF needs no sibling + `tokenizer.json`; byte-verify token ids against the HF tokenizer of the same checkpoint. +- [ ] **LFM2 GGUF name map** — extend `load_from_gguf` to the LFM2/LFM2.5 hybrid (llama.cpp `lfm2` + arch: `shortconv.*` tensor names, `lfm2.*` metadata keys) once a same-weights GGUF is validated. +- [ ] **Quantized-reference parity leg for GGUF loads** — the end-to-end test compares against the bf16 + build (quantization drift bounded, not exact); a strict leg needs llama.cpp itself running the + SAME quantized file (`llama-server` raw `/completion`, `n_probs` logprobs — see the P7 LFM2.5 + reference item's caveats) to assert our dequant matches ggml's compute path token-for-token. - [ ] **GPTQ / AWQ** (HF safetensors) — import the calibration-quantized int4/int8 layouts most "quantized on the Hub" models ship as (a `.safetensors` payload + a quant config), dequant or keep-quant into Burn. - [ ] **ONNX** (optional) — `burn-import` ONNX→Burn for models distributed as ONNX graphs. diff --git a/crates/mummu/src/gguf.rs b/crates/mummu/src/gguf.rs index bad834a..2aa9680 100644 --- a/crates/mummu/src/gguf.rs +++ b/crates/mummu/src/gguf.rs @@ -37,6 +37,10 @@ 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). +const MAX_DEQUANT_BYTES: u64 = 48 << 30; + /// What went wrong reading a GGUF header. #[derive(Debug, thiserror::Error)] pub enum GgufError { @@ -321,6 +325,71 @@ impl GgufFile { Ok(out) } + /// Dequantize every tensor to f32 and serialize the result as an + /// in-memory **safetensors** file — the bridge onto the exact store + /// pipeline (adapters, key remaps, checked load) the safetensors path + /// already trusts. `rename` maps each GGUF tensor name to the name the + /// blob should carry (HF-checkpoint naming, so per-model remap tables + /// apply unchanged); an unmapped tensor is a loud error, never a skip — + /// a name this crate doesn't recognize means weights would silently + /// vanish from the model. + /// + /// Shapes are the GGUF dims **reversed**: ggml orders dims + /// fastest-varying-first, so the raw payload bytes are exactly the + /// row-major layout of the reversed shape — same bytes, HF convention. + pub fn dequant_to_safetensors( + &self, + rename: &dyn Fn(&str) -> 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 { + path: self.path.display().to_string(), + what: "dequantized f32 payload bytes", + count: total_f32_bytes, + bound: MAX_DEQUANT_BYTES, + }); + } + let bad = |index: usize, reason: String| GgufError::BadTensor { + path: self.path.display().to_string(), + index, + reason, + }; + + 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")); + for (index, info) in self.tensors.iter().enumerate() { + let Some(name) = rename(&info.name) else { + return Err(bad(index, format!("unmapped tensor name '{}'", info.name))); + }; + 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 shape: Vec = info.dims.iter().rev().copied().collect(); + let json_name = serde_json::to_string(&name).expect("string serializes"); + if index > 0 { + header.push(','); + } + header.push_str(&format!( + "{json_name}:{{\"dtype\":\"F32\",\"shape\":{shape:?},\"data_offsets\":[{start},{end}]}}", + end = data.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) + } + /// Look up a metadata value by exact key. #[must_use] pub fn get(&self, key: &str) -> Option<&GgufValue> { @@ -568,9 +637,12 @@ impl Reader { // ---- Dequantization ------------------------------------------------------ // -// Exact ports of ggml's reference dequantizers (ggml-quants.c) for the types -// a Q4_K_M file actually carries: plain floats, Q8_0, and the K-quant -// superblocks Q4_K / Q6_K. Layouts follow `GgmlType::bytes_per_block`. +// Exact ports of ggml's reference dequantizers (ggml-quants.c) for every +// dtype llama.cpp stores model weights in: plain floats, the 32-element +// legacy blocks Q4_0/Q4_1/Q5_0/Q5_1/Q8_0, and the 256-element K-quant +// superblocks Q2_K–Q6_K. Layouts follow `GgmlType::bytes_per_block`. +// Q8_1/Q8_K are activation formats (dot-product scratch), never tensor +// storage — they stay loud errors. /// Dequantize a whole tensor payload to f32. `bytes` must be whole blocks of /// `dtype` (guaranteed for payload slices sized by [`GgufTensorInfo::byte_len`]). @@ -594,8 +666,15 @@ pub fn dequantize(dtype: GgmlType, bytes: &[u8]) -> Result, String> { u32::from(u16::from_le_bytes([block[0], block[1]])) << 16, )); } + GgmlType::Q4_0 => dequant_q4_0(block, &mut out), + GgmlType::Q4_1 => dequant_q4_1(block, &mut out), + GgmlType::Q5_0 => dequant_q5_0(block, &mut out), + GgmlType::Q5_1 => dequant_q5_1(block, &mut out), GgmlType::Q8_0 => dequant_q8_0(block, &mut out), + GgmlType::Q2_K => dequant_q2_k(block, &mut out), + GgmlType::Q3_K => dequant_q3_k(block, &mut out), GgmlType::Q4_K => dequant_q4_k(block, &mut out), + GgmlType::Q5_K => dequant_q5_k(block, &mut out), GgmlType::Q6_K => dequant_q6_k(block, &mut out), other => return Err(format!("dequant for {other:?} is not implemented yet")), } @@ -618,6 +697,132 @@ fn dequant_q8_0(block: &[u8], out: &mut Vec) { out.extend(block[2..34].iter().map(|&q| d * f32::from(q as i8))); } +/// Q4_0: f16 scale + 16 bytes of 4-bit quants; `x = d·(q − 8)` — all 16 low +/// nibbles are elements 0..16, the high nibbles elements 16..32. +fn dequant_q4_0(block: &[u8], out: &mut Vec) { + assert_eq!(block.len(), 18, "Q4_0 block is 18 bytes"); + let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); + let qs = &block[2..18]; + out.extend(qs.iter().map(|&b| d * (f32::from(b & 0x0F) - 8.0))); + out.extend(qs.iter().map(|&b| d * (f32::from(b >> 4) - 8.0))); +} + +/// Q4_1: f16 scale + f16 min + 16 bytes of 4-bit quants; `x = d·q + m`. +fn dequant_q4_1(block: &[u8], out: &mut Vec) { + assert_eq!(block.len(), 20, "Q4_1 block is 20 bytes"); + let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); + let m = f16_to_f32(u16::from_le_bytes([block[2], block[3]])); + let qs = &block[4..20]; + out.extend(qs.iter().map(|&b| d * f32::from(b & 0x0F) + m)); + out.extend(qs.iter().map(|&b| d * f32::from(b >> 4) + m)); +} + +/// Q5_0: f16 scale + 4 B of packed 5th bits + 16 B of 4-bit quants; +/// `x = d·(q − 16)` with bit `j` of `qh` topping up element `j`. +fn dequant_q5_0(block: &[u8], out: &mut Vec) { + assert_eq!(block.len(), 22, "Q5_0 block is 22 bytes"); + let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); + let qh = u32::from_le_bytes([block[2], block[3], block[4], block[5]]); + let qs = &block[6..22]; + #[allow(clippy::cast_possible_truncation)] // masked to one nibble bit + out.extend(qs.iter().enumerate().map(|(j, &b)| { + let hi = ((qh >> j) << 4) as u8 & 0x10; + d * (f32::from((b & 0x0F) | hi) - 16.0) + })); + #[allow(clippy::cast_possible_truncation)] // masked to one nibble bit + out.extend(qs.iter().enumerate().map(|(j, &b)| { + let hi = (qh >> (j + 12)) as u8 & 0x10; + d * (f32::from((b >> 4) | hi) - 16.0) + })); +} + +/// Q5_1: f16 scale + f16 min + 4 B packed 5th bits + 16 B quants; `x = d·q + m`. +fn dequant_q5_1(block: &[u8], out: &mut Vec) { + assert_eq!(block.len(), 24, "Q5_1 block is 24 bytes"); + let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); + let m = f16_to_f32(u16::from_le_bytes([block[2], block[3]])); + let qh = u32::from_le_bytes([block[4], block[5], block[6], block[7]]); + let qs = &block[8..24]; + #[allow(clippy::cast_possible_truncation)] // masked to one nibble bit + out.extend(qs.iter().enumerate().map(|(j, &b)| { + let hi = ((qh >> j) << 4) as u8 & 0x10; + d * f32::from((b & 0x0F) | hi) + m + })); + #[allow(clippy::cast_possible_truncation)] // masked to one nibble bit + out.extend(qs.iter().enumerate().map(|(j, &b)| { + let hi = (qh >> (j + 12)) as u8 & 0x10; + d * f32::from((b >> 4) | hi) + m + })); +} + +/// Q2_K: 256-element superblock — 16 packed (scale, min) nibbles + 64 B of +/// 2-bit quants + f16 d + f16 dmin; `x = d·sc·q − dmin·m` over 16 sub-blocks +/// of 16. +fn dequant_q2_k(block: &[u8], out: &mut Vec) { + assert_eq!(block.len(), 84, "Q2_K superblock is 84 bytes"); + let scales = &block[0..16]; + let qs = &block[16..80]; + let d = f16_to_f32(u16::from_le_bytes([block[80], block[81]])); + let dmin = f16_to_f32(u16::from_le_bytes([block[82], block[83]])); + let mut is = 0; + // Two halves of 128 values; each half reads 32 quant bytes at 4 shifts. + for q in [&qs[0..32], &qs[32..64]] { + for shift in [0u8, 2, 4, 6] { + for part in [&q[0..16], &q[16..32]] { + let sc = scales[is]; + is += 1; + let dl = d * f32::from(sc & 0x0F); + let ml = dmin * f32::from(sc >> 4); + out.extend(part.iter().map(|&b| dl * f32::from((b >> shift) & 3) - ml)); + } + } + } + assert_eq!(is, 16, "16 sub-block scales consumed"); +} + +/// Unpack Q3_K's 12 packed scale bytes into 16 signed 6-bit sub-scales +/// (ggml's kmask bit dance, done bytewise). +fn q3_k_scales(packed: &[u8]) -> [i8; 16] { + assert_eq!(packed.len(), 12, "Q3_K scale block is 12 bytes"); + let mut sc = [0i8; 16]; + #[allow(clippy::cast_possible_wrap)] // 6-bit values reinterpret exactly + for j in 0..4 { + let hi = packed[8 + j]; // 2-bit tops for slots j, j+4, j+8, j+12 + sc[j] = ((packed[j] & 0x0F) | ((hi & 3) << 4)) as i8; + sc[j + 4] = ((packed[j + 4] & 0x0F) | (((hi >> 2) & 3) << 4)) as i8; + sc[j + 8] = ((packed[j] >> 4) | (((hi >> 4) & 3) << 4)) as i8; + sc[j + 12] = ((packed[j + 4] >> 4) | ((hi >> 6) << 4)) as i8; + } + sc +} + +/// Q3_K: 256-element superblock — 32 B high-bit mask + 64 B of 2-bit quants +/// + 12 B packed 6-bit sub-scales + f16 d; `x = d·(sc − 32)·(q − hm·4)`. +fn dequant_q3_k(block: &[u8], out: &mut Vec) { + assert_eq!(block.len(), 110, "Q3_K superblock is 110 bytes"); + let hmask = &block[0..32]; + let qs = &block[32..96]; + let scales = q3_k_scales(&block[96..108]); + let d = f16_to_f32(u16::from_le_bytes([block[108], block[109]])); + let mut is = 0; + let mut m: u8 = 1; + for q in [&qs[0..32], &qs[32..64]] { + for shift in [0u8, 2, 4, 6] { + for base in [0usize, 16] { + let dl = d * f32::from(i16::from(scales[is]) - 32); + is += 1; + for l in base..base + 16 { + let low = i16::from((q[l] >> shift) & 3); + let sub = if hmask[l] & m == 0 { 4 } else { 0 }; + out.push(dl * f32::from(low - sub)); + } + } + m <<= 1; // the hmask bit advances per (half, shift) pair + } + } + assert_eq!(is, 16, "16 sub-block scales consumed"); +} + /// The Q4_K/Q5_K 6-bit (scale, min) pair for sub-block `j` — ggml's /// `get_scale_min_k4`. fn scale_min_k4(scales: &[u8], j: usize) -> (f32, f32) { @@ -652,6 +857,38 @@ fn dequant_q4_k(block: &[u8], out: &mut Vec) { } } +/// Q5_K: 256-element superblock — f16 d + f16 dmin + 12 B packed 6-bit +/// (scale, min) pairs + 32 B of 5th bits + 128 B of 4-bit quants; +/// `x = d·sc·q − dmin·m` with two `qh` bits per byte per chunk. +fn dequant_q5_k(block: &[u8], out: &mut Vec) { + assert_eq!(block.len(), 176, "Q5_K superblock is 176 bytes"); + let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]])); + let dmin = f16_to_f32(u16::from_le_bytes([block[2], block[3]])); + let scales = &block[4..16]; + let qh = &block[16..48]; + let qs = &block[48..176]; + let mut u1: u8 = 1; + let mut u2: u8 = 2; + // 4 chunks of 64 values; each chunk reads 32 quant bytes — low nibbles + // first — and one bit-plane pair of the shared 32-byte qh. + for chunk in 0..4 { + let (sc1, m1) = scale_min_k4(scales, chunk * 2); + let (sc2, m2) = scale_min_k4(scales, chunk * 2 + 1); + let q = &qs[chunk * 32..chunk * 32 + 32]; + out.extend(q.iter().zip(qh).map(|(&b, &h)| { + let top = if h & u1 == 0 { 0 } else { 16 }; + d * sc1 * f32::from((b & 0x0F) + top) - dmin * m1 + })); + out.extend(q.iter().zip(qh).map(|(&b, &h)| { + let top = if h & u2 == 0 { 0 } else { 16 }; + d * sc2 * f32::from((b >> 4) + top) - dmin * m2 + })); + u1 <<= 2; + u2 <<= 2; + } + assert_eq!(u1, 0, "four bit-plane pairs consumed"); // 1<<8 wraps to 0 +} + /// Q6_K: 256-element superblock — 128 B low-4 + 64 B high-2 + 16 i8 /// sub-scales + f16 d; `x = d·sc·(q − 32)`. fn dequant_q6_k(block: &[u8], out: &mut Vec) { @@ -766,9 +1003,21 @@ mod tests { self.buf.extend_from_slice(&self.tensors); self.buf } + + /// Header + alignment padding + tensor payload bytes (offsets in the + /// tensor table are relative to the padded data start). + fn build_with_payload(self, payload: &[u8]) -> Vec { + let mut buf = self.build(); + let data_offset = (buf.len() as u64).div_ceil(DEFAULT_ALIGNMENT) * DEFAULT_ALIGNMENT; + buf.resize(usize::try_from(data_offset).expect("small test file"), 0); + buf.extend_from_slice(payload); + buf + } } - fn open_bytes(bytes: &[u8]) -> Result { + /// Write `bytes` to a fresh temp file and run `f` on the parse result + /// while the file still exists (payload reads re-open the path). + fn with_gguf_bytes(bytes: &[u8], f: impl FnOnce(Result) -> R) -> R { use std::sync::atomic::{AtomicU64, Ordering}; // Parallel tests in one process must never share a temp file. static NEXT: AtomicU64 = AtomicU64::new(0); @@ -779,14 +1028,18 @@ mod tests { std::process::id(), NEXT.fetch_add(1, Ordering::Relaxed) )); - let mut f = File::create(&path).expect("temp file"); - f.write_all(bytes).expect("write"); - drop(f); - let result = GgufFile::open(&path); + let mut file = File::create(&path).expect("temp file"); + file.write_all(bytes).expect("write"); + drop(file); + let result = f(GgufFile::open(&path)); let _ = std::fs::remove_file(&path); result } + fn open_bytes(bytes: &[u8]) -> Result { + with_gguf_bytes(bytes, |r| r) + } + #[test] fn minimal_file_round_trips() { let bytes = TestGguf::new() @@ -958,11 +1211,169 @@ mod tests { assert_eq!(out[64], 0.0); } + #[test] + fn q4_0_and_q4_1_blocks_match_hand_computation() { + // Q4_0: symmetric around 8 — low nibbles are elements 0..16. + let mut b = vec![0u8; 18]; + b[0..2].copy_from_slice(&f16_bytes(2.0)); + b[2] = 0x31; // low nibble 1 → elem 0, high nibble 3 → elem 16 + let out = dequantize(GgmlType::Q4_0, &b).unwrap(); + assert_eq!(out.len(), 32); + assert_eq!(out[0], (1.0 - 8.0) * 2.0); + assert_eq!(out[16], (3.0 - 8.0) * 2.0); + assert_eq!(out[1], -16.0); // zero nibble is −8·d, not 0 + + // Q4_1: affine — the min shifts every element. + let mut b = vec![0u8; 20]; + b[0..2].copy_from_slice(&f16_bytes(2.0)); + b[2..4].copy_from_slice(&f16_bytes(1.0)); + b[4] = 0x31; + let out = dequantize(GgmlType::Q4_1, &b).unwrap(); + assert_eq!(out[0], 1.0 * 2.0 + 1.0); + assert_eq!(out[16], 3.0 * 2.0 + 1.0); + assert_eq!(out[1], 1.0); + } + + #[test] + fn q5_0_and_q5_1_high_bits_land_on_the_right_elements() { + // Q5_0: qh bit 0 tops element 0, bit 16 tops element 16. + let mut b = vec![0u8; 22]; + b[0..2].copy_from_slice(&f16_bytes(1.0)); + b[2..6].copy_from_slice(&(1u32 | (1 << 16)).to_le_bytes()); + b[6] = 0x21; // low nibble 1 → elem 0, high nibble 2 → elem 16 + let out = dequantize(GgmlType::Q5_0, &b).unwrap(); + assert_eq!(out.len(), 32); + assert_eq!(out[0], (1.0 + 16.0) - 16.0); + assert_eq!(out[16], (2.0 + 16.0) - 16.0); + assert_eq!(out[1], -16.0); // no high bit, zero nibble + + // Q5_1: same bit packing, affine. + let mut b = vec![0u8; 24]; + b[0..2].copy_from_slice(&f16_bytes(1.0)); + b[2..4].copy_from_slice(&f16_bytes(1.0)); + b[4..8].copy_from_slice(&1u32.to_le_bytes()); + b[8] = 0x01; + let out = dequantize(GgmlType::Q5_1, &b).unwrap(); + assert_eq!(out[0], (1.0 + 16.0) * 1.0 + 1.0); + assert_eq!(out[16], 1.0); // high nibble 0, qh bit 16 unset → just m + } + + #[test] + fn q2_k_superblock_matches_hand_computation() { + let mut b = vec![0u8; 84]; + b[80..82].copy_from_slice(&f16_bytes(1.0)); // d + b[82..84].copy_from_slice(&f16_bytes(0.5)); // dmin + b[0] = 0x12; // sub 0: sc=2, min=1 + b[1] = 0x01; // sub 1: sc=1, min=0 + b[8] = 0x0F; // sub 8 (second half, shift 0): sc=15, min=0 + b[16] = 3; // qs[0] bits 0–1 → elem 0 + b[32] = 1; // qs[16] → elem 16 (sub 1) + b[48] = 2; // qs[32] → elem 128 (second half) + let out = dequantize(GgmlType::Q2_K, &b).unwrap(); + assert_eq!(out.len(), 256); + assert_eq!(out[0], 2.0 * 3.0 - 0.5); // d·sc·q − dmin·m + assert_eq!(out[1], -0.5); // zero quant still subtracts the min + assert_eq!(out[16], 1.0); + assert_eq!(out[128], 15.0 * 2.0); // second half reads qs[32..] + } + + #[test] + fn q3_k_superblock_matches_hand_computation() { + let mut b = vec![0u8; 110]; + b[108..110].copy_from_slice(&f16_bytes(1.0)); // d + // scales: slot 0 = 2|32 = 34 → dl 2; slot 1 = 1|32 = 33 → dl 1; + // slot 8 = (0x32>>4)|0 = 3 → dl 3−32 = −29 (tests the >>4 packing). + b[96] = 0x32; + b[97] = 0x01; + b[104] = 0b10; // top bits of slot 0 + b[105] = 0b10; // top bits of slot 1 + b[0] = 1; // hmask[0] bit 0 → element 0 keeps its high bit (no −4) + b[16] = 1; // hmask[16] bit 0 → element 16 too + b[32] = 3; // qs[0] → elem 0 + b[48] = 2; // qs[16] → elem 16 + let out = dequantize(GgmlType::Q3_K, &b).unwrap(); + assert_eq!(out.len(), 256); + assert_eq!(out[0], 2.0 * 3.0); // high bit set → q unshifted + assert_eq!(out[1], 2.0 * -4.0); // high bit clear → q − 4 + assert_eq!(out[16], 1.0 * 2.0); + // Second half, shift 0 (is=8): hmask[0] bit 4 clear → (0−4)·(3−32). + assert_eq!(out[128], -4.0 * (3.0 - 32.0)); + } + + #[test] + fn q5_k_superblock_matches_hand_computation() { + let mut b = vec![0u8; 176]; + b[0..2].copy_from_slice(&f16_bytes(1.0)); // d + b[2..4].copy_from_slice(&f16_bytes(1.0)); // dmin + b[4] = 2; // sub 0 scale + b[5] = 3; // sub 1 scale + b[8] = 1; // sub 0 min + b[16] = 1; // qh[0] bit 0 → elem 0 gets +16 (u1 = 1) + b[48] = 0x21; // ql[0]: low 1 → elem 0, high 2 → elem 32 + let out = dequantize(GgmlType::Q5_K, &b).unwrap(); + assert_eq!(out.len(), 256); + assert_eq!(out[0], 2.0 * (1.0 + 16.0) - 1.0); + assert_eq!(out[1], -1.0); // zero quant still subtracts the min + assert_eq!(out[32], 3.0 * 2.0); // qh bit 1 unset → no +16; m1 = 0 + } + #[test] fn dequant_rejects_partial_blocks_and_unimplemented_types() { assert!(dequantize(GgmlType::Q8_0, &[0u8; 33]).is_err()); assert!(dequantize(GgmlType::Q8_0, &[]).is_err()); - assert!(dequantize(GgmlType::Q2_K, &[0u8; 84]).is_err()); + // Q8_K is an activation format, never tensor storage. + assert!(dequantize(GgmlType::Q8_K, &[0u8; 292]).is_err()); + } + + #[test] + fn dequant_to_safetensors_reverses_dims_and_round_trips_bytes() { + // One F32 tensor with ggml dims [2, 3] and payload 1..=6. + let payload: Vec = (1..=6).flat_map(|v| (v as f32).to_le_bytes()).collect(); + let bytes = TestGguf::new() + .kv_str("general.architecture", "qwen2") + .tensor("token_embd.weight", &[2, 3], 0, 0) + .build_with_payload(&payload); + with_gguf_bytes(&bytes, |f| { + let f = f.expect("parses"); + let blob = f + .dequant_to_safetensors(&|n| Some(format!("model.{n}"))) + .expect("serializes"); + let header_len = u64::from_le_bytes(blob[0..8].try_into().unwrap()); + let json: serde_json::Value = + serde_json::from_slice(&blob[8..8 + usize::try_from(header_len).unwrap()]) + .expect("header is valid JSON"); + let entry = &json["model.token_embd.weight"]; + assert_eq!(entry["dtype"], "F32"); + assert_eq!(entry["shape"], serde_json::json!([3, 2])); // reversed + assert_eq!(entry["data_offsets"], serde_json::json!([0, 24])); + // F32 → f32 is byte-identical. + assert_eq!( + &blob[8 + usize::try_from(header_len).unwrap()..], + &payload[..] + ); + + // An unmapped tensor is a loud error, never a silent skip. + assert!(matches!( + f.dequant_to_safetensors(&|_| None), + Err(GgufError::BadTensor { .. }) + )); + }); + } + + #[test] + fn dequant_to_safetensors_rejects_rename_collisions() { + let payload = [0u8; 64]; // two 8-element F32 tensors, offsets 0 and 32 + let bytes = TestGguf::new() + .tensor("a.weight", &[8], 0, 0) + .tensor("b.weight", &[8], 0, 32) + .build_with_payload(&payload); + with_gguf_bytes(&bytes, |f| { + let f = f.expect("parses"); + assert!(matches!( + f.dequant_to_safetensors(&|_| Some("same".into())), + Err(GgufError::BadTensor { .. }) + )); + }); } #[test] diff --git a/crates/mummu/src/models/qwen2.rs b/crates/mummu/src/models/qwen2.rs index 941d715..9cfce28 100644 --- a/crates/mummu/src/models/qwen2.rs +++ b/crates/mummu/src/models/qwen2.rs @@ -11,10 +11,11 @@ use std::path::Path; use burn::module::Module; -use burn::nn::{Embedding, EmbeddingConfig, RmsNorm, RmsNormConfig}; +use burn::nn::{Embedding, EmbeddingConfig, Linear, LinearConfig, RmsNorm, RmsNormConfig}; use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, SafetensorsStore}; use burn::tensor::{Int, Tensor, TensorData, backend::Backend}; +use crate::gguf::{GgufFile, GgufValue}; use crate::import::{CastFloatAdapter, ImportError, load_checked, required_file}; use crate::models::CausalLm; use crate::nn::{ @@ -63,6 +64,21 @@ impl EosIds { } } +/// A required GGUF metadata integer, as usize. +fn gguf_usize(f: &GgufFile, key: &str) -> Result { + f.get(key) + .and_then(GgufValue::as_u64) + .and_then(|v| usize::try_from(v).ok()) + .ok_or_else(|| format!("missing or non-integer GGUF metadata '{key}'")) +} + +/// A required GGUF metadata f32. +fn gguf_f32(f: &GgufFile, key: &str) -> Result { + f.get(key) + .and_then(GgufValue::as_f32) + .ok_or_else(|| format!("missing or non-f32 GGUF metadata '{key}'")) +} + impl Qwen2Config { /// Parse `config.json` bytes; derives `head_dim` when absent. pub fn from_json_bytes(bytes: &[u8]) -> Result { @@ -74,6 +90,62 @@ impl Qwen2Config { Ok(cfg) } + /// Hyperparameters from a GGUF header's `qwen2.*` metadata — a GGUF file + /// is self-contained, no `config.json` beside it. `vocab_size` comes from + /// the embedding tensor (llama.cpp may pad it past the tokenizer vocab). + pub fn from_gguf(f: &GgufFile) -> Result { + let arch = f.architecture().unwrap_or(""); + if arch != "qwen2" { + return Err(format!("GGUF architecture '{arch}' is not qwen2")); + } + let hidden_size = gguf_usize(f, "qwen2.embedding_length")?; + let num_attention_heads = gguf_usize(f, "qwen2.attention.head_count")?; + let embd = f + .tensor("token_embd.weight") + .ok_or("GGUF has no token_embd.weight tensor")?; + if embd.dims.len() != 2 || embd.dims[0] != hidden_size as u64 { + return Err(format!( + "token_embd.weight dims {:?} do not match embedding_length {hidden_size}", + embd.dims + )); + } + let vocab_size = usize::try_from(embd.dims[1]).map_err(|_| "vocab too large")?; + if let Some(tokens) = f.get("tokenizer.ggml.tokens").and_then(GgufValue::as_array) + && tokens.len() > vocab_size + { + return Err(format!( + "tokenizer vocab {} exceeds embedding rows {vocab_size}", + tokens.len() + )); + } + let eos_token_id = f + .get("tokenizer.ggml.eos_token_id") + .and_then(GgufValue::as_u64) + .and_then(|v| u32::try_from(v).ok()) + .map_or(EosIds::None, EosIds::One); + let head_dim = f + .get("qwen2.attention.key_length") + .and_then(GgufValue::as_u64) + .and_then(|v| usize::try_from(v).ok()) + .unwrap_or(hidden_size / num_attention_heads); + let cfg = Self { + vocab_size, + hidden_size, + intermediate_size: gguf_usize(f, "qwen2.feed_forward_length")?, + num_hidden_layers: gguf_usize(f, "qwen2.block_count")?, + num_attention_heads, + num_key_value_heads: gguf_usize(f, "qwen2.attention.head_count_kv")?, + head_dim, + rms_norm_eps: f64::from(gguf_f32(f, "qwen2.attention.layer_norm_rms_epsilon")?), + rope_theta: gguf_f32(f, "qwen2.rope.freq_base")?, + // No separate output.weight tensor means the lm-head is tied. + tie_word_embeddings: f.tensor("output.weight").is_none(), + eos_token_id, + }; + cfg.validate()?; + Ok(cfg) + } + fn validate(&self) -> Result<(), String> { if self.num_key_value_heads == 0 || !self @@ -102,12 +174,16 @@ pub struct DecoderLayer { pub post_attention_layernorm: RmsNorm, } -/// The Qwen2 decoder stack (HF's `model.*` subtree; the lm-head is tied). +/// The Qwen2 decoder stack (HF's `model.*` subtree). The lm-head is tied on +/// the small tiers (0.5B/1.5B safetensors); untied checkpoints — the 7B, and +/// every llama.cpp GGUF (which materializes the head as `output.weight`, at +/// higher precision than the embedding) — carry it explicitly. #[derive(Module, Debug)] pub struct Qwen2 { pub embed_tokens: Embedding, pub layers: Vec>, pub norm: RmsNorm, + pub lm_head: Option>, } /// A weight-loaded Qwen2 plus its config — everything a forward needs. @@ -142,10 +218,16 @@ fn build(cfg: &Qwen2Config, device: &B::Device) -> Qwen2 { post_attention_layernorm: norm(device), }) .collect(); + let lm_head = (!cfg.tie_word_embeddings).then(|| { + LinearConfig::new(cfg.hidden_size, cfg.vocab_size) + .with_bias(false) + .init(device) + }); Qwen2 { embed_tokens: EmbeddingConfig::new(cfg.vocab_size, cfg.hidden_size).init(device), layers, norm: norm(device), + lm_head, } } @@ -185,6 +267,69 @@ pub fn load_from_dir( Ok(LoadedQwen2 { model, config }) } +/// GGUF (llama.cpp) tensor names → the HF checkpoint names the safetensors +/// remap chain already handles. `None` for anything unrecognized — the blob +/// writer turns that into a loud error rather than dropping weights. +fn gguf_tensor_to_hf(name: &str) -> Option { + match name { + "token_embd.weight" => return Some("model.embed_tokens.weight".into()), + "output_norm.weight" => return Some("model.norm.weight".into()), + "output.weight" => return Some("lm_head.weight".into()), + _ => {} + } + let rest = name.strip_prefix("blk.")?; + let (layer, field) = rest.split_once('.')?; + let layer: usize = layer.parse().ok()?; + let mapped = match field { + "attn_norm.weight" => "input_layernorm.weight", + "ffn_norm.weight" => "post_attention_layernorm.weight", + "attn_q.weight" => "self_attn.q_proj.weight", + "attn_q.bias" => "self_attn.q_proj.bias", + "attn_k.weight" => "self_attn.k_proj.weight", + "attn_k.bias" => "self_attn.k_proj.bias", + "attn_v.weight" => "self_attn.v_proj.weight", + "attn_v.bias" => "self_attn.v_proj.bias", + "attn_output.weight" => "self_attn.o_proj.weight", + "ffn_gate.weight" => "mlp.gate_proj.weight", + "ffn_up.weight" => "mlp.up_proj.weight", + "ffn_down.weight" => "mlp.down_proj.weight", + _ => return None, + }; + Some(format!("model.layers.{layer}.{mapped}")) +} + +/// Load a Qwen2 model straight from a **GGUF** file (any dtype the dequant +/// suite covers — Q4_K_M, Q8_0, F16, …): hyperparameters from the GGUF +/// metadata, weights dequantized to f32 and driven through the exact store +/// pipeline (adapters + remaps + checked load) the safetensors path uses. +pub fn load_from_gguf( + path: &Path, + device: &B::Device, +) -> Result, ImportError> { + let parse = |reason: String| ImportError::Parse { + file: path.to_path_buf(), + reason, + }; + 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"); + + let mut model = build::(&config, device); + let target_float = Tensor::::zeros([1], device).dtype(); + let mut store = SafetensorsStore::from_bytes(Some(blob)) + .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) + .allow_partial(true) + .with_key_remapping(r"^model\.", "") + .with_key_remapping(r"(input_layernorm)\.weight$", "$1.gamma") + .with_key_remapping(r"(post_attention_layernorm)\.weight$", "$1.gamma") + .with_key_remapping(r"^norm\.weight$", "norm.gamma"); + load_checked(&mut model, &mut store, path)?; + Ok(LoadedQwen2 { model, config }) +} + impl CausalLm for LoadedQwen2 { type Cache = Vec>; @@ -242,10 +387,20 @@ impl CausalLm for LoadedQwen2 { } let x = self.model.norm.forward(x); - // Tied lm-head: logits = last_hidden @ embed_weight^T. let last = x.narrow(1, t - 1, 1).reshape([1, cfg.hidden_size]); - let w = self.model.embed_tokens.weight.val(); // [vocab, hidden] - last.matmul(w.swap_dims(0, 1)) // [1, vocab] + debug_assert!( + self.model.lm_head.is_some() != cfg.tie_word_embeddings, + "lm_head presence must match the config's tie flag" + ); + match &self.model.lm_head { + // Untied: the checkpoint's own head projection. + Some(head) => head.forward(last), // [1, vocab] + // Tied lm-head: logits = last_hidden @ embed_weight^T. + None => { + let w = self.model.embed_tokens.weight.val(); // [vocab, hidden] + last.matmul(w.swap_dims(0, 1)) // [1, vocab] + } + } } } @@ -342,6 +497,119 @@ mod tests { } } + /// A synthetic in-memory GGUF header shaped like the Qwen2.5-1.5B file. + fn toy_gguf() -> GgufFile { + use crate::gguf::{GgmlType, GgufTensorInfo}; + let meta = |k: &str, v: GgufValue| (k.to_string(), v); + GgufFile { + path: std::path::PathBuf::new(), + version: 3, + metadata: vec![ + meta("general.architecture", GgufValue::Str("qwen2".into())), + meta("qwen2.embedding_length", GgufValue::U32(16)), + meta("qwen2.block_count", GgufValue::U32(2)), + meta("qwen2.feed_forward_length", GgufValue::U32(32)), + meta("qwen2.attention.head_count", GgufValue::U32(4)), + meta("qwen2.attention.head_count_kv", GgufValue::U32(2)), + meta( + "qwen2.attention.layer_norm_rms_epsilon", + GgufValue::F32(1e-6), + ), + meta("qwen2.rope.freq_base", GgufValue::F32(1e4)), + meta("tokenizer.ggml.eos_token_id", GgufValue::U32(2)), + ], + tensors: vec![GgufTensorInfo { + name: "token_embd.weight".into(), + dims: vec![16, 64], // ggml order: [hidden, vocab] + dtype: GgmlType::F32, + offset: 0, + }], + alignment: 32, + data_offset: 0, + } + } + + #[test] + fn config_from_gguf_reads_metadata_and_embedding_dims() { + let cfg = Qwen2Config::from_gguf(&toy_gguf()).expect("parses"); + assert_eq!(cfg.vocab_size, 64); // from token_embd dims[1] + assert_eq!(cfg.hidden_size, 16); + assert_eq!(cfg.num_hidden_layers, 2); + assert_eq!(cfg.head_dim, 4); // derived: hidden / heads + assert!(cfg.eos_token_id.contains(2)); + assert!(cfg.tie_word_embeddings); // no output.weight tensor + } + + #[test] + fn config_from_gguf_fails_loudly_on_missing_keys_and_wrong_arch() { + let mut f = toy_gguf(); + f.metadata.retain(|(k, _)| k != "qwen2.block_count"); + let err = Qwen2Config::from_gguf(&f).unwrap_err(); + assert!(err.contains("qwen2.block_count"), "{err}"); + + let mut f = toy_gguf(); + f.metadata[0].1 = GgufValue::Str("llama".into()); + assert!(Qwen2Config::from_gguf(&f).is_err()); + + // Embedding dims that contradict the metadata are rejected. + let mut f = toy_gguf(); + f.tensors[0].dims = vec![8, 64]; + assert!(Qwen2Config::from_gguf(&f).is_err()); + } + + #[test] + fn untied_toy_config_builds_and_uses_an_lm_head() { + let device = Dev::default(); + let mut cfg = toy_config(); + cfg.tie_word_embeddings = false; + let vocab = cfg.vocab_size; + let loaded = LoadedQwen2:: { + model: build(&cfg, &device), + config: cfg, + }; + assert!(loaded.model.lm_head.is_some()); + let mut cache = loaded.new_cache(); + let logits = loaded.forward(&[1, 2], 0, &mut cache, &device); + assert_eq!(logits.dims(), [1, vocab]); + } + + #[test] + fn config_from_gguf_detects_an_untied_head() { + use crate::gguf::{GgmlType, GgufTensorInfo}; + let mut f = toy_gguf(); + f.tensors.push(GgufTensorInfo { + name: "output.weight".into(), + dims: vec![16, 64], + dtype: GgmlType::F32, + offset: 4096, + }); + let cfg = Qwen2Config::from_gguf(&f).expect("parses"); + assert!(!cfg.tie_word_embeddings); + } + + #[test] + fn gguf_names_map_onto_hf_checkpoint_names() { + assert_eq!( + gguf_tensor_to_hf("token_embd.weight").as_deref(), + Some("model.embed_tokens.weight") + ); + assert_eq!( + gguf_tensor_to_hf("blk.27.attn_q.bias").as_deref(), + Some("model.layers.27.self_attn.q_proj.bias") + ); + assert_eq!( + gguf_tensor_to_hf("blk.0.ffn_down.weight").as_deref(), + Some("model.layers.0.mlp.down_proj.weight") + ); + assert_eq!( + gguf_tensor_to_hf("output_norm.weight").as_deref(), + Some("model.norm.weight") + ); + // Unknown names must map to None (the writer errors loudly). + assert_eq!(gguf_tensor_to_hf("rope_freqs.weight"), None); + assert_eq!(gguf_tensor_to_hf("blk.x.attn_q.weight"), None); + } + #[test] fn greedy_generate_respects_max_tokens_bound() { let device = Dev::default(); diff --git a/crates/mummu/tests/real_gguf.rs b/crates/mummu/tests/real_gguf.rs index 5ed947d..f0b2d33 100644 --- a/crates/mummu/tests/real_gguf.rs +++ b/crates/mummu/tests/real_gguf.rs @@ -8,7 +8,10 @@ use std::path::PathBuf; +use mummu::backend::{Gpu, use_gpu}; use mummu::gguf::{GgmlType, GgufFile, GgufValue}; +use mummu::models::CausalLm; +use mummu::models::qwen2; fn gguf_path() -> Option { let path = PathBuf::from(std::env::var_os("MUMMU_GGUF_PATH")?); @@ -159,3 +162,95 @@ fn real_qwen2_gguf_dequant_matches_the_true_weights() { ); } } + +/// END-TO-END: the Q4_K_M GGUF alone (config + weights from the one file) +/// becomes a running model on the GPU — greedy-decodes a correct answer, and +/// its first-token logits agree with the bf16 safetensors build of the same +/// checkpoint (top-1 identical, high cosine; small drift IS the quantization). +/// The models load sequentially — the second only after the first is dropped +/// — so peak VRAM stays one-model-sized. +#[test] +#[ignore = "needs the local GGUF (MUMMU_GGUF_PATH) + safetensors dir (MUMMU_QWEN2_DIR) + GPU"] +fn real_qwen2_gguf_loads_and_decodes_on_gpu() { + let Some(path) = gguf_path() else { + panic!("set MUMMU_GGUF_PATH to the qwen2.5-1.5b-instruct q4_k_m gguf"); + }; + let dir = std::env::var_os("MUMMU_QWEN2_DIR") + .map(PathBuf::from) + .filter(|d| d.join("tokenizer.json").is_file()) + .expect("set MUMMU_QWEN2_DIR to the same model's safetensors dir (tokenizer.json)"); + assert!(use_gpu(), "this proof wants the real GPU"); + let device = burn::tensor::Device::::default(); + + // Tokenizer from the sibling checkpoint — tokenizer-from-GGUF-metadata + // is the next P3 slice. + let tok = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer"); + let chat = mummu::chat::ChatMl::qwen2(); + let prompt_text = chat.render(&[ + mummu::chat::Turn::system("You are a concise assistant."), + mummu::chat::Turn::user("What is 2+2? Answer in one short sentence."), + ]); + let prompt = tok + .encode(prompt_text, true) + .expect("prompt encodes") + .get_ids() + .to_vec(); + + // Leg 1: the GGUF-loaded model decodes a coherent, correct answer. + let gguf_model = qwen2::load_from_gguf::(&path, &device).expect("gguf load is checked"); + assert_eq!(gguf_model.config.vocab_size, 151_936); + assert_eq!(gguf_model.config.num_hidden_layers, 28); + let ids = gguf_model + .greedy_generate(&prompt, 32, &device) + .expect("decode"); + let text = tok.decode(&ids, true).expect("ids decode"); + eprintln!("[real_gguf] Q4_K_M greedy: {text:?}"); + assert!(text.contains('4'), "expected the answer 4 in: {text:?}"); + + let logits_of = |m: &qwen2::LoadedQwen2| -> Vec { + let mut cache = m.new_cache(); + m.forward(&prompt, 0, &mut cache, &device) + .into_data() + .to_vec::() + .expect("logits read back") + }; + let argmax = |v: &[f32]| -> usize { + let (mut best, mut best_v) = (0usize, f32::NEG_INFINITY); + for (i, &x) in v.iter().enumerate() { + if x > best_v { + (best, best_v) = (i, x); + } + } + best + }; + + // Leg 2: first-token logits vs the bf16 safetensors build — sequential + // loads (drop first) keep peak VRAM at one model. + let gguf_logits = logits_of(&gguf_model); + drop(gguf_model); + let st_model = qwen2::load_from_dir::(&dir, &device).expect("safetensors load"); + let st_logits = logits_of(&st_model); + drop(st_model); + + let top5 = |v: &[f32]| -> Vec { + let mut idx: Vec = (0..v.len()).collect(); + idx.sort_by(|&a, &b| v[b].total_cmp(&v[a])); + idx.truncate(5); + idx + }; + let cos = cosine(&gguf_logits, &st_logits); + let (g_top, s_top) = (argmax(&gguf_logits), argmax(&st_logits)); + let (g5, s5) = (top5(&gguf_logits), top5(&st_logits)); + let overlap = g5.iter().filter(|id| s5.contains(id)).count(); + eprintln!( + "[real_gguf] first-token logits: cosine {cos:.5} vs bf16 · top-1 {g_top} vs {s_top} · top-5 overlap {overlap}/5 ({g5:?} vs {s5:?})" + ); + assert_eq!(g_top, s_top, "Q4_K_M must agree with bf16 on the top token"); + assert!(overlap >= 4, "top-5 sets diverge: {g5:?} vs {s5:?}"); + // Measured 0.977 on this checkpoint: 28 layers of Q4_K_M drift (plus the + // bf16 reference's own rounding). A layout/scale decode bug reads ≈ 0. + assert!( + cos > 0.95, + "logit cosine {cos} — quantization noise should be small, layout bugs are not" + ); +} From f6aee6f6def7f5637c5c86f8fb52ab1af532efed Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Mon, 13 Jul 2026 07:56:28 -0500 Subject: [PATCH 3/6] =?UTF-8?q?feat(p3):=20tokenizer-from-GGUF=20metadata?= =?UTF-8?q?=20=E2=80=94=20one=20.gguf=20file=20is=20the=20whole=20model?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - mummu::tokenizer::tokenizer_from_gguf rebuilds the HF tokenizers pipeline from tokenizer.ggml.* metadata: NFC -> Split(per-family pre-tokenizer regex, llama.cpp-style registry keyed by tokenizer.ggml.pre; unknown ids error loudly) -> ByteLevel -> BPE - token id = tokens-array index; CONTROL(3) -> special added tokens, USER_DEFINED(4) -> plain added tokens, UNUSED(5) [PADn] skipped; every added-token id is verified post-build (drift = loud error) - GgufValue::as_i64 for the signed token_type array - the end-to-end GPU test now sources tokenizer + config + weights from the single .gguf file Proof: byte-identical ids vs the checkpoint's tokenizer.json across an 8-prompt battery (ChatML with specials, unicode/CJK/emoji, whitespace runs, contractions, numbers, empty) and identical decodes; end-to-end Q4_K_M decode re-passed ('2+2 equals 4.'). 135 unit tests green, clippy clean. Co-Authored-By: Claude Opus 4.8 --- README.md | 4 +- ROADMAP.md | 10 +- crates/mummu/src/gguf.rs | 13 ++ crates/mummu/src/lib.rs | 1 + crates/mummu/src/tokenizer.rs | 227 ++++++++++++++++++++++++++++++++ crates/mummu/tests/real_gguf.rs | 61 ++++++++- 6 files changed, 311 insertions(+), 5 deletions(-) create mode 100644 crates/mummu/src/tokenizer.rs diff --git a/README.md b/README.md index d4b7d9b..443a5e8 100644 --- a/README.md +++ b/README.md @@ -72,7 +72,9 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys the model's true weights (`tests/real_gguf.rs`): F32 norms **bit-exact** vs the bf16 safetensors of the same checkpoint, Q4_K rows at cosine 0.9975, and the real Qwen2.5-1.5B **Q4_K_M file greedy-decodes "2+2 equals 4." on the GPU** with first-token top-1 identical to the bf16 build (logit cosine 0.977). - Next: tokenizer-from-GGUF metadata and the LFM2 GGUF map (tracked in P3). + The tokenizer comes from the GGUF too (`tokenizer_from_gguf`: NFC → per-family pre-regex → ByteLevel + → BPE, byte-identical ids vs the checkpoint's `tokenizer.json` on an 8-prompt battery) — **one .gguf + file is the whole model**. Next: the LFM2 GGUF map (tracked in P3). - **Hub downloads** — streaming HuggingFace fetches into the model cache: resumable (`.part` + HTTP Range, proven byte-identical after an interrupted transfer), length-verified, shard-index aware, with a per-chunk progress callback; verified end-to-end by downloading all-MiniLM and embedding with it. diff --git a/ROADMAP.md b/ROADMAP.md index 9d7cf54..19b8014 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -203,9 +203,17 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari logit cosine 0.977 (28 layers of Q4_K drift; a layout bug reads ≈ 0). Parity gate re-passed byte-identically after the lm_head change (max |Δlogit| 2.670e-5, Ollama greedy exact); f16 + budget gates green (GPU 107.4 ms / 13.6 tok/s, CPU 14.8 tok/s).* -- [ ] **Tokenizer-from-GGUF metadata** — build the HF `tokenizers` pipeline from `tokenizer.ggml.*` +- [x] **Tokenizer-from-GGUF metadata** — build the HF `tokenizers` pipeline from `tokenizer.ggml.*` (tokens, merges, token types, BPE pre-tokenizer regex) so a GGUF needs no sibling `tokenizer.json`; byte-verify token ids against the HF tokenizer of the same checkpoint. + *(2026-07-13, same run) Shipped (`mummu::tokenizer::tokenizer_from_gguf`): NFC → Split(the + per-family `tokenizer.ggml.pre` regex, llama.cpp-style registry — unknown ids are loud errors) → + ByteLevel → BPE; token id = array index; CONTROL(3)/USER_DEFINED(4) types become special/plain + added tokens with post-build id verification, UNUSED(5) `[PADn]` entries skipped. REAL-FILE + proof: **byte-identical ids vs the checkpoint's `tokenizer.json` on an 8-prompt battery** + (ChatML + specials, unicode/CJK/emoji, whitespace runs, contractions, empty) and identical + decodes; the end-to-end GPU test now runs tokenizer + config + weights from the ONE .gguf file. + Only `gpt2`-model/`qwen2`-pre is registered so far — new families add a regex entry.* - [ ] **LFM2 GGUF name map** — extend `load_from_gguf` to the LFM2/LFM2.5 hybrid (llama.cpp `lfm2` arch: `shortconv.*` tensor names, `lfm2.*` metadata keys) once a same-weights GGUF is validated. - [ ] **Quantized-reference parity leg for GGUF loads** — the end-to-end test compares against the bf16 diff --git a/crates/mummu/src/gguf.rs b/crates/mummu/src/gguf.rs index 2aa9680..b6acb70 100644 --- a/crates/mummu/src/gguf.rs +++ b/crates/mummu/src/gguf.rs @@ -115,6 +115,19 @@ impl GgufValue { } } + /// The value widened to i64, if it is any integer (signed or unsigned) + /// that fits. + #[must_use] + pub fn as_i64(&self) -> Option { + match *self { + Self::I8(v) => Some(i64::from(v)), + Self::I16(v) => Some(i64::from(v)), + Self::I32(v) => Some(i64::from(v)), + Self::I64(v) => Some(v), + _ => self.as_u64().and_then(|v| i64::try_from(v).ok()), + } + } + /// The value as f32, if it is one. #[must_use] pub fn as_f32(&self) -> Option { diff --git a/crates/mummu/src/lib.rs b/crates/mummu/src/lib.rs index 8239a95..0298b20 100644 --- a/crates/mummu/src/lib.rs +++ b/crates/mummu/src/lib.rs @@ -20,3 +20,4 @@ pub mod manage; pub mod models; pub mod nn; pub mod registry; +pub mod tokenizer; diff --git a/crates/mummu/src/tokenizer.rs b/crates/mummu/src/tokenizer.rs new file mode 100644 index 0000000..b720fc0 --- /dev/null +++ b/crates/mummu/src/tokenizer.rs @@ -0,0 +1,227 @@ +//! Tokenizer from GGUF metadata — the piece that makes a GGUF file fully +//! self-contained (no sibling `tokenizer.json` needed). +//! +//! llama.cpp stores the tokenizer as `tokenizer.ggml.*` metadata: the vocab +//! (`tokens`, index = token id), per-token types, BPE `merges`, and a `pre` +//! identifier naming the pre-tokenizer regex (the ecosystem hardcodes the +//! regex per model family, exactly as llama.cpp's `llama_vocab` does). This +//! module rebuilds the equivalent HF `tokenizers` pipeline: +//! NFC → Split(pre regex) → ByteLevel → BPE, with control/user-defined +//! tokens re-added as special/non-special added tokens. +//! +//! Faithfulness is verified against the same checkpoint's `tokenizer.json` +//! (byte-identical ids over a battery of prompts) in `tests/real_gguf.rs`. + +use tokenizers::models::bpe::{BPE, Merges, Vocab}; +use tokenizers::normalizers::unicode::NFC; +use tokenizers::pre_tokenizers::byte_level::ByteLevel; +use tokenizers::pre_tokenizers::sequence::Sequence; +use tokenizers::pre_tokenizers::split::{Split, SplitPattern}; +use tokenizers::{AddedToken, SplitDelimiterBehavior, Tokenizer}; + +use crate::gguf::{GgufFile, GgufValue}; + +/// llama.cpp token types (`llama_token_type`). +const TOKEN_TYPE_NORMAL: i64 = 1; +const TOKEN_TYPE_CONTROL: i64 = 3; +const TOKEN_TYPE_USER_DEFINED: i64 = 4; +const TOKEN_TYPE_UNUSED: i64 = 5; + +/// The GPT-2-style byte-level BPE pre-tokenizer regexes, keyed by +/// `tokenizer.ggml.pre` — the same registry llama.cpp keeps in its vocab +/// loader. Only families we actually run are listed; unknown ids are a loud +/// error (a wrong regex silently produces wrong token ids). +fn pre_tokenizer_regex(pre: &str) -> Option<&'static str> { + match pre { + // Qwen2/2.5 (matches the checkpoint's tokenizer.json byte for byte). + "qwen2" => Some( + r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+", + ), + _ => None, + } +} + +/// A required `tokenizer.ggml.*` metadata array. +fn required_array<'f>(f: &'f GgufFile, key: &str) -> Result<&'f [GgufValue], String> { + f.get(key) + .and_then(GgufValue::as_array) + .ok_or_else(|| format!("missing or non-array GGUF metadata '{key}'")) +} + +/// Build an HF [`Tokenizer`] from a GGUF file's `tokenizer.ggml.*` metadata. +/// +/// Supports the `gpt2` (byte-level BPE) tokenizer model with a known `pre` +/// regex. Token ids are the `tokens` array indexes; control (type 3) tokens +/// become special added tokens, user-defined (type 4) become non-special +/// added tokens, unused (type 5) padding entries are skipped. Every added +/// token's id is verified after construction — a drifted id would silently +/// corrupt every prompt, so it fails loudly instead. +pub fn tokenizer_from_gguf(f: &GgufFile) -> Result { + let model = f + .get("tokenizer.ggml.model") + .and_then(GgufValue::as_str) + .ok_or("missing GGUF metadata 'tokenizer.ggml.model'")?; + if model != "gpt2" { + return Err(format!( + "tokenizer model '{model}' is not supported yet (only gpt2 byte-level BPE)" + )); + } + let pre = f + .get("tokenizer.ggml.pre") + .and_then(GgufValue::as_str) + .ok_or("missing GGUF metadata 'tokenizer.ggml.pre'")?; + let Some(regex) = pre_tokenizer_regex(pre) else { + return Err(format!("unknown pre-tokenizer id '{pre}'")); + }; + + let tokens = required_array(f, "tokenizer.ggml.tokens")?; + let types = required_array(f, "tokenizer.ggml.token_type")?; + if tokens.len() != types.len() { + return Err(format!( + "tokens ({}) and token_type ({}) lengths differ", + tokens.len(), + types.len() + )); + } + + // The BPE vocab: every NORMAL token, id = array index. + let mut vocab = Vocab::default(); + // (index, content, is_special) for everything re-added post-BPE. + let mut added: Vec<(usize, String, bool)> = Vec::new(); + for (index, (token, ty)) in tokens.iter().zip(types).enumerate() { + let text = token + .as_str() + .ok_or_else(|| format!("token {index} is not a string"))?; + let ty = ty + .as_i64() + .ok_or_else(|| format!("token_type {index} is not an integer"))?; + #[allow(clippy::cast_possible_truncation)] // bounded by MAX_ARRAY_LEN + match ty { + TOKEN_TYPE_NORMAL => { + vocab.insert(text.to_string(), index as u32); + } + TOKEN_TYPE_CONTROL => added.push((index, text.to_string(), true)), + TOKEN_TYPE_USER_DEFINED => added.push((index, text.to_string(), false)), + TOKEN_TYPE_UNUSED => {} // vocab-padding entries ([PADn]) + other => return Err(format!("token {index} has unsupported type {other}")), + } + } + assert!(!vocab.is_empty(), "a tokenizer must have normal tokens"); + + let mut merges = Merges::with_capacity(required_array(f, "tokenizer.ggml.merges")?.len()); + for (i, m) in required_array(f, "tokenizer.ggml.merges")? + .iter() + .enumerate() + { + let m = m + .as_str() + .ok_or_else(|| format!("merge {i} is not a string"))?; + let (a, b) = m + .split_once(' ') + .ok_or_else(|| format!("merge {i} ('{m}') is not 'left right'"))?; + merges.push((a.to_string(), b.to_string())); + } + + let bpe = BPE::builder() + .vocab_and_merges(vocab, merges) + .build() + .map_err(|e| format!("BPE build: {e}"))?; + let mut tok = Tokenizer::new(bpe); + tok.with_normalizer(Some(NFC)); + let split = Split::new( + SplitPattern::Regex(regex.to_string()), + SplitDelimiterBehavior::Isolated, + false, + ) + .map_err(|e| format!("pre-tokenizer regex: {e}"))?; + // ByteLevel exactly as the HF checkpoints configure it: no prefix space, + // no offset trimming, regex handled by the Split stage above. + let byte_level = ByteLevel::new(false, false, false); + tok.with_pre_tokenizer(Some(Sequence::new(vec![split.into(), byte_level.into()]))); + tok.with_decoder(Some(byte_level)); + tok.with_post_processor(Some(byte_level)); + + // Re-add control/user-defined tokens in id order, then verify every id + // landed where the GGUF says it lives. + for (_, text, special) in &added { + let t = AddedToken::from(text.clone(), *special); + if *special { + tok.add_special_tokens(&[t]); + } else { + tok.add_tokens(&[t]); + } + } + for (index, text, _) in &added { + let got = tok.token_to_id(text); + #[allow(clippy::cast_possible_truncation)] // bounded by MAX_ARRAY_LEN + if got != Some(*index as u32) { + return Err(format!( + "added token '{text}' resolved to id {got:?}, GGUF says {index} — \ + non-contiguous added-token ids are not supported" + )); + } + } + Ok(tok) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::gguf::GgufValue; + + /// A minimal synthetic GGUF header carrying a 4-token BPE tokenizer. + fn toy_gguf() -> GgufFile { + let strs = |v: &[&str]| { + GgufValue::Array(v.iter().map(|s| GgufValue::Str((*s).to_string())).collect()) + }; + let ints = |v: &[i32]| GgufValue::Array(v.iter().map(|&i| GgufValue::I32(i)).collect()); + GgufFile { + path: std::path::PathBuf::new(), + version: 3, + metadata: vec![ + ("tokenizer.ggml.model".into(), GgufValue::Str("gpt2".into())), + ("tokenizer.ggml.pre".into(), GgufValue::Str("qwen2".into())), + ( + "tokenizer.ggml.tokens".into(), + strs(&["a", "b", "ab", "<|stop|>", "[PAD4]"]), + ), + ("tokenizer.ggml.token_type".into(), ints(&[1, 1, 1, 3, 5])), + ("tokenizer.ggml.merges".into(), strs(&["a b"])), + ], + tensors: vec![], + alignment: 32, + data_offset: 0, + } + } + + #[test] + fn toy_gguf_tokenizer_encodes_merges_and_specials() { + let tok = tokenizer_from_gguf(&toy_gguf()).expect("builds"); + let ids = tok.encode("ab", true).expect("encodes"); + assert_eq!(ids.get_ids(), &[2], "the merge applies"); + assert_eq!(tok.token_to_id("<|stop|>"), Some(3), "control token id"); + let ids = tok.encode("<|stop|>ab", true).expect("encodes"); + assert_eq!(ids.get_ids(), &[3, 2], "special token is never split"); + // Unused padding entries are skipped, not part of the vocab. + assert_eq!(tok.token_to_id("[PAD4]"), None); + } + + #[test] + fn unknown_model_pre_and_bad_merges_fail_loudly() { + let mut f = toy_gguf(); + f.metadata[0].1 = GgufValue::Str("sentencepiece".into()); + assert!(tokenizer_from_gguf(&f).unwrap_err().contains("gpt2")); + + let mut f = toy_gguf(); + f.metadata[1].1 = GgufValue::Str("llama-bpe".into()); + assert!( + tokenizer_from_gguf(&f) + .unwrap_err() + .contains("pre-tokenizer") + ); + + let mut f = toy_gguf(); + f.metadata[4].1 = GgufValue::Array(vec![GgufValue::Str("nospace".into())]); + assert!(tokenizer_from_gguf(&f).unwrap_err().contains("merge")); + } +} diff --git a/crates/mummu/tests/real_gguf.rs b/crates/mummu/tests/real_gguf.rs index f0b2d33..a1832e1 100644 --- a/crates/mummu/tests/real_gguf.rs +++ b/crates/mummu/tests/real_gguf.rs @@ -163,6 +163,60 @@ fn real_qwen2_gguf_dequant_matches_the_true_weights() { } } +/// The tokenizer rebuilt from GGUF metadata must be **byte-identical** to the +/// checkpoint's own `tokenizer.json` — same ids for every prompt shape we +/// throw at it (ChatML with specials, unicode, numbers, whitespace runs), and +/// the same decoded text back. +#[test] +#[ignore = "needs the local GGUF (MUMMU_GGUF_PATH) + safetensors dir (MUMMU_QWEN2_DIR)"] +fn real_qwen2_gguf_tokenizer_matches_the_hf_tokenizer() { + let Some(path) = gguf_path() else { + panic!("set MUMMU_GGUF_PATH to the qwen2.5-1.5b-instruct q4_k_m gguf"); + }; + let dir = std::env::var_os("MUMMU_QWEN2_DIR") + .map(PathBuf::from) + .filter(|d| d.join("tokenizer.json").is_file()) + .expect("set MUMMU_QWEN2_DIR to the same model's safetensors dir (tokenizer.json)"); + + let f = GgufFile::open(&path).expect("header parses"); + let ours = mummu::tokenizer::tokenizer_from_gguf(&f).expect("tokenizer builds from metadata"); + let reference = + tokenizers::Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer.json"); + + let chat = mummu::chat::ChatMl::qwen2(); + let battery = [ + chat.render(&[ + mummu::chat::Turn::system("You are a concise assistant."), + mummu::chat::Turn::user("What is 2+2? Answer in one short sentence."), + ]), + "The quick brown fox jumps over the lazy dog.".into(), + "héllo wörld — 世界 · Ω ≠ ω · 🦀🔥".into(), + " leading spaces, trailing \n\nnewlines\r\nand\ttabs ".into(), + "1234567890 3.14159 1e-6 0x7F".into(), + "<|im_start|>user\nplain<|im_end|><|endoftext|>".into(), + "don't they're we've I'll it's CAN'T".into(), + String::new(), + ]; + for text in &battery { + let a = ours.encode(text.as_str(), true).expect("ours encodes"); + let b = reference.encode(text.as_str(), true).expect("ref encodes"); + assert_eq!( + a.get_ids(), + b.get_ids(), + "token ids diverge on {text:?}: {:?} vs {:?}", + a.get_tokens(), + b.get_tokens() + ); + let da = ours.decode(a.get_ids(), false).expect("ours decodes"); + let db = reference.decode(b.get_ids(), false).expect("ref decodes"); + assert_eq!(da, db, "decoded text diverges on {text:?}"); + } + eprintln!( + "[real_gguf] tokenizer-from-GGUF: {} prompts byte-identical to tokenizer.json", + battery.len() + ); +} + /// END-TO-END: the Q4_K_M GGUF alone (config + weights from the one file) /// becomes a running model on the GPU — greedy-decodes a correct answer, and /// its first-token logits agree with the bf16 safetensors build of the same @@ -182,9 +236,10 @@ fn real_qwen2_gguf_loads_and_decodes_on_gpu() { assert!(use_gpu(), "this proof wants the real GPU"); let device = burn::tensor::Device::::default(); - // Tokenizer from the sibling checkpoint — tokenizer-from-GGUF-metadata - // is the next P3 slice. - let tok = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer"); + // The tokenizer comes from the GGUF itself — the whole model is ONE file + // (byte-verified against tokenizer.json by the sibling test). + let header = GgufFile::open(&path).expect("header parses"); + let tok = mummu::tokenizer::tokenizer_from_gguf(&header).expect("tokenizer from metadata"); let chat = mummu::chat::ChatMl::qwen2(); let prompt_text = chat.render(&[ mummu::chat::Turn::system("You are a concise assistant."), From d194ef407987b126f98b36c322dfc2ce646cca31 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Mon, 13 Jul 2026 07:58:44 -0500 Subject: [PATCH 4/6] =?UTF-8?q?docs:=20research=20fold=20=E2=80=94=20CubeC?= =?UTF-8?q?L=20quant=20primitives=20(P9),=20GPTQ/AWQ=3Dcompressed-tensors?= =?UTF-8?q?=20convention=20(P3),=20llama.cpp=20LFM2.5=20parser=20fix=20ups?= =?UTF-8?q?tream=20(P7)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 4.8 --- ROADMAP.md | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/ROADMAP.md b/ROADMAP.md index 19b8014..d54a691 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -222,6 +222,12 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari reference item's caveats) to assert our dequant matches ggml's compute path token-for-token. - [ ] **GPTQ / AWQ** (HF safetensors) — import the calibration-quantized int4/int8 layouts most "quantized on the Hub" models ship as (a `.safetensors` payload + a quant config), dequant or keep-quant into Burn. + *(2026-07-13 research)* Both are quantization *algorithms*, not formats — the artifact is ordinary + safetensors shards (packed int4 `qweight`/`qzeros`/`scales`, group size 128 is the de-facto + standard) plus `quantization_config` in `config.json`; vLLM's **compressed-tensors** is the + emerging unified on-disk convention to target (one reader covers GPTQ/AWQ/INT8/FP8 exports) — + https://github.com/vllm-project/compressed-tensors · + https://www.digitalapplied.com/blog/gguf-vs-awq-vs-gptq-vs-mlx-llm-quantization-formats-2026 - [ ] **ONNX** (optional) — `burn-import` ONNX→Burn for models distributed as ONNX graphs. - [ ] **Dtype handling** — a `CastFloatAdapter` (bf16→f32/f16); quantized→dequant on import; keep-quantized handed to P9. @@ -376,6 +382,11 @@ that fits the model AND uses every device to the fullest. drive `llama-server` in RAW completion mode (`/completion`, no chat template) and render prompts with our byte-verified `ChatMl::lfm2()` — never through llama.cpp's template stack — https://github.com/ggml-org/llama.cpp/issues/23838 + *(2026-07-13 research)* That parser bug is now FIXED upstream (PR #24178 merged ~June 2026 — + `is_lfm2_template()` detected only the gen-1 `<|tool_list_start|>` tags); raw completion mode + remains the right choice for the logits leg regardless (no template stack in the loop). Also: + LiquidAI officially publishes LFM2.5-1.2B GGUFs (incl. F16) — the same-weights reference artifact + this leg needs — https://github.com/ggml-org/llama.cpp/issues/23838 - [ ] Wire the perf suite (above) into the parity harness so a correctness *or* budget regression fails CI. ### P8 — Model management API @@ -394,7 +405,14 @@ The VRAM lever the P6 planner pulls to make the largest useful model fit the use + CPU paths; quantize on import or on the fly, keyed to the fit target from P6. - [ ] **Import pre-quantized** — run **GGUF K-quants** (Q2_K–Q8_0, per-layer precision) and **GPTQ / AWQ** int4 layouts directly (dequant to the compute dtype, or a keep-quantized matmul where the kernel exists), so a - model already quantized on the Hub loads as-is. + model already quantized on the Hub loads as-is. *(2026-07-13) The dequant-to-f32 leg of the GGUF + path shipped in P3 (`load_from_gguf` — every storage dtype); what remains here is **keep-quantized + in VRAM**, which is the actual fit lever (Q4_K_M currently dequants to the same f32 footprint).* +- [ ] Evaluate **CubeCL's quantization primitives** for the keep-quantized matmul: recent CubeCL ships + block-scaled MMA, global quantization for matmul, quantized tensor views, and FP4/FP2 formats — + the kernel substrate a Q4-weights × f16-activations decode path would ride (vs hand-writing a + dequant-fused kernel); gate any adoption on the parity harness + `bench/BASELINE.md` — + https://github.com/tracel-ai/cubecl/releases · https://burn.dev/blog/release-0.21.0/ - [ ] **Auto-quantize-to-fit** — the planner picks the *highest* precision that fits the detected VRAM (f16 → int8 → int4), reports the quality/size trade, and never silently ships a worse tier than asked. From 2876af9b3389e0346933be6864801a081bc6b5ce Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Mon, 13 Jul 2026 08:26:52 -0500 Subject: [PATCH 5/6] =?UTF-8?q?feat(p3):=20LFM2.5=20GGUF=20import=20?= =?UTF-8?q?=E2=80=94=20the=20hybrid=20loads=20from=20one=20file=20too?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Lfm2Config::from_gguf: hyperparameters from lfm2.* metadata; layer kinds derived from llama.cpp's per-layer head_count_kv array (0 = shortconv, nonzero = attention; i32 in real files, all nonzero must agree); feed_forward_length arrives pre-adjusted - lfm2 tensor-name map incl. per-head q/k norms and shortconv.*; the depthwise conv kernel is the one shape special-case: new GgufMap::Reshape un-squeezes ggml's [K, C] back to the checkpoint's [C, 1, K] (same bytes; element count validated) - tokenizer registry gains the lfm2 pre: digits-in-<=3-groups regex, no NFC, BOS post-processor from tokenizer.ggml.add_bos_token; added tokens at LOW ids (LFM2 puts 500+ specials at 0..) work by seeding the BPE vocab and letting add_tokens reuse the model id Proof (official LiquidAI Q4_K_M vs the same checkpoint's local bf16 safetensors): 5 F32 tensors incl. both conv kernels bit-exact; tokenizer byte-identical over a 6-prompt x 2-mode battery; REAL-GPU end-to-end — the one file greedy-decodes '2 + 2 equals 4.', top-1 identical to the bf16 build, logit cosine 0.9914. 138 unit tests; qwen2 GGUF suite, parity gate (2.670e-5, Ollama byte-identical), and budget gates (GPU 106.7 ms / 13.5 tok/s, CPU 13.5 tok/s) all re-passed. Co-Authored-By: Claude Opus 4.8 --- README.md | 9 +- ROADMAP.md | 15 +- crates/mummu/src/gguf.rs | 56 ++++-- crates/mummu/src/models/lfm2.rs | 285 +++++++++++++++++++++++++++++++ crates/mummu/src/models/qwen2.rs | 20 ++- crates/mummu/src/tokenizer.rs | 83 +++++++-- crates/mummu/tests/real_gguf.rs | 193 +++++++++++++++++++++ 7 files changed, 622 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 443a5e8..a52e7f4 100644 --- a/README.md +++ b/README.md @@ -72,9 +72,12 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys the model's true weights (`tests/real_gguf.rs`): F32 norms **bit-exact** vs the bf16 safetensors of the same checkpoint, Q4_K rows at cosine 0.9975, and the real Qwen2.5-1.5B **Q4_K_M file greedy-decodes "2+2 equals 4." on the GPU** with first-token top-1 identical to the bf16 build (logit cosine 0.977). - The tokenizer comes from the GGUF too (`tokenizer_from_gguf`: NFC → per-family pre-regex → ByteLevel - → BPE, byte-identical ids vs the checkpoint's `tokenizer.json` on an 8-prompt battery) — **one .gguf - file is the whole model**. Next: the LFM2 GGUF map (tracked in P3). + The tokenizer comes from the GGUF too (`tokenizer_from_gguf`: per-family pre-regex → ByteLevel → BPE, + byte-identical ids vs the checkpoint's `tokenizer.json`) — **one .gguf file is the whole model**. + Works for the **LFM2.5 hybrid** as well (`lfm2::load_from_gguf`: layer kinds from the per-layer + kv-head array, conv kernels un-squeezed bit-exactly): the official LiquidAI Q4_K_M greedy-decodes + "2 + 2 equals 4." with top-1 identical to bf16 (logit cosine 0.991). Next: keep-quantized VRAM + (tracked in P9). - **Hub downloads** — streaming HuggingFace fetches into the model cache: resumable (`.part` + HTTP Range, proven byte-identical after an interrupted transfer), length-verified, shard-index aware, with a per-chunk progress callback; verified end-to-end by downloading all-MiniLM and embedding with it. diff --git a/ROADMAP.md b/ROADMAP.md index d54a691..4b7fa0e 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -214,8 +214,21 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari (ChatML + specials, unicode/CJK/emoji, whitespace runs, contractions, empty) and identical decodes; the end-to-end GPU test now runs tokenizer + config + weights from the ONE .gguf file. Only `gpt2`-model/`qwen2`-pre is registered so far — new families add a regex entry.* -- [ ] **LFM2 GGUF name map** — extend `load_from_gguf` to the LFM2/LFM2.5 hybrid (llama.cpp `lfm2` +- [x] **LFM2 GGUF name map** — extend `load_from_gguf` to the LFM2/LFM2.5 hybrid (llama.cpp `lfm2` arch: `shortconv.*` tensor names, `lfm2.*` metadata keys) once a same-weights GGUF is validated. + *(2026-07-13, same run) Shipped (`lfm2::load_from_gguf`): `Lfm2Config::from_gguf` derives + `layer_types` from llama.cpp's per-layer `head_count_kv` array (0 = conv, nonzero = attention; + i32 in real files), `feed_forward_length` arrives pre-adjusted; the name map covers the hybrid's + per-head q/k norms + `shortconv.{conv,in_proj,out_proj}`, with the depthwise conv kernel as the + one shape special-case (`GgufMap::Reshape` — llama.cpp squeezes `[C,1,K]` to ggml `[K,C]`, same + bytes). Tokenizer registry gained the `lfm2` pre (digits-≤3 regex, no NFC, BOS post-processor + from `add_bos_token`; low-id added tokens ride BPE-vocab id reuse). REAL-FILE proof against the + official LiquidAI Q4_K_M (697 MB, downloaded this run) vs the local bf16 safetensors of the same + checkpoint: 5 F32 tensors incl. both conv kernels **bit-exact**; tokenizer byte-identical on a + 6-prompt × 2-mode battery; REAL-GPU end-to-end — the one file greedy-decodes "2 + 2 equals 4.", + first-token top-1 identical to the bf16 build, logit cosine **0.9914**. (Still not the P2/P7 + strict parity gate — that needs the llama.cpp same-quant reference leg.) 138 unit tests; parity + (2.670e-5, byte-identical) + budget gates re-passed.* - [ ] **Quantized-reference parity leg for GGUF loads** — the end-to-end test compares against the bf16 build (quantization drift bounded, not exact); a strict leg needs llama.cpp itself running the SAME quantized file (`llama-server` raw `/completion`, `n_probs` logprobs — see the P7 LFM2.5 diff --git a/crates/mummu/src/gguf.rs b/crates/mummu/src/gguf.rs index b6acb70..8397a8d 100644 --- a/crates/mummu/src/gguf.rs +++ b/crates/mummu/src/gguf.rs @@ -259,6 +259,19 @@ impl GgufTensorInfo { } } +/// How one GGUF tensor lands in the safetensors blob +/// ([`GgufFile::dequant_to_safetensors`]). +#[derive(Debug, Clone)] +pub enum GgufMap { + /// Target name; shape = the GGUF dims reversed (the row-major twin of + /// the ggml layout — right for everything but squeezed kernels). + Rename(String), + /// Target name + explicit row-major shape (same element count, same + /// bytes — e.g. un-squeezing a depthwise conv kernel back to + /// `[channels, 1, k]`). + Reshape(String, Vec), +} + /// A parsed GGUF header: typed metadata + the located tensor table. #[derive(Debug)] pub struct GgufFile { @@ -341,18 +354,21 @@ impl GgufFile { /// Dequantize every tensor to f32 and serialize the result as an /// in-memory **safetensors** file — the bridge onto the exact store /// pipeline (adapters, key remaps, checked load) the safetensors path - /// already trusts. `rename` maps each GGUF tensor name to the name the - /// blob should carry (HF-checkpoint naming, so per-model remap tables - /// apply unchanged); an unmapped tensor is a loud error, never a skip — - /// a name this crate doesn't recognize means weights would silently - /// vanish from the model. + /// already trusts. `map` maps each GGUF tensor to the name (and + /// optionally an explicit shape) the blob should carry (HF-checkpoint + /// naming, so per-model remap tables apply unchanged); an unmapped + /// tensor is a loud error, never a skip — a name this crate doesn't + /// recognize means weights would silently vanish from the model. /// - /// Shapes are the GGUF dims **reversed**: ggml orders dims + /// Default shapes are the GGUF dims **reversed**: ggml orders dims /// fastest-varying-first, so the raw payload bytes are exactly the /// row-major layout of the reversed shape — same bytes, HF convention. + /// [`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). pub fn dequant_to_safetensors( &self, - rename: &dyn Fn(&str) -> Option, + 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 { @@ -374,8 +390,25 @@ impl GgufFile { let mut data: Vec = Vec::with_capacity(usize::try_from(total_f32_bytes).expect("bounded above")); for (index, info) in self.tensors.iter().enumerate() { - let Some(name) = rename(&info.name) else { - return Err(bad(index, format!("unmapped tensor name '{}'", info.name))); + let (name, shape) = match map(info) { + Some(GgufMap::Rename(name)) => { + (name, info.dims.iter().rev().copied().collect::>()) + } + Some(GgufMap::Reshape(name, shape)) => { + if shape.iter().product::() != info.element_count() { + return Err(bad( + index, + format!( + "reshape of '{}' to {shape:?} changes the element count", + info.name + ), + )); + } + (name, shape) + } + None => { + return Err(bad(index, format!("unmapped tensor name '{}'", info.name))); + } }; if !names.insert(name.clone()) { return Err(bad(index, format!("rename collision on '{name}'"))); @@ -383,7 +416,6 @@ 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 shape: Vec = info.dims.iter().rev().copied().collect(); let json_name = serde_json::to_string(&name).expect("string serializes"); if index > 0 { header.push(','); @@ -1349,7 +1381,7 @@ mod tests { with_gguf_bytes(&bytes, |f| { let f = f.expect("parses"); let blob = f - .dequant_to_safetensors(&|n| Some(format!("model.{n}"))) + .dequant_to_safetensors(&|i| Some(GgufMap::Rename(format!("model.{}", i.name)))) .expect("serializes"); let header_len = u64::from_le_bytes(blob[0..8].try_into().unwrap()); let json: serde_json::Value = @@ -1383,7 +1415,7 @@ mod tests { with_gguf_bytes(&bytes, |f| { let f = f.expect("parses"); assert!(matches!( - f.dequant_to_safetensors(&|_| Some("same".into())), + f.dequant_to_safetensors(&|_| Some(GgufMap::Rename("same".into()))), Err(GgufError::BadTensor { .. }) )); }); diff --git a/crates/mummu/src/models/lfm2.rs b/crates/mummu/src/models/lfm2.rs index 4675fd9..12ef98d 100644 --- a/crates/mummu/src/models/lfm2.rs +++ b/crates/mummu/src/models/lfm2.rs @@ -15,6 +15,7 @@ use burn::nn::{Embedding, EmbeddingConfig, RmsNorm, RmsNormConfig}; 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::models::CausalLm; use crate::models::qwen2::EosIds; @@ -79,6 +80,101 @@ impl Lfm2Config { Ok(cfg) } + /// Hyperparameters from a GGUF header's `lfm2.*` metadata. Layer kinds + /// come from llama.cpp's per-layer `attention.head_count_kv` array: + /// `0` marks a shortconv layer, nonzero an attention layer. + /// `feed_forward_length` in a GGUF is the already-adjusted SwiGLU dim, + /// so auto-adjust is off. + pub fn from_gguf(f: &GgufFile) -> Result { + let arch = f.architecture().unwrap_or(""); + if arch != "lfm2" { + return Err(format!("GGUF architecture '{arch}' is not lfm2")); + } + let meta_usize = |key: &str| -> Result { + f.get(key) + .and_then(GgufValue::as_u64) + .and_then(|v| usize::try_from(v).ok()) + .ok_or_else(|| format!("missing or non-integer GGUF metadata '{key}'")) + }; + let hidden_size = meta_usize("lfm2.embedding_length")?; + let num_hidden_layers = meta_usize("lfm2.block_count")?; + + // Per-layer kv-head counts: 0 = conv, nonzero = attention (all + // nonzero entries must agree — one KV geometry per model). + let kv_per_layer = f + .get("lfm2.attention.head_count_kv") + .and_then(GgufValue::as_array) + .ok_or("missing per-layer array 'lfm2.attention.head_count_kv'")?; + if kv_per_layer.len() != num_hidden_layers { + return Err(format!( + "head_count_kv has {} entries for {num_hidden_layers} layers", + kv_per_layer.len() + )); + } + let mut layer_types = Vec::with_capacity(num_hidden_layers); + let mut kv_heads: Option = None; + for (i, v) in kv_per_layer.iter().enumerate() { + let n = v + .as_i64() + .and_then(|n| u64::try_from(n).ok()) // i32 array in real files + .ok_or_else(|| format!("head_count_kv[{i}] is not a non-negative integer"))?; + if n == 0 { + layer_types.push("conv".to_string()); + } else { + if kv_heads.is_some_and(|k| k != n) { + return Err(format!("head_count_kv mixes {kv_heads:?} and {n}")); + } + kv_heads = Some(n); + layer_types.push("full_attention".to_string()); + } + } + let Some(kv_heads) = kv_heads else { + return Err("no attention layers in head_count_kv".into()); + }; + + let embd = f + .tensor("token_embd.weight") + .ok_or("GGUF has no token_embd.weight tensor")?; + if embd.dims.len() != 2 || embd.dims[0] != hidden_size as u64 { + return Err(format!( + "token_embd.weight dims {:?} do not match embedding_length {hidden_size}", + embd.dims + )); + } + if f.tensor("output.weight").is_some() { + return Err("LFM2 GGUF carries an untied output.weight — unsupported".into()); + } + let eps = f + .get("lfm2.attention.layer_norm_rms_epsilon") + .and_then(GgufValue::as_f32) + .ok_or("missing 'lfm2.attention.layer_norm_rms_epsilon'")?; + let theta = f + .get("lfm2.rope.freq_base") + .and_then(GgufValue::as_f32) + .ok_or("missing 'lfm2.rope.freq_base'")?; + let cfg = Self { + vocab_size: usize::try_from(embd.dims[1]).map_err(|_| "vocab too large")?, + hidden_size, + num_hidden_layers, + num_attention_heads: meta_usize("lfm2.attention.head_count")?, + num_key_value_heads: usize::try_from(kv_heads).map_err(|_| "kv heads too large")?, + norm_eps: f64::from(eps), + rope_theta: theta, + conv_l_cache: meta_usize("lfm2.shortconv.l_cache")?, + block_ff_dim: meta_usize("lfm2.feed_forward_length")?, + block_multiple_of: 1, + block_auto_adjust_ff_dim: false, + layer_types, + eos_token_id: f + .get("tokenizer.ggml.eos_token_id") + .and_then(GgufValue::as_u64) + .and_then(|v| u32::try_from(v).ok()) + .map_or(EosIds::None, EosIds::One), + }; + cfg.validate()?; + Ok(cfg) + } + fn validate(&self) -> Result<(), String> { if self.layer_types.len() != self.num_hidden_layers { return Err(format!( @@ -219,6 +315,90 @@ pub fn load_from_dir( Ok(LoadedLfm2 { model, config }) } +/// GGUF (llama.cpp `lfm2` arch) tensor names → the HF checkpoint names the +/// safetensors remap chain already handles. The depthwise conv kernel is the +/// one shape special-case: llama.cpp stores it squeezed (`[k, channels]` in +/// ggml dims), the checkpoint as `[channels, 1, k]` — same bytes. +fn gguf_tensor_to_hf(info: &GgufTensorInfo) -> Option { + match info.name.as_str() { + "token_embd.weight" => { + return Some(GgufMap::Rename("model.embed_tokens.weight".into())); + } + "token_embd_norm.weight" => { + return Some(GgufMap::Rename("model.embedding_norm.weight".into())); + } + _ => {} + } + let rest = info.name.strip_prefix("blk.")?; + let (layer, field) = rest.split_once('.')?; + let layer: usize = layer.parse().ok()?; + if field == "shortconv.conv.weight" { + // ggml dims [k, channels] → row-major [channels, k] → the + // checkpoint's depthwise-Conv1d shape [channels, 1, k]. + let (&k, &channels) = (info.dims.first()?, info.dims.get(1)?); + return Some(GgufMap::Reshape( + format!("model.layers.{layer}.conv.conv.weight"), + vec![channels, 1, k], + )); + } + let mapped = match field { + "attn_norm.weight" => "operator_norm.weight", + "ffn_norm.weight" => "ffn_norm.weight", + "attn_q.weight" => "self_attn.q_proj.weight", + "attn_k.weight" => "self_attn.k_proj.weight", + "attn_v.weight" => "self_attn.v_proj.weight", + "attn_output.weight" => "self_attn.out_proj.weight", + "attn_q_norm.weight" => "self_attn.q_layernorm.weight", + "attn_k_norm.weight" => "self_attn.k_layernorm.weight", + "ffn_gate.weight" => "feed_forward.w1.weight", + "ffn_down.weight" => "feed_forward.w2.weight", + "ffn_up.weight" => "feed_forward.w3.weight", + "shortconv.in_proj.weight" => "conv.in_proj.weight", + "shortconv.out_proj.weight" => "conv.out_proj.weight", + _ => return None, + }; + Some(GgufMap::Rename(format!("model.layers.{layer}.{mapped}"))) +} + +/// Load an LFM2/LFM2.5 model straight from a **GGUF** file: hyperparameters +/// from the `lfm2.*` metadata (layer kinds from the per-layer kv-head +/// array), weights dequantized and driven through the exact store pipeline +/// the safetensors path uses. +pub fn load_from_gguf( + path: &Path, + device: &B::Device, +) -> Result, ImportError> { + let parse = |reason: String| ImportError::Parse { + file: path.to_path_buf(), + reason, + }; + 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"); + + let mut model = build::(&config, device); + let target_float = Tensor::::zeros([1], device).dtype(); + let mut store = SafetensorsStore::from_bytes(Some(blob)) + .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) + .allow_partial(true) + .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") + .with_key_remapping(r"(self_attn)\.k_layernorm\.weight$", "$1.k_norm.gamma") + .with_key_remapping(r"(feed_forward)\.w1\.", "$1.gate_proj.") + .with_key_remapping(r"(feed_forward)\.w2\.", "$1.down_proj.") + .with_key_remapping(r"(feed_forward)\.w3\.", "$1.up_proj.") + .with_key_remapping( + r"(operator_norm|ffn_norm|embedding_norm)\.weight$", + "$1.gamma", + ); + load_checked(&mut model, &mut store, path)?; + Ok(LoadedLfm2 { model, config }) +} + impl CausalLm for LoadedLfm2 { type Cache = Vec>; @@ -394,4 +574,109 @@ mod tests { let out = loaded.greedy_generate(&[1, 2], 3, &device).unwrap(); assert!(out.len() <= 3); } + + /// A synthetic GGUF header shaped like the LFM2.5-1.2B file. + fn toy_gguf() -> GgufFile { + use crate::gguf::GgmlType; + let meta = |k: &str, v: GgufValue| (k.to_string(), v); + let kv_array = GgufValue::Array( + [0u32, 2, 0] + .iter() + .map(|&v| GgufValue::U32(v)) + .collect::>(), + ); + GgufFile { + path: std::path::PathBuf::new(), + version: 3, + metadata: vec![ + meta("general.architecture", GgufValue::Str("lfm2".into())), + meta("lfm2.embedding_length", GgufValue::U32(16)), + meta("lfm2.block_count", GgufValue::U32(3)), + meta("lfm2.feed_forward_length", GgufValue::U32(32)), + meta("lfm2.attention.head_count", GgufValue::U32(4)), + meta("lfm2.attention.head_count_kv", kv_array), + meta( + "lfm2.attention.layer_norm_rms_epsilon", + GgufValue::F32(1e-5), + ), + meta("lfm2.rope.freq_base", GgufValue::F32(1e6)), + meta("lfm2.shortconv.l_cache", GgufValue::U32(3)), + meta("tokenizer.ggml.eos_token_id", GgufValue::U32(7)), + ], + tensors: vec![GgufTensorInfo { + name: "token_embd.weight".into(), + dims: vec![16, 48], + dtype: GgmlType::F32, + offset: 0, + }], + alignment: 32, + data_offset: 0, + } + } + + #[test] + fn config_from_gguf_derives_layer_types_from_kv_array() { + let cfg = Lfm2Config::from_gguf(&toy_gguf()).expect("parses"); + assert_eq!( + cfg.layer_types, + vec!["conv", "full_attention", "conv"], + "0 = conv, nonzero = attention" + ); + assert_eq!(cfg.num_key_value_heads, 2); + assert_eq!(cfg.vocab_size, 48); // from token_embd dims[1] + assert_eq!(cfg.ff_dim(), 32); // GGUF value is pre-adjusted + assert!(cfg.eos_token_id.contains(7)); + } + + #[test] + fn config_from_gguf_rejects_mixed_kv_and_missing_array() { + let mut f = toy_gguf(); + f.metadata[5].1 = GgufValue::Array(vec![ + GgufValue::U32(4), + GgufValue::U32(2), + GgufValue::U32(0), + ]); + assert!(Lfm2Config::from_gguf(&f).unwrap_err().contains("mixes")); + + let mut f = toy_gguf(); + f.metadata + .retain(|(k, _)| k != "lfm2.attention.head_count_kv"); + assert!(Lfm2Config::from_gguf(&f).is_err()); + } + + #[test] + fn gguf_names_map_onto_hf_checkpoint_names_incl_conv_reshape() { + use crate::gguf::GgmlType; + let info = |name: &str, dims: &[u64]| GgufTensorInfo { + name: name.into(), + dims: dims.to_vec(), + dtype: GgmlType::F32, + offset: 0, + }; + let name_of = |i: &GgufTensorInfo| match gguf_tensor_to_hf(i) { + Some(GgufMap::Rename(n)) => n, + other => panic!("expected Rename, got {other:?}"), + }; + assert_eq!( + name_of(&info("blk.2.attn_q_norm.weight", &[64])), + "model.layers.2.self_attn.q_layernorm.weight" + ); + assert_eq!( + name_of(&info("blk.0.shortconv.in_proj.weight", &[2048, 6144])), + "model.layers.0.conv.in_proj.weight" + ); + assert_eq!( + name_of(&info("token_embd_norm.weight", &[2048])), + "model.embedding_norm.weight" + ); + // The conv kernel un-squeezes to the checkpoint's [channels, 1, k]. + match gguf_tensor_to_hf(&info("blk.0.shortconv.conv.weight", &[3, 2048])) { + Some(GgufMap::Reshape(n, shape)) => { + assert_eq!(n, "model.layers.0.conv.conv.weight"); + assert_eq!(shape, vec![2048, 1, 3]); + } + other => panic!("expected Reshape, got {other:?}"), + } + assert!(gguf_tensor_to_hf(&info("rope_freqs.weight", &[64])).is_none()); + } } diff --git a/crates/mummu/src/models/qwen2.rs b/crates/mummu/src/models/qwen2.rs index 9cfce28..e09ce41 100644 --- a/crates/mummu/src/models/qwen2.rs +++ b/crates/mummu/src/models/qwen2.rs @@ -15,7 +15,7 @@ use burn::nn::{Embedding, EmbeddingConfig, Linear, LinearConfig, RmsNorm, RmsNor use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, SafetensorsStore}; use burn::tensor::{Int, Tensor, TensorData, backend::Backend}; -use crate::gguf::{GgufFile, GgufValue}; +use crate::gguf::{GgufFile, GgufMap, GgufTensorInfo, GgufValue}; use crate::import::{CastFloatAdapter, ImportError, load_checked, required_file}; use crate::models::CausalLm; use crate::nn::{ @@ -270,7 +270,11 @@ pub fn load_from_dir( /// GGUF (llama.cpp) tensor names → the HF checkpoint names the safetensors /// remap chain already handles. `None` for anything unrecognized — the blob /// writer turns that into a loud error rather than dropping weights. -fn gguf_tensor_to_hf(name: &str) -> Option { +fn gguf_tensor_to_hf(info: &GgufTensorInfo) -> Option { + qwen2_gguf_name(&info.name).map(GgufMap::Rename) +} + +fn qwen2_gguf_name(name: &str) -> Option { match name { "token_embd.weight" => return Some("model.embed_tokens.weight".into()), "output_norm.weight" => return Some("model.norm.weight".into()), @@ -590,24 +594,24 @@ mod tests { #[test] fn gguf_names_map_onto_hf_checkpoint_names() { assert_eq!( - gguf_tensor_to_hf("token_embd.weight").as_deref(), + qwen2_gguf_name("token_embd.weight").as_deref(), Some("model.embed_tokens.weight") ); assert_eq!( - gguf_tensor_to_hf("blk.27.attn_q.bias").as_deref(), + qwen2_gguf_name("blk.27.attn_q.bias").as_deref(), Some("model.layers.27.self_attn.q_proj.bias") ); assert_eq!( - gguf_tensor_to_hf("blk.0.ffn_down.weight").as_deref(), + qwen2_gguf_name("blk.0.ffn_down.weight").as_deref(), Some("model.layers.0.mlp.down_proj.weight") ); assert_eq!( - gguf_tensor_to_hf("output_norm.weight").as_deref(), + qwen2_gguf_name("output_norm.weight").as_deref(), Some("model.norm.weight") ); // Unknown names must map to None (the writer errors loudly). - assert_eq!(gguf_tensor_to_hf("rope_freqs.weight"), None); - assert_eq!(gguf_tensor_to_hf("blk.x.attn_q.weight"), None); + assert_eq!(qwen2_gguf_name("rope_freqs.weight"), None); + assert_eq!(qwen2_gguf_name("blk.x.attn_q.weight"), None); } #[test] diff --git a/crates/mummu/src/tokenizer.rs b/crates/mummu/src/tokenizer.rs index b720fc0..a46c3eb 100644 --- a/crates/mummu/src/tokenizer.rs +++ b/crates/mummu/src/tokenizer.rs @@ -17,6 +17,7 @@ use tokenizers::normalizers::unicode::NFC; use tokenizers::pre_tokenizers::byte_level::ByteLevel; use tokenizers::pre_tokenizers::sequence::Sequence; use tokenizers::pre_tokenizers::split::{Split, SplitPattern}; +use tokenizers::processors::template::TemplateProcessing; use tokenizers::{AddedToken, SplitDelimiterBehavior, Tokenizer}; use crate::gguf::{GgufFile, GgufValue}; @@ -27,16 +28,31 @@ const TOKEN_TYPE_CONTROL: i64 = 3; const TOKEN_TYPE_USER_DEFINED: i64 = 4; const TOKEN_TYPE_UNUSED: i64 = 5; -/// The GPT-2-style byte-level BPE pre-tokenizer regexes, keyed by -/// `tokenizer.ggml.pre` — the same registry llama.cpp keeps in its vocab -/// loader. Only families we actually run are listed; unknown ids are a loud -/// error (a wrong regex silently produces wrong token ids). -fn pre_tokenizer_regex(pre: &str) -> Option<&'static str> { +/// How a model family's byte-level BPE pipeline is configured — the part a +/// GGUF names by `tokenizer.ggml.pre` id instead of carrying explicitly. +struct PreSpec { + /// The pre-tokenizer split regex (from the family's `tokenizer.json`). + regex: &'static str, + /// Whether the pipeline NFC-normalizes first. + nfc: bool, +} + +/// The per-family registry, keyed by `tokenizer.ggml.pre` — the same registry +/// llama.cpp keeps in its vocab loader. Only families we actually run are +/// listed; unknown ids are a loud error (a wrong regex silently produces +/// wrong token ids). +fn pre_spec(pre: &str) -> Option { match pre { // Qwen2/2.5 (matches the checkpoint's tokenizer.json byte for byte). - "qwen2" => Some( - r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+", - ), + "qwen2" => Some(PreSpec { + regex: r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+", + nfc: true, + }), + // LFM2/LFM2.5: digits split in groups of ≤3, no normalizer. + "lfm2" => Some(PreSpec { + regex: r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+", + nfc: false, + }), _ => None, } } @@ -70,7 +86,7 @@ pub fn tokenizer_from_gguf(f: &GgufFile) -> Result { .get("tokenizer.ggml.pre") .and_then(GgufValue::as_str) .ok_or("missing GGUF metadata 'tokenizer.ggml.pre'")?; - let Some(regex) = pre_tokenizer_regex(pre) else { + let Some(spec) = pre_spec(pre) else { return Err(format!("unknown pre-tokenizer id '{pre}'")); }; @@ -84,7 +100,10 @@ pub fn tokenizer_from_gguf(f: &GgufFile) -> Result { )); } - // The BPE vocab: every NORMAL token, id = array index. + // The BPE vocab: every non-padding token at id = array index. Control/ + // user-defined tokens go in TOO — `add_tokens` below then reuses the + // model id it finds, which is what lets added tokens live at LOW ids + // (LFM2 puts its 500+ specials at 0..) instead of only after the vocab. let mut vocab = Vocab::default(); // (index, content, is_special) for everything re-added post-BPE. let mut added: Vec<(usize, String, bool)> = Vec::new(); @@ -100,8 +119,10 @@ pub fn tokenizer_from_gguf(f: &GgufFile) -> Result { TOKEN_TYPE_NORMAL => { vocab.insert(text.to_string(), index as u32); } - TOKEN_TYPE_CONTROL => added.push((index, text.to_string(), true)), - TOKEN_TYPE_USER_DEFINED => added.push((index, text.to_string(), false)), + TOKEN_TYPE_CONTROL | TOKEN_TYPE_USER_DEFINED => { + vocab.insert(text.to_string(), index as u32); + added.push((index, text.to_string(), ty == TOKEN_TYPE_CONTROL)); + } TOKEN_TYPE_UNUSED => {} // vocab-padding entries ([PADn]) other => return Err(format!("token {index} has unsupported type {other}")), } @@ -127,9 +148,11 @@ pub fn tokenizer_from_gguf(f: &GgufFile) -> Result { .build() .map_err(|e| format!("BPE build: {e}"))?; let mut tok = Tokenizer::new(bpe); - tok.with_normalizer(Some(NFC)); + if spec.nfc { + tok.with_normalizer(Some(NFC)); + } let split = Split::new( - SplitPattern::Regex(regex.to_string()), + SplitPattern::Regex(spec.regex.to_string()), SplitDelimiterBehavior::Isolated, false, ) @@ -139,7 +162,37 @@ pub fn tokenizer_from_gguf(f: &GgufFile) -> Result { let byte_level = ByteLevel::new(false, false, false); tok.with_pre_tokenizer(Some(Sequence::new(vec![split.into(), byte_level.into()]))); tok.with_decoder(Some(byte_level)); - tok.with_post_processor(Some(byte_level)); + + // `tokenizer.ggml.add_bos_token` → a BOS-prepending template processor + // (what the family's tokenizer.json does); otherwise the offsets-only + // ByteLevel processor. + let add_bos = f + .get("tokenizer.ggml.add_bos_token") + .and_then(|v| match *v { + GgufValue::Bool(b) => Some(b), + _ => None, + }) + .unwrap_or(false); + if add_bos { + let bos_id = f + .get("tokenizer.ggml.bos_token_id") + .and_then(GgufValue::as_u64) + .and_then(|v| u32::try_from(v).ok()) + .ok_or("add_bos_token is set but tokenizer.ggml.bos_token_id is missing")?; + let bos = tokens + .get(bos_id as usize) + .and_then(GgufValue::as_str) + .ok_or_else(|| format!("bos_token_id {bos_id} is out of vocab range"))?; + let template = TemplateProcessing::builder() + .try_single(format!("{bos} $A")) + .map_err(|e| format!("BOS template: {e}"))? + .special_tokens(vec![(bos.to_string(), bos_id)]) + .build() + .map_err(|e| format!("BOS template: {e}"))?; + tok.with_post_processor(Some(template)); + } else { + tok.with_post_processor(Some(byte_level)); + } // Re-add control/user-defined tokens in id order, then verify every id // landed where the GGUF says it lives. diff --git a/crates/mummu/tests/real_gguf.rs b/crates/mummu/tests/real_gguf.rs index a1832e1..bf5b022 100644 --- a/crates/mummu/tests/real_gguf.rs +++ b/crates/mummu/tests/real_gguf.rs @@ -217,6 +217,199 @@ fn real_qwen2_gguf_tokenizer_matches_the_hf_tokenizer() { ); } +// ---- LFM2.5 (the hybrid conv+attention architecture) ---------------------- + +fn lfm2_gguf_path() -> Option { + let path = PathBuf::from(std::env::var_os("MUMMU_LFM2_GGUF_PATH")?); + path.is_file().then_some(path) +} + +fn lfm2_dir() -> Option { + std::env::var_os("MUMMU_LFM2_DIR") + .map(PathBuf::from) + .filter(|d| d.join("model.safetensors").is_file()) +} + +/// LFM2.5 mapping proof against the TRUE weights: every F32 tensor in the +/// GGUF (norms, per-head q/k norms, and the depthwise conv kernels — the +/// reshape special-case) must equal the bf16 safetensors bit-exactly. +#[test] +#[ignore = "needs the local LFM2.5 GGUF (MUMMU_LFM2_GGUF_PATH) + safetensors (MUMMU_LFM2_DIR)"] +fn real_lfm2_gguf_f32_tensors_match_the_true_weights() { + let Some(path) = lfm2_gguf_path() else { + panic!("set MUMMU_LFM2_GGUF_PATH to the lfm2.5-1.2b q4_k_m gguf"); + }; + let Some(dir) = lfm2_dir() else { + panic!("set MUMMU_LFM2_DIR to the same model's safetensors dir"); + }; + let st = dir.join("model.safetensors"); + let f = GgufFile::open(&path).expect("header parses"); + + // (gguf name, safetensors name) — covers the plain rename, the per-head + // norms, and the conv-kernel reshape (same bytes, unsqueezed shape). + let pairs = [ + ("token_embd_norm.weight", "model.embedding_norm.weight"), + ( + "blk.0.attn_norm.weight", + "model.layers.0.operator_norm.weight", + ), + ( + "blk.2.attn_q_norm.weight", + "model.layers.2.self_attn.q_layernorm.weight", + ), + ( + "blk.0.shortconv.conv.weight", + "model.layers.0.conv.conv.weight", + ), + ( + "blk.15.shortconv.conv.weight", + "model.layers.15.conv.conv.weight", + ), + ]; + for (ours_name, ref_name) in pairs { + let ours = f.read_tensor_f32(ours_name).expect("dequantizes"); + let reference = safetensors_bf16_f32(&st, ref_name); + assert_eq!(ours.len(), reference.len(), "{ours_name}: same size"); + let exact = ours + .iter() + .zip(&reference) + .all(|(a, b)| a.to_bits() == b.to_bits()); + assert!(exact, "{ours_name} must be bit-exact vs {ref_name}"); + } + eprintln!( + "[real_gguf/lfm2] {} F32 tensors bit-exact vs safetensors (incl. conv kernels)", + pairs.len() + ); +} + +/// The LFM2.5 tokenizer rebuilt from GGUF metadata must match the +/// checkpoint's `tokenizer.json` — with AND without the BOS-adding +/// post-processor (LFM2 sets `add_bos_token`, unlike Qwen2). +#[test] +#[ignore = "needs the local LFM2.5 GGUF (MUMMU_LFM2_GGUF_PATH) + safetensors dir (MUMMU_LFM2_DIR)"] +fn real_lfm2_gguf_tokenizer_matches_the_hf_tokenizer() { + let Some(path) = lfm2_gguf_path() else { + panic!("set MUMMU_LFM2_GGUF_PATH to the lfm2.5-1.2b q4_k_m gguf"); + }; + let dir = lfm2_dir().expect("set MUMMU_LFM2_DIR to the safetensors dir"); + let f = GgufFile::open(&path).expect("header parses"); + let ours = mummu::tokenizer::tokenizer_from_gguf(&f).expect("tokenizer builds"); + let reference = + tokenizers::Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer.json"); + + let chat = mummu::chat::ChatMl::lfm2(); + let battery = [ + chat.render(&[ + mummu::chat::Turn::system("You are a concise assistant."), + mummu::chat::Turn::user("What is 2+2? Answer in one short sentence."), + ]), + "The quick brown fox jumps over the lazy dog.".into(), + "héllo wörld — 世界 · Ω ≠ ω · 🦀🔥".into(), + "1234567890 3.14159 1e-6 0x7F".into(), // digits split in ≤3 groups + " leading spaces, trailing \n\nnewlines\r\nand\ttabs ".into(), + "don't they're we've I'll it's CAN'T".into(), + ]; + for text in &battery { + for add_special in [true, false] { + let a = ours.encode(text.as_str(), add_special).expect("encodes"); + let b = reference + .encode(text.as_str(), add_special) + .expect("encodes"); + assert_eq!( + a.get_ids(), + b.get_ids(), + "ids diverge (add_special={add_special}) on {text:?}" + ); + } + } + eprintln!( + "[real_gguf/lfm2] tokenizer-from-GGUF: {} prompts × 2 modes byte-identical", + battery.len() + ); +} + +/// END-TO-END for the hybrid: the LFM2.5 Q4_K_M GGUF alone becomes a running +/// model on the GPU (conv layers, attention layers, per-head norms — all +/// mapped from llama.cpp naming), greedy-decodes a correct answer, and its +/// first-token logits agree with the bf16 safetensors build. +#[test] +#[ignore = "needs the local LFM2.5 GGUF (MUMMU_LFM2_GGUF_PATH) + safetensors dir (MUMMU_LFM2_DIR) + GPU"] +fn real_lfm2_gguf_loads_and_decodes_on_gpu() { + use mummu::models::lfm2; + let Some(path) = lfm2_gguf_path() else { + panic!("set MUMMU_LFM2_GGUF_PATH to the lfm2.5-1.2b q4_k_m gguf"); + }; + let dir = lfm2_dir().expect("set MUMMU_LFM2_DIR to the safetensors dir"); + assert!(use_gpu(), "this proof wants the real GPU"); + let device = burn::tensor::Device::::default(); + + let header = GgufFile::open(&path).expect("header parses"); + let tok = mummu::tokenizer::tokenizer_from_gguf(&header).expect("tokenizer from metadata"); + let chat = mummu::chat::ChatMl::lfm2(); + let prompt_text = chat.render(&[ + mummu::chat::Turn::system("You are a concise assistant."), + mummu::chat::Turn::user("What is 2+2? Answer in one short sentence."), + ]); + // The rendered template already carries <|startoftext|>; no auto-BOS. + let prompt = tok + .encode(prompt_text, false) + .expect("prompt encodes") + .get_ids() + .to_vec(); + + let gguf_model = lfm2::load_from_gguf::(&path, &device).expect("gguf load is checked"); + assert_eq!(gguf_model.config.num_hidden_layers, 16); + assert_eq!( + gguf_model + .config + .layer_types + .iter() + .filter(|t| *t == "conv") + .count(), + 10, + "the 1.2B is 10 conv + 6 attention" + ); + let ids = gguf_model + .greedy_generate(&prompt, 32, &device) + .expect("decode"); + let text = tok.decode(&ids, true).expect("ids decode"); + eprintln!("[real_gguf/lfm2] Q4_K_M greedy: {text:?}"); + assert!(text.contains('4'), "expected the answer 4 in: {text:?}"); + + let logits_of = |m: &lfm2::LoadedLfm2| -> Vec { + let mut cache = m.new_cache(); + m.forward(&prompt, 0, &mut cache, &device) + .into_data() + .to_vec::() + .expect("logits read back") + }; + let gguf_logits = logits_of(&gguf_model); + drop(gguf_model); + let st_model = lfm2::load_from_dir::(&dir, &device).expect("safetensors load"); + let st_logits = logits_of(&st_model); + drop(st_model); + + let argmax = |v: &[f32]| -> usize { + let (mut best, mut best_v) = (0usize, f32::NEG_INFINITY); + for (i, &x) in v.iter().enumerate() { + if x > best_v { + (best, best_v) = (i, x); + } + } + best + }; + let cos = cosine(&gguf_logits, &st_logits); + let (g_top, s_top) = (argmax(&gguf_logits), argmax(&st_logits)); + eprintln!( + "[real_gguf/lfm2] first-token logits: cosine {cos:.5} vs bf16 · top-1 {g_top} vs {s_top}" + ); + assert_eq!(g_top, s_top, "Q4_K_M must agree with bf16 on the top token"); + assert!( + cos > 0.95, + "logit cosine {cos} — quantization noise should be small, layout bugs are not" + ); +} + /// END-TO-END: the Q4_K_M GGUF alone (config + weights from the one file) /// becomes a running model on the GPU — greedy-decodes a correct answer, and /// its first-token logits agree with the bf16 safetensors build of the same From f557f6099ea6d27321a1a46c57d9588843ed5d8c Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Mon, 13 Jul 2026 08:36:05 -0500 Subject: [PATCH 6/6] =?UTF-8?q?feat(p3):=20registry=20learns=20GGUF=20?= =?UTF-8?q?=E2=80=94=20single-file=20model=20specs=20install=20end-to-end?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ModelSpec.format: WeightFormat::Safetensors | Gguf { file } (serde-defaulted, so manifests written before the field still parse) - GGUF specs fetch their one file through the resumable/verified hub path; gguf_path() names the local file; validation rejects empty, traversal, absolute, and non-.gguf file names - catalog: single-file Q4_K_M entries for Qwen2.5-1.5B-Instruct and LFM2.5-1.2B-Instruct (quarter the download of the safetensors) REAL-NETWORK proof (real_hub.rs): the LFM2.5 GGUF spec installs end-to-end — 697 MB fetched via spec.fetch, header parses as lfm2 (148 tensors), tokenizer builds from its metadata. 140 unit tests, clippy clean. Co-Authored-By: Claude Opus 4.8 --- ROADMAP.md | 6 ++ crates/mummu/src/registry.rs | 128 ++++++++++++++++++++++++++++++--- crates/mummu/tests/real_hub.rs | 35 +++++++++ 3 files changed, 161 insertions(+), 8 deletions(-) diff --git a/ROADMAP.md b/ROADMAP.md index 4b7fa0e..b292b3f 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -262,6 +262,12 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari `spec.fetch(models_root, progress)` onto the hub downloader; built-in catalog: Qwen2.5-1.5B/0.5B, LFM2.5-1.2B, all-MiniLM (repo ids match laurelane's validated constants); the network proof now fetches spec-driven. Weight-format/dtype/chat-template fields accrete as those import paths land.* + *(2026-07-13) **Weight-format field landed**: `ModelSpec.format` (`WeightFormat::Safetensors` | + `Gguf { file }`, serde-defaulted so old manifests still parse); GGUF specs fetch the one file via + the resumable hub path, `gguf_path()` names where it lands, validation rejects unsafe file names. + Catalog gains single-file Q4_K_M entries for Qwen2.5-1.5B and LFM2.5-1.2B (quarter the download of + the safetensors). REAL-NETWORK proof (`real_hub.rs`): the LFM2.5 GGUF spec installed end-to-end — + 697 MB fetched, header parses as `lfm2` (148 tensors), tokenizer built from its metadata.* - [ ] **Import validation** — checked load + a first-token parity smoke against a reference before a model is marked trusted; a clear error taxonomy (missing file, bad shard, key mismatch, unsupported dtype). diff --git a/crates/mummu/src/registry.rs b/crates/mummu/src/registry.rs index 22bf095..bb5d615 100644 --- a/crates/mummu/src/registry.rs +++ b/crates/mummu/src/registry.rs @@ -19,6 +19,23 @@ pub enum Architecture { MiniLm, } +/// How the checkpoint's weights are stored — which fetch + load path a spec +/// takes. +#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub enum WeightFormat { + /// `config.json` + `tokenizer.json` + `model.safetensors` (or a shard + /// index); fetched by [`hub::fetch_model`], loaded by `load_from_dir`. + #[default] + Safetensors, + /// One self-contained `.gguf` file in the repo (config + tokenizer + + /// weights in the metadata); loaded by the architecture's + /// `load_from_gguf` + [`crate::tokenizer::tokenizer_from_gguf`]. + Gguf { + /// The file name inside the repo, e.g. `qwen2.5-1.5b-instruct-q4_k_m.gguf`. + file: String, + }, +} + /// A declarative model manifest entry. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct ModelSpec { @@ -29,6 +46,9 @@ pub struct ModelSpec { /// Git revision (tag, branch, or commit) — pin for reproducibility. pub revision: String, pub architecture: Architecture, + /// Weight storage (absent in older manifests = safetensors). + #[serde(default)] + pub format: WeightFormat, /// Rough on-disk size, for settings UIs and fit checks (0 = unknown). pub disk_bytes_estimate: u64, } @@ -50,6 +70,15 @@ impl ModelSpec { if self.revision.is_empty() { return Err("revision must be non-empty (pin something)".into()); } + if let WeightFormat::Gguf { file } = &self.format { + let safe = !file.is_empty() + && !file.contains("..") + && !file.starts_with('/') + && file.ends_with(".gguf"); + if !safe { + return Err(format!("bad gguf file name {file:?}")); + } + } Ok(()) } @@ -59,21 +88,39 @@ impl ModelSpec { models_root.join(&self.name) } + /// For a GGUF spec: the local path of the model file after [`Self::fetch`]. + #[must_use] + pub fn gguf_path(&self, models_root: &Path) -> Option { + match &self.format { + WeightFormat::Gguf { file } => Some(self.dir(models_root).join(file)), + WeightFormat::Safetensors => None, + } + } + /// Download this model into `models_root` (resumable, cache-first; see - /// [`hub::fetch_model`]) and return its directory, ready for the - /// architecture's `load_from_dir`. + /// [`hub::fetch_model`] / [`hub::fetch_file`]) and return its directory, + /// ready for the architecture's `load_from_dir` / `load_from_gguf`. pub fn fetch( &self, models_root: &Path, on_progress: impl FnMut(Progress<'_>), ) -> Result { assert!(self.validate().is_ok(), "fetch of an invalid spec"); - hub::fetch_model( - &self.repo, - &self.revision, - &self.dir(models_root), - on_progress, - ) + let dir = self.dir(models_root); + match &self.format { + WeightFormat::Safetensors => { + hub::fetch_model(&self.repo, &self.revision, &dir, on_progress) + } + WeightFormat::Gguf { file } => { + let url = hub::hub_file_url(&self.repo, &self.revision, file); + std::fs::create_dir_all(&dir).map_err(|e| HubError::Io { + path: dir.clone(), + reason: e.to_string(), + })?; + hub::fetch_file(&url, &dir.join(file), on_progress)?; + Ok(dir) + } + } } } @@ -87,6 +134,7 @@ pub fn catalog() -> Vec { repo: "Qwen/Qwen2.5-1.5B-Instruct".into(), revision: "main".into(), architecture: Architecture::Qwen2, + format: WeightFormat::Safetensors, disk_bytes_estimate: 3_100_000_000, }, ModelSpec { @@ -94,6 +142,7 @@ pub fn catalog() -> Vec { repo: "Qwen/Qwen2.5-0.5B-Instruct".into(), revision: "main".into(), architecture: Architecture::Qwen2, + format: WeightFormat::Safetensors, disk_bytes_estimate: 1_000_000_000, }, ModelSpec { @@ -101,6 +150,7 @@ pub fn catalog() -> Vec { repo: "LiquidAI/LFM2.5-1.2B-Instruct".into(), revision: "main".into(), architecture: Architecture::Lfm2, + format: WeightFormat::Safetensors, disk_bytes_estimate: 2_400_000_000, }, ModelSpec { @@ -108,8 +158,31 @@ pub fn catalog() -> Vec { repo: "sentence-transformers/all-MiniLM-L6-v2".into(), revision: "main".into(), architecture: Architecture::MiniLm, + format: WeightFormat::Safetensors, disk_bytes_estimate: 91_000_000, }, + // Single-file GGUF variants — quarter the download, same model + // (proven vs the bf16 safetensors builds in tests/real_gguf.rs). + ModelSpec { + name: "qwen2.5-1.5b-instruct-q4km".into(), + repo: "Qwen/Qwen2.5-1.5B-Instruct-GGUF".into(), + revision: "main".into(), + architecture: Architecture::Qwen2, + format: WeightFormat::Gguf { + file: "qwen2.5-1.5b-instruct-q4_k_m.gguf".into(), + }, + disk_bytes_estimate: 1_120_000_000, + }, + ModelSpec { + name: "lfm2.5-1.2b-q4km".into(), + repo: "LiquidAI/LFM2.5-1.2B-Instruct-GGUF".into(), + revision: "main".into(), + architecture: Architecture::Lfm2, + format: WeightFormat::Gguf { + file: "LFM2.5-1.2B-Instruct-Q4_K_M.gguf".into(), + }, + disk_bytes_estimate: 731_000_000, + }, ]; debug_assert!( entries.iter().all(|s| s.validate().is_ok()), @@ -164,4 +237,43 @@ mod tests { assert_eq!(back.name, spec.name); assert_eq!(back.architecture, spec.architecture); } + + #[test] + fn gguf_specs_round_trip_and_old_manifests_default_to_safetensors() { + let gguf = catalog() + .into_iter() + .find(|s| matches!(s.format, WeightFormat::Gguf { .. })) + .expect("catalog has a gguf entry"); + let json = serde_json::to_string(&gguf).unwrap(); + let back: ModelSpec = serde_json::from_str(&json).unwrap(); + assert_eq!(back.format, gguf.format); + let file = match &gguf.format { + WeightFormat::Gguf { file } => file.clone(), + WeightFormat::Safetensors => unreachable!(), + }; + assert_eq!( + gguf.gguf_path(Path::new("root")), + Some(Path::new("root").join(&gguf.name).join(file)) + ); + + // A manifest written before `format` existed still deserializes. + let old = r#"{"name":"m","repo":"a/b","revision":"main", + "architecture":"Qwen2","disk_bytes_estimate":1}"#; + let back: ModelSpec = serde_json::from_str(old).unwrap(); + assert_eq!(back.format, WeightFormat::Safetensors); + assert_eq!(back.gguf_path(Path::new("root")), None); + } + + #[test] + fn bad_gguf_file_names_are_rejected() { + let mut spec = catalog()[0].clone(); + for bad in ["", "../up.gguf", "/abs.gguf", "weights.bin"] { + spec.format = WeightFormat::Gguf { file: bad.into() }; + assert!(spec.validate().is_err(), "{bad:?} must not validate"); + } + spec.format = WeightFormat::Gguf { + file: "ok-model.q4_k_m.gguf".into(), + }; + assert!(spec.validate().is_ok()); + } } diff --git a/crates/mummu/tests/real_hub.rs b/crates/mummu/tests/real_hub.rs index daedb2c..04aa153 100644 --- a/crates/mummu/tests/real_hub.rs +++ b/crates/mummu/tests/real_hub.rs @@ -197,3 +197,38 @@ fn hub_fetches_the_cpu_tier_qwen() { } eprintln!("[real_hub] 0.5B fetched into {}", dir.display()); } + +/// Registry → single-file GGUF install: the catalog's LFM2.5 Q4_K_M spec +/// downloads through `fetch` (one ~700 MB file, resumable/cache-first like +/// every hub fetch), lands where `gguf_path` says, parses as a valid GGUF of +/// the right architecture, and its metadata builds the tokenizer — the whole +/// app-facing "install a GGUF model" path in one proof. +#[test] +#[ignore = "needs network (MUMMU_HUB_DEST names the download dir; ~700 MB)"] +fn hub_gguf_spec_downloads_and_parses() { + use mummu::registry::WeightFormat; + let Some(dest) = std::env::var_os("MUMMU_HUB_DEST").map(PathBuf::from) else { + panic!("set MUMMU_HUB_DEST to a scratch dir for the ~700 MB download"); + }; + let spec = mummu::registry::catalog() + .into_iter() + .find(|s| matches!(s.format, WeightFormat::Gguf { .. }) && s.name.starts_with("lfm2.5")) + .expect("the catalog has the LFM2.5 GGUF entry"); + + let mut events = 0u64; + let dir = spec.fetch(&dest, |_| events += 1).expect("gguf fetch"); + let path = spec.gguf_path(&dest).expect("gguf specs have a file path"); + assert!(path.starts_with(&dir), "file lives in the spec's dir"); + assert!(path.is_file(), "downloaded file exists at {path:?}"); + + let f = mummu::gguf::GgufFile::open(&path).expect("valid GGUF"); + assert_eq!(f.architecture(), Some("lfm2")); + assert!(!f.tensors.is_empty()); + let tok = mummu::tokenizer::tokenizer_from_gguf(&f).expect("tokenizer from metadata"); + assert!(tok.token_to_id("<|im_end|>").is_some()); + eprintln!( + "[real_hub/gguf] {} → {} tensors, tokenizer ok ({events} progress events)", + spec.name, + f.tensors.len() + ); +}