From 7b73d01ac0f66724e6397ebd1ba66f65a06e9b30 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sat, 11 Jul 2026 06:47:30 -0500 Subject: [PATCH 1/6] chore(deps): compatible cargo update (bytemuck, cc, exr, pollster, rand, sha1, tinyvec) wgpu 30 / tokenizers 0.23 majors were offered by cargo-upgrade and reverted: both are intentional pins (burn 0.21 + wgpu 29 + tokenizers 0.22 is the parity-validated combination; burn 0.21 itself resolves wgpu 29, a direct wgpu 30 handle would split the dependency tree). All 80 unit tests green, clippy clean. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 50 +++++++++++++++++++++++++------------------------- 1 file changed, 25 insertions(+), 25 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ec336e5..fa86462 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -802,18 +802,18 @@ dependencies = [ [[package]] name = "bytemuck" -version = "1.25.0" +version = "1.25.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +checksum = "d6aedf8ae72766347502cf3cb4f41cf5e9cc37d28bee90f1fdaaae15f9cf9424" dependencies = [ "bytemuck_derive", ] [[package]] name = "bytemuck_derive" -version = "1.10.2" +version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +checksum = "f65693059b6b9c588b9f62fed1cedbf0a8b805631457ea162d68f0de186f3de5" dependencies = [ "proc-macro2", "quote", @@ -875,7 +875,7 @@ dependencies = [ "memmap2", "num-traits", "num_cpus", - "rand 0.9.4", + "rand 0.9.5", "rand_distr 0.5.1", "rayon", "safetensors 0.7.0", @@ -911,9 +911,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.66" +version = "1.2.67" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f5d6cac793997bd970000024b2934968efe83b382de4fdcf4fcb46b6ee4ad996" +checksum = "e17dd265a7d0f31ef544e1b20e03add05d3b45b491b633b10d67145d2acc1a38" dependencies = [ "find-msvc-tools", "jobserver", @@ -2179,9 +2179,9 @@ dependencies = [ [[package]] name = "exr" -version = "1.74.1" +version = "1.74.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6be87932f10230a4339ab394edd8e4611fcb72553d8295b4d52ea55249b3daa5" +checksum = "711fe42c9964295e01ee3fba3f9fe0e1d24b98886950d68efe81b1c76e21adf3" dependencies = [ "bit_field", "half", @@ -2267,7 +2267,7 @@ checksum = "c2d1f04709a8ac06e8e8042875a3c466cc4832d3c1a18dbcb9dba3c6e83046bc" dependencies = [ "half", "num-traits", - "rand 0.9.4", + "rand 0.9.5", "rand_distr 0.5.1", ] @@ -2645,7 +2645,7 @@ dependencies = [ "cfg-if", "crunchy", "num-traits", - "rand 0.9.4", + "rand 0.9.5", "rand_distr 0.5.1", "serde", "zerocopy", @@ -3931,9 +3931,9 @@ dependencies = [ [[package]] name = "pollster" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2560ec4d09f0d4db1fb45518d099d74954dd3da5855a0fcb84f2779b5b41a59e" +checksum = "bc6355899e1c9462875b6757c79f3caa011a1fdae12bbb1a2e72dd1f234f8336" [[package]] name = "portable-atomic" @@ -4144,9 +4144,9 @@ checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] name = "rand" -version = "0.8.6" +version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" dependencies = [ "libc", "rand_chacha 0.3.1", @@ -4155,9 +4155,9 @@ dependencies = [ [[package]] name = "rand" -version = "0.9.4" +version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" dependencies = [ "rand_chacha 0.9.0", "rand_core 0.9.5", @@ -4225,7 +4225,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" dependencies = [ "num-traits", - "rand 0.9.4", + "rand 0.9.5", ] [[package]] @@ -4280,7 +4280,7 @@ dependencies = [ "num-traits", "paste", "profiling", - "rand 0.9.4", + "rand 0.9.5", "rand_chacha 0.9.0", "simd_helpers", "thiserror 2.0.18", @@ -4729,9 +4729,9 @@ dependencies = [ [[package]] name = "sha1" -version = "0.10.6" +version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +checksum = "a978451301f4db1d02937a4ab3ccce137717b81826e79b7d49ffe3244a13c3b8" dependencies = [ "cfg-if", "cpufeatures 0.2.17", @@ -4961,7 +4961,7 @@ dependencies = [ "lazy_static", "libc", "ndarray 0.16.1", - "rand 0.8.6", + "rand 0.8.7", "safetensors 0.3.3", "thiserror 1.0.69", "torch-sys", @@ -5098,9 +5098,9 @@ dependencies = [ [[package]] name = "tinyvec" -version = "1.11.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" dependencies = [ "tinyvec_macros", ] @@ -5131,7 +5131,7 @@ dependencies = [ "monostate", "onig", "paste", - "rand 0.9.4", + "rand 0.9.5", "rayon", "rayon-cond", "regex", From 92b5ab7e08669aa759a9d2f42df5167344df1115 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sat, 11 Jul 2026 06:59:21 -0500 Subject: [PATCH 2/6] =?UTF-8?q?feat(hub):=20sha256=20download=20integrity?= =?UTF-8?q?=20=E2=80=94=20stream-verify=20against=20the=20Hub's=20X-Linked?= =?UTF-8?q?-ETag=20(P3)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A redirect-stopped HEAD reads the LFS sha256 from the resolve/ endpoint before the CDN handoff would replace the headers; the GET then stream-hashes every chunk (sha2, SHA-NI), folding a resumed .part prefix in first so a resumed download verifies as one whole file. A wrong-hash .part is deleted, never resumed. FetchOptions::verify_cached re-hashes an existing destination and self-heals exactly once on mismatch (delete + verified refetch); files with no announced sha256 (non-LFS, other hosts) stay length-verified. Real-network proof (tests/real_hub.rs, ignored): fresh 90.8 MB MiniLM fetch stream-verified; clean cache re-verified untouched; a flipped byte mid-file was caught + healed with a full re-stream; a 45.4 MB-seeded resume verified whole. 83 unit tests green (3 new), clippy clean. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 1 + Cargo.toml | 3 + ROADMAP.md | 10 +- crates/mummu/Cargo.toml | 1 + crates/mummu/src/hub.rs | 277 +++++++++++++++++++++++++++++++-- crates/mummu/tests/real_hub.rs | 64 ++++++++ 6 files changed, 343 insertions(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fa86462..fc02467 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3402,6 +3402,7 @@ dependencies = [ "pollster", "serde", "serde_json", + "sha2", "thiserror 2.0.18", "tokenizers", "ureq 3.3.0", diff --git a/Cargo.toml b/Cargo.toml index 242574f..2e8a63f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,9 @@ tokenizers = "0.22" # Blocking HTTP (rustls) for the P3 model downloader; small, sync like the rest # of the library surface. ureq = "3" +# Stream-verifies downloads against the Hub's announced LFS sha256 +# (X-Linked-ETag); SHA-NI accelerated, ~GB/s — negligible next to the network. +sha2 = "0.10" serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/ROADMAP.md b/ROADMAP.md index 23bf7bd..e0340c7 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -129,8 +129,14 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari proof: all-MiniLM (90.8 MB) downloaded → checked-load → unit-norm embedding; a half-seeded `.part` resumed at byte 249,507/466,247 and finished byte-identical. Local paths are already first-class (`load_from_dir`); bundled-resources-dir precedence is app wiring.* -- [ ] Stronger download integrity: verify the Hub's LFS sha256 (`X-Linked-ETag`) instead of length-only; - re-verify on cache hits behind a flag. +- [x] Stronger download integrity: verify the Hub's LFS sha256 (`X-Linked-ETag`) instead of length-only; + re-verify on cache hits behind a flag. *(2026-07-11) Every download stream-hashes (sha2, SHA-NI) + against the sha256 a redirect-stopped HEAD reads from `X-Linked-ETag`; resumes fold the `.part` + prefix into the hash; a mismatched `.part` is deleted, never resumed. `FetchOptions::verify_cached` + re-hashes cache hits and self-heals once (delete + verified refetch). Real-network proof on the + 90.8 MB MiniLM weights: a flipped byte mid-file (invisible to the length check) was caught and + healed; a 45.4 MB-seeded resume re-verified whole. Non-LFS files (no announced sha256) stay + length-verified.* - [x] **safetensors** *(ex-laurelane)* — `burn-store` `SafetensorsStore` + `PyTorchToBurnAdapter`; the primary path. *(2026-07-09) `import::{CastFloatAdapter, load_checked}`: bf16→backend-float cast + fail-loud load; proven by loading the real 3.1 GB Qwen2.5-1.5B and 2.3 GB LFM2.5 checkpoints with zero missing keys.* diff --git a/crates/mummu/Cargo.toml b/crates/mummu/Cargo.toml index 36b3b16..f521e06 100644 --- a/crates/mummu/Cargo.toml +++ b/crates/mummu/Cargo.toml @@ -16,5 +16,6 @@ serde = { workspace = true } serde_json = { workspace = true } thiserror = { workspace = true } ureq = { workspace = true } +sha2 = { workspace = true } once_cell = { workspace = true } burn-flex = { workspace = true } diff --git a/crates/mummu/src/hub.rs b/crates/mummu/src/hub.rs index 3424c40..d32747f 100644 --- a/crates/mummu/src/hub.rs +++ b/crates/mummu/src/hub.rs @@ -1,14 +1,18 @@ //! Model downloads: HuggingFace Hub (or any HTTP host) → the local model //! cache. Streaming, **resumable** (a `.part` picks up where a killed -//! download stopped, via HTTP `Range`), length-verified, and -//! **sharded-checkpoint aware** (`model.safetensors.index.json` → fetch every -//! shard). Progress surfaces through a callback so app settings UIs can show -//! it (P8). Completed files are cache-first: an existing destination is never -//! re-fetched. +//! download stopped, via HTTP `Range`), **integrity-checked** (streamed +//! sha256 against the Hub's announced LFS `X-Linked-ETag`, length as the +//! fallback), and **sharded-checkpoint aware** (`model.safetensors.index.json` +//! → fetch every shard). Progress surfaces through a callback so app settings +//! UIs can show it (P8). Completed files are cache-first: an existing +//! destination is never re-fetched unless [`FetchOptions::verify_cached`] +//! asks for a re-hash. use std::io::{Read, Write}; use std::path::{Path, PathBuf}; +use sha2::{Digest, Sha256}; + /// Streaming copy granularity: big enough to amortize syscalls, small enough /// to keep progress callbacks responsive. const CHUNK_BYTES: usize = 64 * 1024; @@ -36,6 +40,22 @@ pub enum HubError { }, #[error("shard index {path}: {reason}")] BadIndex { path: PathBuf, reason: String }, + #[error("{url}: sha256 mismatch — announced {expected}, computed {computed}")] + Corrupt { + url: String, + expected: String, + computed: String, + }, +} + +/// Options for [`fetch_file_with`] / [`fetch_model_with`]. +#[derive(Debug, Clone, Copy, Default)] +pub struct FetchOptions { + /// Re-hash an already-complete destination against the server's announced + /// sha256 (one extra HEAD request per file); on a mismatch the corrupt + /// copy is deleted and re-fetched once. Off by default: a completed file + /// was already verified as it streamed in. + pub verify_cached: bool, } /// Download progress for one file, reported after every chunk. @@ -95,18 +115,154 @@ fn part_path(dest: &Path) -> PathBuf { PathBuf::from(p) } +/// Lowercase hex of a digest. +fn hex64(bytes: &[u8]) -> String { + use std::fmt::Write as _; + debug_assert_eq!(bytes.len(), 32, "sha256 digests are 32 bytes"); + let mut s = String::with_capacity(bytes.len() * 2); + for b in bytes { + let _ = write!(s, "{b:02x}"); + } + s +} + +/// A quoted 64-hex etag value is a content sha256 (the Hub's `X-Linked-ETag` +/// for LFS files). Git-style etags (40-hex sha1, or non-hex) parse to `None` — +/// they name a revision, not the bytes, and cannot verify a stream. +fn parse_sha256_etag(raw: &str) -> Option { + let v = raw.trim().trim_start_matches("W/").trim_matches('"'); + let is_sha256 = v.len() == 64 && v.bytes().all(|b| b.is_ascii_hexdigit()); + is_sha256.then(|| v.to_ascii_lowercase()) +} + +/// Ask the server for the file's content sha256: a redirect-stopped HEAD reads +/// the Hub's `X-Linked-ETag` from the `resolve/` endpoint itself, before the +/// CDN handoff would replace the headers. `Ok(None)` when nothing usable is +/// announced (non-LFS files, other hosts, HEAD rejected) — those downloads +/// stay length-verified only. Transport failures are loud: the GET would fail +/// the same way. +fn announced_sha256(url: &str) -> Result, HubError> { + assert!(url.starts_with("https://"), "refusing non-https url: {url}"); + let agent: ureq::Agent = ureq::Agent::config_builder() + .max_redirects(0) + .http_status_as_error(false) + .build() + .into(); + let resp = agent.head(url).call().map_err(|e| HubError::Http { + url: url.into(), + reason: e.to_string(), + })?; + Ok(["x-linked-etag", "etag"].iter().find_map(|h| { + resp.headers() + .get(*h) + .and_then(|v| v.to_str().ok()) + .and_then(parse_sha256_etag) + })) +} + +/// Streaming sha256 of a file on disk, as lowercase hex. +fn sha256_hex_of_file(path: &Path) -> Result { + let io_err = |e: std::io::Error| HubError::Io { + path: path.to_path_buf(), + reason: e.to_string(), + }; + let mut f = std::fs::File::open(path).map_err(io_err)?; + let mut hasher = Sha256::new(); + let mut buf = vec![0u8; CHUNK_BYTES]; + let mut hashed = 0u64; + loop { + let n = f.read(&mut buf).map_err(io_err)?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + hashed += n as u64; + assert!( + hashed <= MAX_FILE_BYTES, + "{path:?}: exceeds the file bound while hashing" + ); + } + Ok(hex64(&hasher.finalize())) +} + +/// Feed the already-downloaded `.part` prefix into the stream hasher so a +/// resumed download still verifies as one whole file. +fn hash_part_prefix(part: &Path, resume_from: u64, hasher: &mut Sha256) -> Result<(), HubError> { + assert!(resume_from > 0, "no prefix to hash"); + let mut f = std::fs::File::open(part).map_err(|e| HubError::Io { + path: part.to_path_buf(), + reason: e.to_string(), + })?; + let mut remaining = resume_from; + let mut buf = vec![0u8; CHUNK_BYTES]; + while remaining > 0 { + let want = remaining.min(CHUNK_BYTES as u64) as usize; + let n = f.read(&mut buf[..want]).map_err(|e| HubError::Io { + path: part.to_path_buf(), + reason: e.to_string(), + })?; + // The prefix length came from this file's own metadata an instant ago. + assert!(n > 0, "{part:?}: prefix ended {remaining} bytes early"); + hasher.update(&buf[..n]); + remaining -= n as u64; + } + Ok(()) +} + /// Fetch `url` into `dest`, streaming through `.part` and resuming any /// earlier partial download. No-op when `dest` already exists (cache-first). -/// `on_progress` fires after every chunk with cumulative counts. +/// `on_progress` fires after every chunk with cumulative counts. The stream +/// is verified against the server's announced sha256 when there is one +/// (Hub LFS files), else by length. pub fn fetch_file( url: &str, dest: &Path, + on_progress: impl FnMut(Progress<'_>), +) -> Result<(), HubError> { + fetch_file_with(url, dest, FetchOptions::default(), on_progress) +} + +/// [`fetch_file`] with explicit [`FetchOptions`]. With `verify_cached`, an +/// existing `dest` is re-hashed against the announced sha256; a mismatch +/// deletes the corrupt copy (and any stale `.part` that would poison a +/// resume) and re-fetches once — self-healing, never silent. +pub fn fetch_file_with( + url: &str, + dest: &Path, + opts: FetchOptions, mut on_progress: impl FnMut(Progress<'_>), ) -> Result<(), HubError> { assert!(url.starts_with("https://"), "refusing non-https url: {url}"); if dest.exists() { - return Ok(()); // cache hit — never re-fetch a completed file + if !opts.verify_cached { + return Ok(()); // cache hit — never re-fetch a completed file + } + let Some(expected) = announced_sha256(url)? else { + return Ok(()); // nothing announced — nothing to re-verify against + }; + if sha256_hex_of_file(dest)? == expected { + return Ok(()); + } + for stale in [dest.to_path_buf(), part_path(dest)] { + if stale.exists() { + std::fs::remove_file(&stale).map_err(|e| HubError::Io { + path: stale.clone(), + reason: e.to_string(), + })?; + } + } } + download(url, dest, &mut on_progress) +} + +/// The streaming GET behind [`fetch_file_with`]: resume, hash, length-check, +/// then atomically rename `.part` → `dest`. +fn download( + url: &str, + dest: &Path, + on_progress: &mut impl FnMut(Progress<'_>), +) -> Result<(), HubError> { + debug_assert!(!dest.exists(), "download() requires a vacant destination"); let file_label = dest .file_name() .map(|n| n.to_string_lossy().into_owned()) @@ -121,6 +277,9 @@ pub fn fetch_file( let part = part_path(dest); let resume_from = std::fs::metadata(&part).map(|m| m.len()).unwrap_or(0); + // One cheap HEAD up front: with an announced sha256 the whole stream + // (resumed prefix included) is verified; without one, length still is. + let expected_sha = announced_sha256(url)?; let mut req = ureq::get(url); if resume_from > 0 { @@ -147,6 +306,11 @@ pub fn fetch_file( ); } + let mut hasher = expected_sha.as_ref().map(|_| Sha256::new()); + if resumed && let Some(h) = hasher.as_mut() { + hash_part_prefix(&part, resume_from, h)?; + } + let mut out = std::fs::OpenOptions::new() .create(true) .append(resumed) @@ -173,6 +337,9 @@ pub fn fetch_file( path: part.clone(), reason: e.to_string(), })?; + if let Some(h) = hasher.as_mut() { + h.update(&buf[..n]); + } received += n as u64; assert!( received <= MAX_FILE_BYTES, @@ -196,6 +363,21 @@ pub fn fetch_file( received, }); } + if let (Some(expected), Some(h)) = (expected_sha, hasher) { + let computed = hex64(&h.finalize()); + if computed != expected { + // A wrong-hash .part must never seed a resume — drop it. + std::fs::remove_file(&part).map_err(|e| HubError::Io { + path: part.clone(), + reason: e.to_string(), + })?; + return Err(HubError::Corrupt { + url: url.into(), + expected, + computed, + }); + } + } std::fs::rename(&part, dest).map_err(|e| HubError::Io { path: dest.to_path_buf(), reason: e.to_string(), @@ -210,19 +392,39 @@ pub fn fetch_model( repo: &str, revision: &str, dest_dir: &Path, + on_progress: impl FnMut(Progress<'_>), +) -> Result { + fetch_model_with( + repo, + revision, + dest_dir, + FetchOptions::default(), + on_progress, + ) +} + +/// [`fetch_model`] with explicit [`FetchOptions`] (e.g. re-verify cached +/// files' sha256 before trusting them). +pub fn fetch_model_with( + repo: &str, + revision: &str, + dest_dir: &Path, + opts: FetchOptions, mut on_progress: impl FnMut(Progress<'_>), ) -> Result { for file in ["config.json", "tokenizer.json"] { - fetch_file( + fetch_file_with( &hub_file_url(repo, revision, file), &dest_dir.join(file), + opts, &mut on_progress, )?; } // Single-file first (the common case for the small-model tiers we target). - let single = fetch_file( + let single = fetch_file_with( &hub_file_url(repo, revision, "model.safetensors"), &dest_dir.join("model.safetensors"), + opts, &mut on_progress, ); if single.is_ok() { @@ -232,9 +434,10 @@ pub fn fetch_model( // the original single-file error (the more useful signal). let index_name = "model.safetensors.index.json"; let index_dest = dest_dir.join(index_name); - if fetch_file( + if fetch_file_with( &hub_file_url(repo, revision, index_name), &index_dest, + opts, &mut on_progress, ) .is_err() @@ -246,9 +449,10 @@ pub fn fetch_model( reason: e.to_string(), })?; for shard in shards_from_index(&index_bytes, &index_dest)? { - fetch_file( + fetch_file_with( &hub_file_url(repo, revision, &shard), &dest_dir.join(&shard), + opts, &mut on_progress, )?; } @@ -295,6 +499,57 @@ mod tests { assert!(matches!(err, Err(HubError::BadIndex { .. }))); } + #[test] + fn sha256_etag_parsing_accepts_only_content_hashes() { + let sha = "a".repeat(64); + // The Hub quotes LFS etags; weak etags carry a W/ prefix. + assert_eq!(parse_sha256_etag(&format!("\"{sha}\"")), Some(sha.clone())); + assert_eq!( + parse_sha256_etag(&format!("W/\"{sha}\"")), + Some(sha.clone()) + ); + assert_eq!(parse_sha256_etag(&sha.to_uppercase()), Some(sha)); + // Git-style 40-hex sha1, non-hex, and empty values are not sha256s. + assert_eq!(parse_sha256_etag(&format!("\"{}\"", "b".repeat(40))), None); + assert_eq!(parse_sha256_etag(&format!("\"{}\"", "z".repeat(64))), None); + assert_eq!(parse_sha256_etag(""), None); + } + + #[test] + fn file_hash_matches_the_reference_vector() { + // FIPS 180-2 test vector: sha256("abc"). + let dir = std::env::temp_dir().join("mummu-hub-test-sha"); + std::fs::create_dir_all(&dir).unwrap(); + let p = dir.join("abc.txt"); + std::fs::write(&p, b"abc").unwrap(); + assert_eq!( + sha256_hex_of_file(&p).unwrap(), + "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad" + ); + std::fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn prefix_plus_remainder_hash_equals_whole_file_hash() { + // The resume path hashes the .part prefix, then the streamed tail; + // together they must equal one pass over the whole file. + let dir = std::env::temp_dir().join("mummu-hub-test-prefix"); + std::fs::create_dir_all(&dir).unwrap(); + let whole: Vec = (0u32..200_000).map(|i| (i % 251) as u8).collect(); + let p = dir.join("whole.bin"); + std::fs::write(&p, &whole).unwrap(); + let reference = sha256_hex_of_file(&p).unwrap(); + + let split = whole.len() / 3; + let part = dir.join("whole.bin.part"); + std::fs::write(&part, &whole[..split]).unwrap(); + let mut h = Sha256::new(); + hash_part_prefix(&part, split as u64, &mut h).unwrap(); + h.update(&whole[split..]); + assert_eq!(hex64(&h.finalize()), reference); + std::fs::remove_dir_all(&dir).ok(); + } + #[test] fn part_path_appends_suffix() { assert_eq!( diff --git a/crates/mummu/tests/real_hub.rs b/crates/mummu/tests/real_hub.rs index f6599e1..daedb2c 100644 --- a/crates/mummu/tests/real_hub.rs +++ b/crates/mummu/tests/real_hub.rs @@ -114,6 +114,70 @@ fn hub_resume_completes_a_partial_download_byte_identical() { ); } +#[test] +#[ignore = "needs network (MUMMU_HUB_DEST names the download dir)"] +fn hub_sha256_verification_catches_and_heals_a_corrupt_cache() { + let Some(dest) = std::env::var_os("MUMMU_HUB_DEST").map(PathBuf::from) else { + panic!("set MUMMU_HUB_DEST to a scratch dir for the ~90 MB download"); + }; + let dir = dest.join("verify-proof"); + let url = hub::hub_file_url(REPO, "main", "model.safetensors"); + let weights = dir.join("model.safetensors"); + let verify = hub::FetchOptions { + verify_cached: true, + }; + + // Fresh (or cached) download; the stream itself hash-verifies against the + // Hub's X-Linked-ETag, so success here is already an integrity proof. + hub::fetch_file(&url, &weights, |_| {}).expect("verified fetch"); + let healthy_len = weights.metadata().expect("weights exist").len(); + + // A clean cache re-verifies without being touched. + let modified_before = weights.metadata().unwrap().modified().unwrap(); + hub::fetch_file_with(&url, &weights, verify, |_| {}).expect("clean cache re-verifies"); + assert_eq!( + weights.metadata().unwrap().modified().unwrap(), + modified_before, + "a matching cached file must not be rewritten" + ); + + // Corrupt one byte mid-file (length unchanged — only the hash can see it), + // then watch verify_cached self-heal by re-downloading. + let mut bytes = std::fs::read(&weights).expect("read weights"); + let mid = bytes.len() / 2; + bytes[mid] ^= 0xFF; + std::fs::write(&weights, &bytes).expect("plant corruption"); + let mut streamed = 0u64; + hub::fetch_file_with(&url, &weights, verify, |p| streamed = p.received_bytes) + .expect("corrupt cache heals"); + assert_eq!( + weights.metadata().unwrap().len(), + healthy_len, + "healed file has the healthy length" + ); + assert_eq!( + streamed, healthy_len, + "healing must have re-streamed the whole file" + ); + eprintln!( + "[real_hub] sha256 gate: clean cache untouched; flipped byte at {mid} was caught and healed ({streamed} bytes re-streamed)" + ); + + // Resumed LFS download: the .part prefix must fold into the stream hash + // (a 206 + hash_part_prefix pass over a real X-Linked-ETag file). + let healthy = std::fs::read(&weights).expect("healthy bytes"); + std::fs::remove_file(&weights).expect("drop completed file"); + let half = healthy.len() / 2; + std::fs::write(dir.join("model.safetensors.part"), &healthy[..half]).expect("seed part"); + hub::fetch_file(&url, &weights, |_| {}).expect("resumed fetch hash-verifies"); + assert_eq!( + weights.metadata().unwrap().len(), + healthy_len, + "resumed file has the healthy length" + ); + eprintln!("[real_hub] sha256 gate: resume from {half} bytes re-verified the whole file"); +} + /// One-shot helper the nightly uses to pull the CPU-tier Qwen into a cache /// dir (also a second real proof of the sharded/single-file fetch path on a /// 1 GB checkpoint). From 1b465627d3349733689e7923df4d95e9e4f62da2 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sat, 11 Jul 2026 07:18:32 -0500 Subject: [PATCH 3/6] =?UTF-8?q?feat(p6):=20f16=20mixed-precision=20island?= =?UTF-8?q?=20=E2=80=94=20coherent=20f16=20inference=20at=20half=20the=20V?= =?UTF-8?q?RAM?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One island sufficed: Qwen-class q·k^T attention scores overflow f16's 65504 max, collapsing softmax to NaN. Scores + mask + softmax now compute in f32 (per-tensor cast, Burn 0.21 multi-dtype — llama.cpp pins the same matmul to f32), and the probabilities return to the ambient dtype for the value matmul. Burn's RmsNorm already reduces in f32 upstream; the logit path needed nothing. Verified on the 4070 Ti SUPER: - f16 gate (tests/real_f16.rs) passes all 3 claims: no crash, 6.75 GiB whole-card / ~3.6 GiB runner (vs ~7.9 GiB f32), coherent greedy output. - Parity gate re-passed both legs on f32, where the casts are no-ops: top-5 logits vs Candle max |dlogit| 2.670e-5 (unchanged), 24-token greedy byte-identical to ollama qwen2.5:1.5b-instruct-fp16. - Benches (criterion, new gpu-f16 group): f16 88.0 ms TTFT / 14.1 tok/s — speed parity with f32, VRAM halved. f32 improved: TTFT 100.5 -> 88.4 ms, decode 13.3 -> 14.1 tok/s; budget gates re-passed (GPU 96.8 ms / 10.2 tok/s, CPU 8.66 tok/s). Confirms decode is dispatch-bound, not bandwidth-bound (f16 == f32 speed); SPIR-V is the remaining lever. Co-Authored-By: Claude Opus 4.8 --- README.md | 13 +++++--- ROADMAP.md | 37 ++++++++++++++++------- bench/BASELINE.md | 34 ++++++++++++++------- crates/mummu-bench/benches/runner.rs | 44 ++++++++++++++++++---------- crates/mummu/src/nn/attention.rs | 18 +++++++++--- 5 files changed, 101 insertions(+), 45 deletions(-) diff --git a/README.md b/README.md index f8df458..b396478 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys - **One binary, every device** — compile both `Wgpu` (Vulkan/DX12/Metal, no CUDA toolchain) and `burn-flex` (CPU); a runtime probe enumerates **all** adapters + the CPU and places the model across them — a single GPU, **several GPUs together**, or GPU + CPU hybrid. No feature-split builds, no per-vendor path. - **Models from scratch, generic over `B: Backend`** — a growing zoo (Qwen2/2.5, LFM2/2.5 hybrid conv+attention, all-MiniLM embedder) built on shared blocks (RmsNorm · GQA · RoPE · SwiGLU · tied lm-head · depthwise causal conv), with a clean trait to add more. - **Trustworthy reimplementations** — every port must pass a **parity gate**: single-forward top-k logits *and* a short greedy sequence match a reference (Candle, or a local Ollama of the same model) exactly. -- **Fast** — per-layer KV cache (+ conv-state cache for hybrids), on-GPU argmax (sync only the winning index), sampling, **token streaming**, cooperative cancellation; kernel `fusion` + `autotune`; an opt-in **f16** path to roughly halve VRAM. +- **Fast** — per-layer KV cache (+ conv-state cache for hybrids), on-GPU argmax (sync only the winning index), sampling, **token streaming**, cooperative cancellation; kernel `fusion` + `autotune`; an **f16** path (f32 attention-score island for numeric safety) that halves VRAM at full speed. - **A full model-import suite** — pull a model from HuggingFace (by repo id) or from disk and load it: **safetensors**, **PyTorch** state dicts, and **GGUF** (llama.cpp, dequantized) weights; `config.json`-driven hyperparameters; tokenizer + chat-template import (HF `tokenizers` / SentencePiece / BPE); per-architecture weight-name remapping with a **checked load** (fail loudly on a key mismatch, never silently zero-init); resumable, shard-aware downloads into a per-user cache; and a declarative **model registry** so adding a model is a manifest entry, not new code. - **Quantize to fit, fill the hardware** — a planner probes every GPU + the CPU (VRAM / RAM), then imports or **quantizes on the fly** (GGUF K-quants, GPTQ / AWQ, or Burn's own int8/int4) and chooses precision + **layer placement** so the *largest model that fits* runs and every device is used — sharded across GPUs, spilling cold layers to CPU when needed. Plus a model-management API (download progress, disk usage, remove) apps surface in their settings UI. - **Local embeddings** — a from-scratch MiniLM-class sentence embedder (CPU) for fully-offline semantic search. @@ -38,9 +38,14 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys - **Sampling, streaming, cancellation** — temperature / top-k / top-p sampling (deterministic per seed), per-token streaming through a `ControlFlow` callback, and cooperative between-token cancellation; greedy decoding keeps the argmax on-device. -- **Benchmarked** — Qwen2.5-1.5B f32 on the reference GPU: **TTFT 100.5 ms, decode 13.3 tok/s, 11.9 GiB - whole-card peak** (~7.9 GiB runner) — recorded with budgets in [bench/BASELINE.md](bench/BASELINE.md), - enforced by an opt-in regression gate (`mummu-bench/tests/budget.rs`). +- **f16 inference, validated** — Qwen2.5-1.5B runs coherently on `GpuF16` (weights + KV in f16, the + q·kᵀ attention scores + softmax computed in an f32 island to stop f16 overflow): **~3.6 GiB runner + VRAM vs ~7.9 GiB f32, at identical speed** (14.1 tok/s / 88 ms TTFT); the parity gate re-passes + unchanged on f32, where the island casts are no-ops ([bench/BASELINE.md](bench/BASELINE.md)). +- **Benchmarked** — Qwen2.5-1.5B on the reference GPU: **TTFT 88.4 ms, decode 14.1 tok/s** (f32, 11.9 GiB + whole-card peak ≈ 7.9 GiB runner; f16: 88.0 ms, 14.1 tok/s, 6.75 GiB ≈ 3.6 GiB runner) — recorded with + budgets in [bench/BASELINE.md](bench/BASELINE.md), enforced by an opt-in regression gate + (`mummu-bench/tests/budget.rs`). - **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`. diff --git a/ROADMAP.md b/ROADMAP.md index e0340c7..100ca37 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -44,7 +44,9 @@ a benchmark holds/improves its budget; README perf claims link an artifact. weights + GPU) passing at 110.4 ms / 11.8 tok/s.* - [ ] Decode is dispatch-bound, not bandwidth-bound: 75 ms/token streams ~6.2 GB of f32 weights at only ~83 GB/s vs the 4070 Ti SUPER's ~672 GB/s — the SPIR-V compiler feature (P6 item) and f16 are the - levers to chase; re-baseline after each. + levers to chase; re-baseline after each. *(2026-07-11) Confirmed empirically: f16 (half the weight + traffic) decodes at exactly f32's speed — 70.9 vs 70.7 ms/token — so bandwidth isn't the limiter; + SPIR-V (TensorCores at f16) is the remaining lever.* - [ ] Evaluate Burn 0.21's `burn.toml` project config — per-subsystem tuning + a CubeCL kernel-validation layer without recompiling; useful as a debug switch for kernel-level parity hunts — https://burn.dev/blog/release-0.21.0/ @@ -149,6 +151,11 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari sub-scales + fp16 d (210 B). Rust references: llama.cpp ggml-quants + the `rage-quant` crate (Q8_0/Q4_K/Q6_K dequant + SIMD dot) — https://haroldbenoit.com/notes/ml/llms/quantization/llama.cpp/k-quants-implementation · https://crates.io/crates/rage-quant + *(2026-07-11 research)* `pmetal-gguf` 0.5 (May 2026, MIT/Apache-2.0, standalone — no candle/burn + deps) is the most complete Rust GGUF implementation yet: read/write + dequant for the K-quants + (Q2K–Q8K) AND IQ-quants, SIMD-optimized, importance-matrix support, HF-compatible config + generation — evaluate as dependency-or-reference before hand-porting ggml-quants — + https://docs.rs/pmetal-gguf/latest/pmetal_gguf/ - [ ] **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. @@ -218,17 +225,21 @@ that fits the model AND uses every device to the fullest. - [ ] **Device inventory** — enumerate every GPU (`wgpu` adapters: name, backend, VRAM) and the CPU (cores, RAM); a stable device set cached at startup, reported so the apps can show it in settings. - [ ] **Precision selection** — pick a per-device dtype (f32 / **f16** / int8 / int4) that fits: f16 via - `Wgpu` (needs wgpu ≥ 27 `SHADER_F16` polyfill — *laurelane compiles it + a startup - `SHADER_F16` diagnostic; finish on-GPU validation here*: no naga crash, ~halved VRAM, coherent output); - drop to int8/int4 (P9) when f16 still won't fit. *(2026-07-10) On-GPU validation ran - (`tests/real_f16.rs`, the standing gate): **2 of 3 claims hold** — shaders compile + run on - Vulkan/SHADER_F16, VRAM drops 11.9 → 8.7 GiB whole-card (~7.9 → ~4.7 GiB runner), but logits - collapse to NaN (the GPU argmax returns the out-of-vocab sentinel 151936 = vocab_size; now caught - loudly by a decode guard). Coherent-output remains open below.* -- [ ] **f16 mixed-precision islands** — Qwen2.5-1.5B in pure f16 NaNs out (overflow in the + `Wgpu`; drop to int8/int4 (P9) when f16 still won't fit. *(2026-07-11) The f16 + backend itself is now **fully validated** (all 3 claims — see the islands item below); what remains + here is the *picking* logic, which rides the placement-plan item + P9.* +- [x] **f16 mixed-precision islands** — Qwen2.5-1.5B in pure f16 NaNs out (overflow in the softmax/RmsNorm/logit reductions; f16 max is 65 504). Keep weights + matmuls f16 but compute the numerically hot reductions (attention softmax, RmsNorm accumulation, final logits) in f32, then - re-run the f16 gate and the parity harness. + re-run the f16 gate and the parity harness. *(2026-07-11) **One island sufficed**: the q·kᵀ + attention scores overflow f16 — scores + mask + softmax now compute in f32 (per-tensor + `cast(DType::F32)`, Burn 0.21 multi-dtype; llama.cpp pins the same matmul to f32), probs return to + the ambient dtype; Burn's RmsNorm already reduces in f32 upstream, and the logit path needed + nothing. The f16 gate passes all 3 claims: no crash, **6.75 GiB whole-card / ~3.6 GiB runner** + (vs ~7.9 GiB f32), coherent greedy output ("2+2 equals 4."). Casts are no-ops on f32: the parity + gate re-passed both legs (max |Δlogit| 2.670e-5, unchanged; Ollama greedy byte-identical), and f32 + perf *improved* (TTFT 100.5 → 88.4 ms, decode 13.3 → 14.1 tok/s). f16 benches recorded in + `bench/BASELINE.md`: 88.0 ms TTFT, 14.1 tok/s — speed parity with f32, VRAM halved.* - [ ] Evaluate burn-wgpu's **`spirv` compiler feature** on Vulkan (CubeCL SPIR-V backend instead of WGSL/naga): claims significantly faster matmul incl. TensorCores at f16 — could be the cheapest decode-tok/s lever on the dev GPU; gate on the parity harness + `bench/BASELINE.md` — @@ -260,7 +271,11 @@ that fits the model AND uses every device to the fullest. `crates/mummu/tests/fixtures/`; fp16 Ollama tag pulled and validated.* - [ ] LFM2.5 same-weights reference for the parity gate: no Candle port exists — candidate routes are llama.cpp `logprobs` on the fp16 GGUF, or a one-shot HF `transformers` logits dump matched to the - safetensors revision. + safetensors revision. *(2026-07-11 research)* Liquid officially documents running LFM2.5-1.2B + GGUFs under `llama-server`; its completion API returns per-token top logprobs via `n_probs` + (temperature 0 for the greedy leg) — that plus an fp16 GGUF of the same revision is a workable + logits leg without Python — https://docs.liquid.ai/deployment/on-device/llama-cpp · + https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md - [ ] Wire the perf suite (above) into the parity harness so a correctness *or* budget regression fails CI. ### P8 — Model management API diff --git a/bench/BASELINE.md b/bench/BASELINE.md index 8440f32..8079d4a 100644 --- a/bench/BASELINE.md +++ b/bench/BASELINE.md @@ -1,30 +1,42 @@ # Benchmark baseline & budgets -Reference machine: Ryzen 9 7950X3D · 128 GB · **RTX 4070 Ti SUPER 16 GB** (wgpu/Vulkan, `Fusion`, -f32). Bench: `MUMMU_QWEN2_DIR= cargo bench -p mummu-bench` (criterion, `benches/runner.rs`; +Reference machine: Ryzen 9 7950X3D · 128 GB · **RTX 4070 Ti SUPER 16 GB** (wgpu/Vulkan, `Fusion`). +Bench: `MUMMU_QWEN2_DIR= cargo bench -p mummu-bench` (criterion, `benches/runner.rs`; fixed ~36-token ChatML prompt). A change that pushes a budget over its ceiling does not ship; update the recorded numbers (and this file's date) only on a legitimate improvement. ## Qwen2.5-1.5B-Instruct · single GPU · f32 -| Metric | Recorded (2026-07-10) | Budget | +| Metric | Recorded (2026-07-11) | Budget | | --- | --- | --- | -| TTFT (fresh cache: full prefill + first token) | 100.5 ms | ≤ 150 ms | -| Decode latency (32 greedy tokens, warm KV cache) | 2.414 s → **75.4 ms/token ≈ 13.3 tok/s** | ≥ 10 tok/s | +| TTFT (fresh cache: full prefill + first token) | 88.4 ms | ≤ 150 ms | +| Decode latency (32 greedy tokens, warm KV cache) | 2.263 s → **70.7 ms/token ≈ 14.1 tok/s** | ≥ 10 tok/s | | Peak GPU memory during decode (whole card, ~4.0 GiB desktop ambient → ~7.9 GiB runner) | 11.9 GiB | ≤ 13 GiB whole-card | +## Qwen2.5-1.5B-Instruct · single GPU · **f16** (weights + KV; f32 attention-score island) + +| Metric | Recorded (2026-07-11) | Budget | +| --- | --- | --- | +| TTFT (fresh cache: full prefill + first token) | 88.0 ms | ≤ 150 ms | +| Decode latency (32 greedy tokens, warm KV cache) | 2.270 s → **70.9 ms/token ≈ 14.1 tok/s** | ≥ 10 tok/s | +| Peak GPU memory during decode (whole card, 3.1 GiB ambient → **~3.6 GiB runner**) | 6.75 GiB | ≤ 8 GiB whole-card | + +f16 decode speed matches f32 — the WGSL decode path is dispatch-bound, not bandwidth-bound (see Notes), +so halved weight traffic buys nothing yet; the win is VRAM (**~7.9 → ~3.6 GiB runner**, room for larger +models/contexts). The SPIR-V compiler feature (ROADMAP P6) is the identified speed lever for both dtypes. + ## Qwen2.5-0.5B-Instruct · CPU (burn-flex) · f32 | Metric | Recorded (2026-07-10) | Budget | | --- | --- | --- | | Decode (8 greedy tokens, warm KV cache; `mummu-bench/tests/budget_cpu.rs`) | **11.7 tok/s** | ≥ 6 tok/s | -Datapoint (not yet a budget): the same model on `GpuF16` peaks at **8.7 GiB whole-card (~4.7 GiB -runner)** — VRAM roughly halves as expected — but decodes NaN today (see the ROADMAP P6 -mixed-precision-islands item), so no f16 perf row exists yet. - Notes -- Effective weight-streaming bandwidth at 75 ms/token over ~6.2 GB of f32 weights is ~83 GB/s vs the +- 2026-07-11: the f32 attention-score island (NaN fix for f16) coincided with an f32 *improvement* + (TTFT 100.5 → 88.4 ms, decode 13.3 → 14.1 tok/s) — softmax now always runs in f32 with fusion + re-tuning around it; both budget gates re-passed (`budget.rs` 96.8 ms / 10.2 tok/s, `budget_cpu.rs` + 8.66 tok/s measured concurrently with the GPU bench). +- Effective weight-streaming bandwidth at ~71 ms/token over ~6.2 GB of f32 weights is ~88 GB/s vs the card's ~672 GB/s — the decode path is kernel/dispatch-bound, not bandwidth-bound. The CubeCL SPIR-V - compiler feature and the f16 path (ROADMAP P6) are the identified levers. + compiler feature (ROADMAP P6) is the identified lever. - `harness_smoke` (sub-ns) exists only to keep `cargo bench` green without the multi-GB weights. diff --git a/crates/mummu-bench/benches/runner.rs b/crates/mummu-bench/benches/runner.rs index fd53000..75ef992 100644 --- a/crates/mummu-bench/benches/runner.rs +++ b/crates/mummu-bench/benches/runner.rs @@ -1,6 +1,6 @@ //! Real-model benchmarks: TTFT and decode tok/s for Qwen2.5-1.5B on the -//! machine's default GPU. Budgets and the last recorded numbers live in -//! `bench/BASELINE.md`. Run with +//! machine's default GPU, in f32 and (where `SHADER_F16` exists) f16. +//! Budgets and the last recorded numbers live in `bench/BASELINE.md`. Run with //! //! ```text //! MUMMU_QWEN2_DIR=path/to/qwen2.5-1.5b cargo bench -p mummu-bench @@ -10,11 +10,12 @@ //! so `cargo bench` stays green on machines without the multi-GB checkpoint. use std::hint::black_box; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::time::{Duration, Instant}; +use burn::tensor::backend::Backend; use criterion::{Criterion, criterion_group, criterion_main}; -use mummu::backend::Gpu; +use mummu::backend::{Gpu, GpuF16, inventory}; use mummu::decode::argmax_id; use mummu::models::CausalLm; use mummu::models::qwen2::{self, LoadedQwen2}; @@ -30,7 +31,7 @@ fn qwen2_dir() -> Option { } /// The benchmark prompt — fixed so numbers are comparable across runs. -fn prompt_ids(dir: &std::path::Path) -> Vec { +fn prompt_ids(dir: &Path) -> Vec { let text = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nExplain, in three sentences, why the sky is blue.<|im_end|>\n<|im_start|>assistant\n"; let tok = Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer.json loads"); let ids = tok.encode(text, false).expect("encodes").get_ids().to_vec(); @@ -44,19 +45,19 @@ fn bench_harness_smoke(c: &mut Criterion) { }); } -fn bench_qwen2_real(c: &mut Criterion) { - let Some(dir) = qwen2_dir() else { - eprintln!("[mummu-bench] MUMMU_QWEN2_DIR not set — skipping real-model benches"); - return; - }; - let device = burn::tensor::Device::::default(); - let loaded: LoadedQwen2 = - qwen2::load_from_dir(&dir, &device).expect("weights load checked"); - let ids = prompt_ids(&dir); +/// TTFT + decode for one backend; the model drops (and its VRAM frees) when +/// this returns, so backends bench sequentially without stacking checkpoints. +fn bench_qwen2_on(c: &mut Criterion, dir: &Path, group_name: &str) +where + B::Device: Default, +{ + let device = B::Device::default(); + let loaded: LoadedQwen2 = qwen2::load_from_dir(dir, &device).expect("weights load checked"); + let ids = prompt_ids(dir); // TTFT: fresh cache, full prefill, first token argmax (the id readback is // the GPU sync point, so the measured span covers real work end-to-end). - let mut group = c.benchmark_group("qwen2.5-1.5b/gpu"); + let mut group = c.benchmark_group(group_name); group.sample_size(10); group.bench_function("ttft_prefill_first_token", |b| { b.iter(|| { @@ -89,5 +90,18 @@ fn bench_qwen2_real(c: &mut Criterion) { group.finish(); } +fn bench_qwen2_real(c: &mut Criterion) { + let Some(dir) = qwen2_dir() else { + eprintln!("[mummu-bench] MUMMU_QWEN2_DIR not set — skipping real-model benches"); + return; + }; + bench_qwen2_on::(c, &dir, "qwen2.5-1.5b/gpu"); + if inventory().any_shader_f16() { + bench_qwen2_on::(c, &dir, "qwen2.5-1.5b/gpu-f16"); + } else { + eprintln!("[mummu-bench] no SHADER_F16 adapter — skipping the f16 benches"); + } +} + criterion_group!(benches, bench_harness_smoke, bench_qwen2_real); criterion_main!(benches); diff --git a/crates/mummu/src/nn/attention.rs b/crates/mummu/src/nn/attention.rs index 8eb910d..1ea481c 100644 --- a/crates/mummu/src/nn/attention.rs +++ b/crates/mummu/src/nn/attention.rs @@ -4,7 +4,7 @@ use burn::module::Module; use burn::nn::{Linear, LinearConfig, RmsNorm, RmsNormConfig}; -use burn::tensor::{Tensor, TensorData, activation, backend::Backend}; +use burn::tensor::{DType, Tensor, TensorData, activation, backend::Backend}; use super::MAX_CONTEXT_TOKENS; use super::rope::apply_rope; @@ -185,12 +185,22 @@ impl GqaAttention { let k = repeat_kv(k_all, group); let v = repeat_kv(v_all, group); + // f32 island: Qwen-class attention logits overflow f16 (max 65504) + // in the q·kᵀ scores, collapsing softmax to NaN — llama.cpp pins this + // same matmul to f32 precision for the same reason. Scores + mask + + // softmax run in f32, the probabilities (all in [0, 1]) return to the + // ambient dtype for the value matmul. Every cast is a no-op on f32 + // backends. + let ambient = q.dtype(); let scale = 1.0 / (hd as f32).sqrt(); - let mut scores = q.matmul(k.swap_dims(2, 3)).mul_scalar(scale); + let mut scores = q + .cast(DType::F32) + .matmul(k.cast(DType::F32).swap_dims(2, 3)) + .mul_scalar(scale); if let Some(m) = mask { - scores = scores.add(m.clone()); + scores = scores.add(m.clone().cast(DType::F32)); } - let probs = activation::softmax(scores, 3); + let probs = activation::softmax(scores, 3).cast(ambient); let ctx = probs.matmul(v).swap_dims(1, 2).reshape([b, t, nh * hd]); self.o_proj.forward(ctx) } From c22d25729f1718787862426793fb62b60ae3c046 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sat, 11 Jul 2026 07:27:42 -0500 Subject: [PATCH 4/6] =?UTF-8?q?feat(p4):=20Hermes-style=20function=20calli?= =?UTF-8?q?ng=20=E2=80=94=20tools=20render,=20tool-call=20parse,=20real-GP?= =?UTF-8?q?U=20proof?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit chat gains the tool-use half of the Qwen2.5/Qwen3 chat template: - ToolSpec + render_with_tools: advertises function signatures in the system turn's '# Tools' section, byte-matching the trained template wording ({"type":"function","function":{...}} per line in ). - Turn::assistant_tool_calls re-renders emitted calls faithfully into history; Turn::tool_response results merge consecutive tool turns into ONE user turn of blocks, exactly as the template does. - parse_tool_calls extracts every block (bounded at 64) plus the surrounding prose, with a loud error taxonomy (Unclosed, BadJson, TooMany); plain text parses as (no calls, text), not an error. REAL-GPU proof (tests/real_toolcall.rs, ignored): Qwen2.5-1.5B, given a rendered get_weather prompt, greedy-emitted a well-formed {"name": "get_weather", "arguments": {"city": "Paris"}} and the parser round-tripped it. 90 unit tests green (10 new), clippy clean. LFM2.5's bracket notation stays a P4 [ ] item. Co-Authored-By: Claude Opus 4.8 --- README.md | 5 + ROADMAP.md | 16 +- crates/mummu/src/chat.rs | 320 +++++++++++++++++++++++++++- crates/mummu/tests/real_toolcall.rs | 77 +++++++ 4 files changed, 407 insertions(+), 11 deletions(-) create mode 100644 crates/mummu/tests/real_toolcall.rs diff --git a/README.md b/README.md index b396478..3b53bb1 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,11 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys - **Sampling, streaming, cancellation** — temperature / top-k / top-p sampling (deterministic per seed), per-token streaming through a `ControlFlow` callback, and cooperative between-token cancellation; greedy decoding keeps the argmax on-device. +- **Function calling (Hermes-style)** — advertise `ToolSpec`s through `render_with_tools` (the exact + `# Tools`/`` template Qwen2.5/Qwen3 are trained on), feed results back as merged + `` turns, and extract calls with a bounded `parse_tool_calls`; proven end-to-end on + the real GPU (Qwen2.5-1.5B emitted a parseable `get_weather({"city": "Paris"})` call, + `tests/real_toolcall.rs`). - **f16 inference, validated** — Qwen2.5-1.5B runs coherently on `GpuF16` (weights + KV in f16, the q·kᵀ attention scores + softmax computed in an f32 island to stop f16 overflow): **~3.6 GiB runner VRAM vs ~7.9 GiB f32, at identical speed** (14.1 tok/s / 88 ms TTFT); the parity gate re-passes diff --git a/ROADMAP.md b/ROADMAP.md index 100ca37..6155989 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -188,11 +188,19 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari (`qwen2` plain, `lfm2` with `<|startoftext|>` BOS); byte-verified — the Qwen2 parity gate now renders its prompt through the template and still matches the Candle fixture and the Ollama fp16 greedy leg exactly. EOS stays config-driven (`EosIds`); tool-use templates are the next item.* -- [ ] Hermes-style tool-use chat template (Qwen3 ships it in `tokenizer_config.json`) + LFM2.5's - bracket-notation tool-call output — the two top scorers (0.880 agent score; LFM2.5-1.2B also the - fastest at ~1.5 s) on 2026's 21-model local tool-calling benchmark; function calling is why the - apps want a local runner — https://mikeveerman.be/blog/github-2026-02-06-tool-calling-benchmark/ · +- [x] Hermes-style tool-use chat template (the format Qwen2.5/Qwen3 ship in `tokenizer_config.json`) — + function calling is why the apps want a local runner — https://qwen.readthedocs.io/en/latest/framework/function_call.html + *(2026-07-11) `chat`: `ToolSpec` → `render_with_tools` (byte-matches the Qwen template's `# Tools` + /``/`` wording), `Turn::{assistant_tool_calls, tool_response}` (consecutive tool + results merge into one user turn, per the template), and a bounded `parse_tool_calls` extractor + (calls + prose, loud error taxonomy). REAL-GPU proof (`tests/real_toolcall.rs`): Qwen2.5-1.5B + greedy-emitted `{"name": "get_weather", "arguments": {"city": "Paris"}}` + from a rendered prompt and the parser round-tripped it. 10 new unit tests.* +- [ ] LFM2.5 bracket-notation tool-call template + parser (`<|tool_list_start|>` special tokens, + Python-ish call syntax) — with Hermes/Qwen2.5 (0.880 agent score) done, LFM2.5-1.2B (same score, + fastest at ~1.5 s on 2026's 21-model local tool-calling benchmark) is the other target — + https://mikeveerman.be/blog/github-2026-02-06-tool-calling-benchmark/ *(2026-07-10 research)* 2026 community numbers back the plan: Qwen3-8B keeps tool-calling score through Q4_K_M (0.919 quantized vs 0.933 full — quant does NOT cost tool reliability, good news for P9); BFCL shows a capability cliff below ~7B (Qwen3.5-9B 66.1% vs 4B 50.3%), so the zoo's diff --git a/crates/mummu/src/chat.rs b/crates/mummu/src/chat.rs index 10fbfe6..fdbce37 100644 --- a/crates/mummu/src/chat.rs +++ b/crates/mummu/src/chat.rs @@ -5,7 +5,11 @@ //! the exact prompt committed in the Candle logits fixture). //! //! Both zoo LLMs speak ChatML; LFM2.5 additionally prefixes `<|startoftext|>`. -//! Tool-use templates (Hermes-style, LFM2.5 bracket notation) are P4 follow-ups. +//! Tool use follows the Hermes convention Qwen2.5/Qwen3 ship in their chat +//! template: tool signatures in a `` block of the system turn, calls +//! emitted as `{json}`, results returned inside +//! `` blocks of a user turn. LFM2.5's bracket notation is a +//! P4 follow-up. /// Who is speaking in a [`Turn`]. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -13,6 +17,9 @@ pub enum Role { System, User, Assistant, + /// A tool result going back to the model. Hermes-style templates render + /// these inside a *user* turn as `` blocks. + Tool, } impl Role { @@ -21,6 +28,7 @@ impl Role { Self::System => "system", Self::User => "user", Self::Assistant => "assistant", + Self::Tool => "user", // Hermes: tool results ride in a user turn } } } @@ -56,6 +64,118 @@ impl Turn { content: content.into(), } } + + /// An assistant turn that invokes tools: each call becomes a Hermes + /// `` block in the turn body (what the model itself would + /// have emitted), so histories containing calls re-render faithfully. + #[must_use] + pub fn assistant_tool_calls(calls: &[ToolCall]) -> Self { + assert!(!calls.is_empty(), "assistant_tool_calls: no calls"); + assert!( + calls.len() <= MAX_TOOL_CALLS, + "assistant_tool_calls: {} calls exceeds the {MAX_TOOL_CALLS} bound", + calls.len() + ); + let blocks: Vec = calls + .iter() + .map(|c| { + let json = serde_json::to_string(c).unwrap_or_default(); + debug_assert!(!json.is_empty(), "a ToolCall always serializes"); + format!("\n{json}\n") + }) + .collect(); + Self { + role: Role::Assistant, + content: blocks.join("\n"), + } + } + + /// A tool's result going back to the model; renders as a + /// `` block (consecutive ones merge into one user turn). + #[must_use] + pub fn tool_response(content: impl Into) -> Self { + Self { + role: Role::Tool, + content: content.into(), + } + } +} + +/// A callable tool signature, serialized into the system prompt exactly as +/// the Hermes-style templates expect: `{"type": "function", "function": +/// {"name": …, "description": …, "parameters": }}`. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct ToolSpec { + pub name: String, + pub description: String, + /// JSON schema of the arguments object. + pub parameters: serde_json::Value, +} + +/// The Hermes wire shape for one tool (field order matters for byte-stable +/// rendering, so this is a struct, not a `json!` map). +#[derive(serde::Serialize)] +struct ToolWire<'a> { + r#type: &'static str, + function: &'a ToolSpec, +} + +/// One tool invocation, as emitted by the model inside `` tags. +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct ToolCall { + pub name: String, + #[serde(default)] + pub arguments: serde_json::Value, +} + +/// Most tool calls a single response may contain (or a single assistant +/// history turn may carry) — far past anything a small model emits. +pub const MAX_TOOL_CALLS: usize = 64; + +/// Most tools one render will advertise. +const MAX_TOOLS: usize = 128; + +/// What went wrong extracting tool calls from a model response. +#[derive(Debug, thiserror::Error)] +pub enum ToolCallError { + #[error("tool call {index}: unclosed tag")] + Unclosed { index: usize }, + #[error("tool call {index}: {reason}")] + BadJson { index: usize, reason: String }, + #[error("more than {MAX_TOOL_CALLS} tool calls in one response")] + TooMany, +} + +/// Extract Hermes-style tool calls from a model response: every +/// `…` block parses as a [`ToolCall`]; the text +/// outside the blocks (the model's prose, trimmed) comes back alongside. +/// Text with no blocks is simply `(vec![], text)` — not an error. +pub fn parse_tool_calls(text: &str) -> Result<(Vec, String), ToolCallError> { + const OPEN: &str = ""; + const CLOSE: &str = ""; + let mut calls = Vec::new(); + let mut prose = String::new(); + let mut rest = text; + while let Some(start) = rest.find(OPEN) { + if calls.len() == MAX_TOOL_CALLS { + return Err(ToolCallError::TooMany); + } + prose.push_str(&rest[..start]); + let after_open = &rest[start + OPEN.len()..]; + let Some(end) = after_open.find(CLOSE) else { + return Err(ToolCallError::Unclosed { index: calls.len() }); + }; + let body = after_open[..end].trim(); + let call: ToolCall = serde_json::from_str(body).map_err(|e| ToolCallError::BadJson { + index: calls.len(), + reason: e.to_string(), + })?; + calls.push(call); + rest = &after_open[end + CLOSE.len()..]; + } + prose.push_str(rest); + debug_assert!(calls.len() <= MAX_TOOL_CALLS, "bound enforced in the loop"); + Ok((calls, prose.trim().to_string())) } /// Longest conversation a single render will wrap — a generous bound that @@ -102,17 +222,84 @@ impl ChatMl { a trailing assistant turn would double it" ); let mut out = String::from(self.bos.unwrap_or("")); - for turn in turns { - out.push_str("<|im_start|>"); - out.push_str(turn.role.tag()); - out.push('\n'); - out.push_str(&turn.content); - out.push_str("<|im_end|>\n"); + let mut i = 0; + while i < turns.len() { + if turns[i].role == Role::Tool { + // Hermes: consecutive tool results merge into ONE user turn, + // each wrapped in its own block. + out.push_str("<|im_start|>user"); + while i < turns.len() && turns[i].role == Role::Tool { + out.push_str("\n\n"); + out.push_str(&turns[i].content); + out.push_str("\n"); + i += 1; + } + out.push_str("<|im_end|>\n"); + } else { + out.push_str("<|im_start|>"); + out.push_str(turns[i].role.tag()); + out.push('\n'); + out.push_str(&turns[i].content); + out.push_str("<|im_end|>\n"); + i += 1; + } } out.push_str("<|im_start|>assistant\n"); debug_assert!(out.ends_with("assistant\n"), "render must open a turn"); out } + + /// [`render`](Self::render) with Hermes-style function calling: the tool + /// signatures are advertised in a `# Tools` section of the system turn — + /// the exact wording and tag structure Qwen2.5/Qwen3 ship in their chat + /// template. An existing leading system turn provides the preamble; a + /// conversation without one gets a neutral "You are a helpful assistant." + #[must_use] + pub fn render_with_tools(&self, tools: &[ToolSpec], turns: &[Turn]) -> String { + assert!( + !tools.is_empty(), + "render_with_tools: no tools — use render()" + ); + assert!( + tools.len() <= MAX_TOOLS, + "render_with_tools: {} tools exceeds the {MAX_TOOLS} bound", + tools.len() + ); + assert!( + tools.iter().all(|t| !t.name.is_empty()), + "render_with_tools: every tool needs a name" + ); + + let (preamble, rest) = match turns.first() { + Some(t) if t.role == Role::System => (t.content.as_str(), &turns[1..]), + _ => ("You are a helpful assistant.", turns), + }; + let mut system = String::from(preamble); + system.push_str( + "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\n\ + You are provided with function signatures within XML tags:\n", + ); + for tool in tools { + let json = serde_json::to_string(&ToolWire { + r#type: "function", + function: tool, + }) + .unwrap_or_default(); + debug_assert!(!json.is_empty(), "a ToolSpec always serializes"); + system.push('\n'); + system.push_str(&json); + } + system.push_str( + "\n\n\nFor each function call, return a json object with function name and \ + arguments within XML tags:\n\n{\"name\": \ + , \"arguments\": }\n", + ); + + let mut wrapped = Vec::with_capacity(rest.len() + 1); + wrapped.push(Turn::system(system)); + wrapped.extend_from_slice(rest); + self.render(&wrapped) + } } #[cfg(test)] @@ -167,6 +354,125 @@ mod tests { let _ = ChatMl::qwen2().render(&[]); } + fn weather_tool() -> ToolSpec { + ToolSpec { + name: "get_weather".into(), + description: "Get the current weather for a city.".into(), + parameters: serde_json::json!({ + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"] + }), + } + } + + /// The tools section must match the Qwen2.5/Qwen3 chat template's wording + /// and tag structure byte-for-byte (the model was trained on this text). + /// Inside a tool's `parameters` schema, keys serialize in serde_json's + /// canonical (sorted) order — key order isn't part of the trained text. + #[test] + fn tools_render_matches_the_hermes_template_shape() { + let raw = ChatMl::qwen2().render_with_tools( + &[weather_tool()], + &[ + Turn::system("You are a helpful assistant."), + Turn::user("Weather in Paris?"), + ], + ); + let expected = "<|im_start|>system\nYou are a helpful assistant.\n\n# Tools\n\n\ + You may call one or more functions to assist with the user query.\n\n\ + You are provided with function signatures within XML tags:\n\ + \n\ + {\"type\":\"function\",\"function\":{\"name\":\"get_weather\",\"description\":\"Get the current weather for a city.\",\"parameters\":{\"properties\":{\"city\":{\"type\":\"string\"}},\"required\":[\"city\"],\"type\":\"object\"}}}\n\ + \n\n\ + For each function call, return a json object with function name and arguments within XML tags:\n\ + \n{\"name\": , \"arguments\": }\n<|im_end|>\n\ + <|im_start|>user\nWeather in Paris?<|im_end|>\n\ + <|im_start|>assistant\n"; + assert_eq!(raw, expected); + } + + #[test] + fn tools_render_without_a_system_turn_injects_a_neutral_preamble() { + let raw = ChatMl::qwen2().render_with_tools(&[weather_tool()], &[Turn::user("hi")]); + assert!(raw.starts_with("<|im_start|>system\nYou are a helpful assistant.\n\n# Tools")); + // The user turn survives un-consumed. + assert!(raw.contains("<|im_start|>user\nhi<|im_end|>")); + } + + #[test] + fn consecutive_tool_responses_merge_into_one_user_turn() { + let calls = [ToolCall { + name: "get_weather".into(), + arguments: serde_json::json!({"city": "Paris"}), + }]; + let raw = ChatMl::qwen2().render(&[ + Turn::user("Weather in Paris and Lyon?"), + Turn::assistant_tool_calls(&calls), + Turn::tool_response("{\"temp_c\": 21}"), + Turn::tool_response("{\"temp_c\": 24}"), + ]); + // The assistant history turn carries the block it emitted. + assert!(raw.contains( + "<|im_start|>assistant\n\n\ + {\"name\":\"get_weather\",\"arguments\":{\"city\":\"Paris\"}}\n\ + <|im_end|>\n" + )); + // Both results ride in ONE user turn, each in its own block. + assert!(raw.contains( + "<|im_start|>user\n\ + \n{\"temp_c\": 21}\n\n\ + \n{\"temp_c\": 24}\n<|im_end|>\n" + )); + assert_eq!(raw.matches("<|im_start|>user").count(), 2); + } + + #[test] + fn parse_extracts_calls_and_prose() { + let text = "Let me check.\n\n{\"name\": \"get_weather\", \ + \"arguments\": {\"city\": \"Paris\"}}\n\n\ + \n{\"name\": \"get_weather\", \"arguments\": \ + {\"city\": \"Lyon\"}}\n"; + let (calls, prose) = parse_tool_calls(text).unwrap(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].name, "get_weather"); + assert_eq!(calls[0].arguments["city"], "Paris"); + assert_eq!(calls[1].arguments["city"], "Lyon"); + assert_eq!(prose, "Let me check."); + } + + #[test] + fn parse_of_plain_text_is_empty_not_an_error() { + let (calls, prose) = parse_tool_calls("The answer is 4.").unwrap(); + assert!(calls.is_empty()); + assert_eq!(prose, "The answer is 4."); + } + + #[test] + fn parse_rejects_unclosed_and_bad_json() { + assert!(matches!( + parse_tool_calls("\n{\"name\": \"x\"}"), + Err(ToolCallError::Unclosed { index: 0 }) + )); + assert!(matches!( + parse_tool_calls("\nnot json\n"), + Err(ToolCallError::BadJson { index: 0, .. }) + )); + } + + /// The whole loop: a rendered history turn re-parses to the same calls. + #[test] + fn tool_calls_round_trip_through_render_and_parse() { + let calls = vec![ToolCall { + name: "lookup".into(), + arguments: serde_json::json!({"q": "primes", "k": 5}), + }]; + let turn = Turn::assistant_tool_calls(&calls); + let (parsed, prose) = parse_tool_calls(&turn.content).unwrap(); + assert_eq!(parsed, calls); + assert!(prose.is_empty()); + } + #[test] #[should_panic(expected = "double it")] fn trailing_assistant_turn_is_rejected() { diff --git a/crates/mummu/tests/real_toolcall.rs b/crates/mummu/tests/real_toolcall.rs new file mode 100644 index 0000000..d8cfeb6 --- /dev/null +++ b/crates/mummu/tests/real_toolcall.rs @@ -0,0 +1,77 @@ +//! Real-model tool-use proof: render a Hermes-style tools prompt, let +//! Qwen2.5-1.5B decode on the GPU, and parse the `` block it +//! actually emits. Ignored by default; run with +//! +//! ```text +//! MUMMU_QWEN2_DIR=path/to/qwen2.5-1.5b cargo test -p mummu --release --test real_toolcall -- --ignored --nocapture +//! ``` + +use std::path::PathBuf; + +use mummu::backend::Gpu; +use mummu::chat::{ChatMl, ToolSpec, Turn, parse_tool_calls}; +use mummu::models::CausalLm; +use mummu::models::qwen2; +use tokenizers::Tokenizer; + +fn qwen2_dir() -> Option { + let dir = PathBuf::from(std::env::var_os("MUMMU_QWEN2_DIR")?); + dir.is_dir().then_some(dir) +} + +#[test] +#[ignore = "needs multi-GB local weights (MUMMU_QWEN2_DIR) + the reference GPU"] +fn qwen2_emits_a_parseable_tool_call() { + let Some(dir) = qwen2_dir() else { + panic!("set MUMMU_QWEN2_DIR to a dir with config.json/tokenizer.json/model.safetensors"); + }; + + let tools = [ToolSpec { + name: "get_weather".into(), + description: "Get the current weather for a city.".into(), + parameters: serde_json::json!({ + "type": "object", + "properties": { + "city": { "type": "string", "description": "City name" } + }, + "required": ["city"] + }), + }]; + let raw = ChatMl::qwen2().render_with_tools( + &tools, + &[Turn::user("What is the weather in Paris right now?")], + ); + + let tok = Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer.json loads"); + let prompt = tok + .encode(raw.as_str(), false) + .expect("encodes") + .get_ids() + .to_vec(); + + let device = burn::tensor::Device::::default(); + let loaded = qwen2::load_from_dir::(&dir, &device).expect("weights load checked"); + let ids = loaded + .greedy_generate(&prompt, 64, &device) + .expect("greedy decode"); + let text = tok.decode(&ids, false).expect("decode"); + eprintln!("[real_toolcall] model emitted: {text:?}"); + + let (calls, prose) = parse_tool_calls(&text).expect("emitted tool call parses"); + assert_eq!( + calls.len(), + 1, + "expected exactly one call, prose: {prose:?}" + ); + assert_eq!(calls[0].name, "get_weather"); + assert_eq!( + calls[0].arguments["city"].as_str(), + Some("Paris"), + "arguments: {:?}", + calls[0].arguments + ); + eprintln!( + "[real_toolcall] parsed: {} with {}", + calls[0].name, calls[0].arguments + ); +} From 35438c073b3977b3a003f6d38952417365f456b1 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sat, 11 Jul 2026 07:38:55 -0500 Subject: [PATCH 5/6] =?UTF-8?q?feat(p3):=20PyTorch=20state-dict=20import?= =?UTF-8?q?=20=E2=80=94=20pytorch=5Fmodel.bin=20through=20the=20checked-lo?= =?UTF-8?q?ad=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit burn-store's PytorchStore (pickle/zip reader, built-in PyTorchToBurnAdapter) joins safetensors behind one import surface: - import::weights_file picks model.safetensors first, falls back to pytorch_model.bin, and reports the safetensors name when neither exists. - load_checked generalizes over any ModuleStore, keeping the fail-loud missing/errored-param contract for every format. - MiniLM loads either format through one shared KEY_REMAPS table (the .bin path notes its no-cast limitation: .bin-era checkpoints are f32). - Direct burn-store workspace dep (same 0.21 crate burn re-exports, feature-unified) turns on its pytorch feature. REAL-WEIGHTS proof (tests/real_pytorch.rs, ignored): the Hub's actual MiniLM pytorch_model.bin embeds byte-identically (max |delta| = 0e0, cosine 1.0) to the safetensors copy of the same weights on the flex CPU backend. 91 unit tests green (1 new), clippy clean. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 293 ++++++++++++++++++++++++++--- Cargo.toml | 8 + README.md | 6 +- ROADMAP.md | 11 +- crates/mummu/Cargo.toml | 1 + crates/mummu/src/import.rs | 75 ++++++-- crates/mummu/src/models/minilm.rs | 95 ++++++---- crates/mummu/tests/real_pytorch.rs | 75 ++++++++ 8 files changed, 494 insertions(+), 70 deletions(-) create mode 100644 crates/mummu/tests/real_pytorch.rs diff --git a/Cargo.lock b/Cargo.lock index fc02467..4f26f44 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -24,10 +24,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" dependencies = [ "cfg-if", - "cipher", + "cipher 0.4.4", "cpufeatures 0.2.17", ] +[[package]] +name = "aes" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1fc76eaeac4c9164506c466d4ffdd8ec9d0c5bf57ee97177c4d8eceb3a0e138" +dependencies = [ + "cipher 0.5.2", + "cpubits", + "cpufeatures 0.3.0", +] + [[package]] name = "ahash" version = "0.8.12" @@ -336,6 +347,16 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", + "zeroize", +] + [[package]] name = "block2" version = "0.6.2" @@ -484,6 +505,7 @@ dependencies = [ "serde", "serde_json", "spin", + "thiserror 2.0.18", "uuid", ] @@ -733,7 +755,10 @@ dependencies = [ "memmap2", "regex", "safetensors 0.7.0", + "serde", + "tar", "textdistance", + "zip 8.6.0", ] [[package]] @@ -851,6 +876,15 @@ dependencies = [ "libc", ] +[[package]] +name = "bzip2" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3a53fac24f34a81bc9954b5d6cfce0c21e18ec6959f44f56e8e90e4bb7c346c" +dependencies = [ + "libbz2-rs-sys", +] + [[package]] name = "bzip2-sys" version = "0.1.13+1.0.8" @@ -986,8 +1020,18 @@ version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" dependencies = [ - "crypto-common", - "inout", + "crypto-common 0.1.7", + "inout 0.1.4", +] + +[[package]] +name = "cipher" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8cf2a2c93cd704877c0858356ed03480ff301ee950b43f1cbe4573b088bfa6c" +dependencies = [ + "crypto-common 0.2.2", + "inout 0.2.2", ] [[package]] @@ -1026,6 +1070,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" +[[package]] +name = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + [[package]] name = "codespan-reporting" version = "0.13.1" @@ -1102,6 +1152,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "const-random" version = "0.1.18" @@ -1128,6 +1184,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "245097e9a4535ee1e3e3931fcfcd55a796a44c643e8596ff6566d68f09b87bbc" +[[package]] +name = "constant_time_eq" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" + [[package]] name = "constcat" version = "0.6.1" @@ -1152,6 +1214,12 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "cpubits" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15b85f9c39137c3a891689859392b1bd49812121d0d61c9caf00d46ed5ce06ae" + [[package]] name = "cpufeatures" version = "0.2.17" @@ -1261,6 +1329,24 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "ctutils" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +dependencies = [ + "cmov", +] + [[package]] name = "cubecl" version = "0.10.0" @@ -1839,6 +1925,12 @@ version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" +[[package]] +name = "deflate64" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac6b926516df9c60bfa16e107b21086399f8285a44ca9711344b9e553c5146e2" + [[package]] name = "deranged" version = "0.5.8" @@ -1922,11 +2014,24 @@ version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.7", "subtle", ] +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "const-oid", + "crypto-common 0.2.2", + "ctutils", + "zeroize", +] + [[package]] name = "dirs" version = "6.0.0" @@ -2245,6 +2350,7 @@ checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ "crc32fast", "miniz_oxide", + "zlib-rs", ] [[package]] @@ -2712,7 +2818,16 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" dependencies = [ - "digest", + "digest 0.10.7", +] + +[[package]] +name = "hmac" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f" +dependencies = [ + "digest 0.11.3", ] [[package]] @@ -2754,6 +2869,15 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" +[[package]] +name = "hybrid-array" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "1.10.1" @@ -3001,6 +3125,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "inout" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4250ce6452e92010fdf7268ccc5d14faa80bb12fc741938534c58f16804e03c7" +dependencies = [ + "hybrid-array", +] + [[package]] name = "interpolate_name" version = "0.2.4" @@ -3120,6 +3253,12 @@ version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7a79a3332a6609480d7d0c9eab957bca6b455b91bb84e66d19f5ff66294b85b8" +[[package]] +name = "libbz2-rs-sys" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34b357333733e8260735ba5894eb928c02ecc69c78715f01a8019e7fa7f2db4c" + [[package]] name = "libc" version = "0.2.186" @@ -3240,6 +3379,15 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" +[[package]] +name = "lzma-rust2" +version = "0.16.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca93e534d1142d1d0dcca6d25fe302508a5dfb40b302802904577725ea0b695b" +dependencies = [ + "sha2 0.11.0", +] + [[package]] name = "macerator" version = "0.3.3" @@ -3397,12 +3545,13 @@ version = "0.1.0" dependencies = [ "burn", "burn-flex", + "burn-store", "half", "once_cell", "pollster", "serde", "serde_json", - "sha2", + "sha2 0.10.9", "thiserror 2.0.18", "tokenizers", "ureq 3.3.0", @@ -3853,10 +4002,20 @@ version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "83a0692ec44e4cf1ef28ca317f14f8f07da2d95ec3fa01f86e4467b725e60917" dependencies = [ - "digest", - "hmac", + "digest 0.10.7", + "hmac 0.12.1", "password-hash", - "sha2", + "sha2 0.10.9", +] + +[[package]] +name = "pbkdf2" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112d82ceb8c5bf524d9af484d4e4970c9fd5a0cc15ba14ad93dccd28873b0629" +dependencies = [ + "digest 0.11.3", + "hmac 0.13.0", ] [[package]] @@ -3969,6 +4128,12 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" +[[package]] +name = "ppmd-rust" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "efca4c95a19a79d1c98f791f10aebd5c1363b473244630bb7dbde1dc98455a24" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -4736,7 +4901,18 @@ checksum = "a978451301f4db1d02937a4ab3ccce137717b81826e79b7d49ffe3244a13c3b8" dependencies = [ "cfg-if", "cpufeatures 0.2.17", - "digest", + "digest 0.10.7", +] + +[[package]] +name = "sha1" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", ] [[package]] @@ -4747,7 +4923,18 @@ checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", "cpufeatures 0.2.17", - "digest", + "digest 0.10.7", +] + +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", ] [[package]] @@ -5056,6 +5243,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "18dfaaeddcb932337b5e7866ee7d0ce9b76d2fd092997146f187ec09b4558a50" dependencies = [ "deranged", + "js-sys", "num-conv", "powerfmt", "serde_core", @@ -5307,7 +5495,7 @@ dependencies = [ "reqwest", "serde", "serde_json", - "sha2", + "sha2 0.10.9", "tar", "walkdir", ] @@ -6308,18 +6496,18 @@ version = "0.6.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "760394e246e4c28189f19d488c058bf16f564016aefac5d32bb1f3b51d5e9261" dependencies = [ - "aes", + "aes 0.8.4", "byteorder", - "bzip2", - "constant_time_eq", + "bzip2 0.4.4", + "constant_time_eq 0.1.5", "crc32fast", "crossbeam-utils", "flate2", - "hmac", - "pbkdf2", - "sha1", + "hmac 0.12.1", + "pbkdf2 0.11.0", + "sha1 0.10.7", "time", - "zstd", + "zstd 0.11.2+zstd.1.5.2", ] [[package]] @@ -6334,19 +6522,73 @@ dependencies = [ "typed-path", ] +[[package]] +name = "zip" +version = "8.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d04a6b5381502aa6087c94c669499eb1602eb9c5e8198e534de571f7154809b" +dependencies = [ + "aes 0.9.1", + "bzip2 0.6.1", + "constant_time_eq 0.4.2", + "crc32fast", + "deflate64", + "flate2", + "getrandom 0.4.3", + "hmac 0.13.0", + "indexmap", + "lzma-rust2", + "memchr", + "pbkdf2 0.13.0", + "ppmd-rust", + "sha1 0.11.0", + "time", + "typed-path", + "zeroize", + "zopfli", + "zstd 0.13.3", +] + +[[package]] +name = "zlib-rs" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b142a20ec14a91d5bc708c1dc21b080c550113d8aa77afa29635673a65dd02c5" + [[package]] name = "zmij" version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +[[package]] +name = "zopfli" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f05cd8797d63865425ff89b5c4a48804f35ba0ce8d125800027ad6017d2b5249" +dependencies = [ + "bumpalo", + "crc32fast", + "log", + "simd-adler32", +] + [[package]] name = "zstd" version = "0.11.2+zstd.1.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20cc960326ece64f010d2d2107537f26dc589a6573a316bd5b1dba685fa5fde4" dependencies = [ - "zstd-safe", + "zstd-safe 5.0.2+zstd.1.5.2", +] + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe 7.2.4", ] [[package]] @@ -6359,6 +6601,15 @@ dependencies = [ "zstd-sys", ] +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + [[package]] name = "zstd-sys" version = "2.0.16+zstd.1.5.7" diff --git a/Cargo.toml b/Cargo.toml index 2e8a63f..9767c26 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,14 @@ burn = { version = "0.21", default-features = false, features = [ # gate passed on it bit-for-bit-equivalent (cosine 0.99999994, max component # diff 1.3e-7) with all 80 unit tests green. burn-flex = "0.21" +# Direct handle on burn-store (the same crate burn's `store` feature pulls) to +# turn on its `pytorch` feature: .pth / pytorch_model.bin state-dict loading +# for models never re-shipped as safetensors (feature-unified, no duplicate). +burn-store = { version = "0.21", default-features = false, features = [ + "std", + "safetensors", + "pytorch", +] } # Direct handle on wgpu (the same version Burn resolves) for the cheap startup # adapter probe: enumerate GPUs to decide GPU vs CPU before building a backend. diff --git a/README.md b/README.md index 3b53bb1..0e372dc 100644 --- a/README.md +++ b/README.md @@ -24,8 +24,10 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys - **Shared blocks, generic over `B: Backend`** — cache-aware GQA attention (optional per-head q/k RMSNorm), manual RoPE, SwiGLU, and LFM2's double-gated causal short-conv with rolling decode state; unit tests prove prefill+decode ≡ full-forward for both cache kinds. -- **Checked safetensors import** — bf16→backend-float cast adapter, per-architecture key remaps, and a - fail-loud load (never silently zero-init); `config.json`-driven hyperparameters. +- **Checked safetensors + PyTorch import** — bf16→backend-float cast adapter, per-architecture key + remaps, and a fail-loud load (never silently zero-init); `config.json`-driven hyperparameters. + `pytorch_model.bin` state dicts load through the same checked path (safetensors preferred when both + exist) — proven byte-identical on MiniLM's real Hub checkpoint in both formats. - **Three models ported and running on real weights** — Qwen2/2.5, the LFM2/2.5 hybrid, and the all-MiniLM sentence embedder; Qwen2.5-1.5B and LFM2.5-1.2B load and greedy-decode correctly on the reference GPU (wgpu/Vulkan). diff --git a/ROADMAP.md b/ROADMAP.md index 6155989..554da34 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -142,7 +142,16 @@ The subsystem that turns "a model on HuggingFace or on disk" into a loaded, pari - [x] **safetensors** *(ex-laurelane)* — `burn-store` `SafetensorsStore` + `PyTorchToBurnAdapter`; the primary path. *(2026-07-09) `import::{CastFloatAdapter, load_checked}`: bf16→backend-float cast + fail-loud load; proven by loading the real 3.1 GB Qwen2.5-1.5B and 2.3 GB LFM2.5 checkpoints with zero missing keys.* -- [ ] **PyTorch state dicts** (`.pth` / `pytorch_model*.bin`) — for models not shipped as safetensors. +- [x] **PyTorch state dicts** (`.pth` / `pytorch_model*.bin`) — for models not shipped as safetensors. + *(2026-07-11)* `burn-store`'s `PytorchStore` wired through the shared checked-load path: + `import::weights_file` picks `model.safetensors` first, falls back to `pytorch_model.bin`; + `load_checked` is now generic over any `ModuleStore`; MiniLM loads either format through one + remap table. REAL-WEIGHTS proof (`tests/real_pytorch.rs`): the Hub's actual MiniLM + `pytorch_model.bin` embeds **byte-identically** (max |Δ| = 0) to the safetensors copy of the same + weights. Remaining follow-ups: sharded `.bin` indexes, a bf16-cast on this path (PytorchStore has + 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, 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)* diff --git a/crates/mummu/Cargo.toml b/crates/mummu/Cargo.toml index f521e06..4aded66 100644 --- a/crates/mummu/Cargo.toml +++ b/crates/mummu/Cargo.toml @@ -19,3 +19,4 @@ ureq = { workspace = true } sha2 = { workspace = true } once_cell = { workspace = true } burn-flex = { workspace = true } +burn-store = { workspace = true } diff --git a/crates/mummu/src/import.rs b/crates/mummu/src/import.rs index 828fa2c..a16339e 100644 --- a/crates/mummu/src/import.rs +++ b/crates/mummu/src/import.rs @@ -1,15 +1,17 @@ -//! Weight import: safetensors → Burn modules, checked and loud. +//! Weight import: checkpoint files → Burn modules, checked and loud. //! //! The pieces every model load shares (P3): a dtype-cast adapter (HF ships -//! bf16, which wgpu can't ingest directly), and a checked-load wrapper that -//! **fails on missing or errored params** instead of silently zero-initing — -//! a partial load is a quietly broken model. +//! bf16, which wgpu can't ingest directly), a weights-file picker +//! (**safetensors** preferred, the PyTorch state dict `pytorch_model.bin` as +//! the fallback for models never re-shipped as safetensors), and a +//! checked-load wrapper that **fails on missing or errored params** instead +//! of silently zero-initing — a partial load is a quietly broken model. use std::path::{Path, PathBuf}; use std::rc::Rc; use burn::module::Module; -use burn::store::{ModuleAdapter, ModuleSnapshot, SafetensorsStore, TensorSnapshot}; +use burn::store::{ModuleAdapter, ModuleSnapshot, ModuleStore, TensorSnapshot}; use burn::tensor::DType; /// Everything that can go wrong turning files on disk into a loaded model. @@ -80,15 +82,21 @@ impl ModuleAdapter for CastFloatAdapter { } } -/// Load `store` into `module`, refusing partial results: any missing param or -/// per-tensor error is an [`ImportError::Incomplete`] carrying the store's own -/// readable report. Unused checkpoint tensors are *allowed* (e.g. BERT's +/// Load `store` (any format: safetensors, PyTorch state dict, …) into +/// `module`, refusing partial results: any missing param or per-tensor error +/// is an [`ImportError::Incomplete`] carrying the store's own readable +/// report. Unused checkpoint tensors are *allowed* (e.g. BERT's /// intentionally-skipped `pooler.*`) — callers that care inspect the report. -pub fn load_checked + ModuleSnapshot, B: burn::tensor::backend::Backend>( +pub fn load_checked( module: &mut M, - store: &mut SafetensorsStore, + store: &mut S, weights_path: &Path, -) -> Result<(), ImportError> { +) -> Result<(), ImportError> +where + M: Module + ModuleSnapshot, + B: burn::tensor::backend::Backend, + S: ModuleStore, +{ let report = module.load_from(store).map_err(|e| ImportError::Load { file: weights_path.to_path_buf(), reason: e.to_string(), @@ -120,6 +128,30 @@ pub fn required_file(dir: &Path, file: &str) -> Result { } } +/// The weights checkpoint found in a model dir. +#[derive(Debug, Clone)] +pub enum WeightsFile { + /// `model.safetensors` — the primary format. + Safetensors(PathBuf), + /// `pytorch_model.bin` — the PyTorch state dict older checkpoints ship. + PytorchBin(PathBuf), +} + +/// Pick the weights file in `dir`: `model.safetensors` when present, else +/// `pytorch_model.bin`. Reports the *safetensors* name when neither exists +/// (it's the file a fresh download would produce). +pub fn weights_file(dir: &Path) -> Result { + let safetensors = dir.join("model.safetensors"); + if safetensors.is_file() { + return Ok(WeightsFile::Safetensors(safetensors)); + } + let bin = dir.join("pytorch_model.bin"); + if bin.is_file() { + return Ok(WeightsFile::PytorchBin(bin)); + } + Err(ImportError::MissingFile(safetensors)) +} + #[cfg(test)] mod tests { use super::*; @@ -140,4 +172,25 @@ mod tests { fn cast_adapter_rejects_non_float_target() { let _ = CastFloatAdapter::new(DType::I32); } + + #[test] + fn weights_file_prefers_safetensors_falls_back_to_pytorch() { + let dir = std::env::temp_dir().join("mummu_weights_file_test"); + std::fs::create_dir_all(&dir).unwrap(); + // Neither present: missing, reported by the safetensors name. + assert!(matches!( + weights_file(&dir), + Err(ImportError::MissingFile(p)) if p.ends_with("model.safetensors") + )); + // Only the state dict: the PyTorch path. + std::fs::write(dir.join("pytorch_model.bin"), b"pt").unwrap(); + assert!(matches!(weights_file(&dir), Ok(WeightsFile::PytorchBin(_)))); + // Both: safetensors wins. + std::fs::write(dir.join("model.safetensors"), b"st").unwrap(); + assert!(matches!( + weights_file(&dir), + Ok(WeightsFile::Safetensors(_)) + )); + std::fs::remove_dir_all(&dir).unwrap(); + } } diff --git a/crates/mummu/src/models/minilm.rs b/crates/mummu/src/models/minilm.rs index a56593f..02be78f 100644 --- a/crates/mummu/src/models/minilm.rs +++ b/crates/mummu/src/models/minilm.rs @@ -15,10 +15,12 @@ use std::path::Path; use burn::module::Module; use burn::nn::{Embedding, EmbeddingConfig, LayerNorm, LayerNormConfig, Linear, LinearConfig}; -use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, SafetensorsStore}; +use burn::store::{ModuleAdapter, PyTorchToBurnAdapter, PytorchStore, SafetensorsStore}; use burn::tensor::{Int, Tensor, TensorData, activation, backend::Backend}; -use crate::import::{CastFloatAdapter, ImportError, load_checked, required_file}; +use crate::import::{ + CastFloatAdapter, ImportError, WeightsFile, load_checked, required_file, weights_file, +}; /// BERT hyperparameters, read from the checkpoint's `config.json`. #[derive(Debug, Clone, serde::Deserialize)] @@ -117,17 +119,49 @@ fn build(cfg: &BertConfig, device: &B::Device) -> Bert { } } -/// Build from `dir/config.json` and load `dir/model.safetensors`, checked. +/// HF BERT checkpoint names → our field paths, shared by every weight format. +const KEY_REMAPS: &[(&str, &str)] = &[ + (r"^bert\.", ""), // some checkpoints prefix `bert.` + (r"^embeddings\.LayerNorm\.", "embeddings.layer_norm."), + ( + r"^encoder\.layer\.(\d+)\.attention\.self\.(query|key|value)\.", + "layers.$1.$2.", + ), + ( + r"^encoder\.layer\.(\d+)\.attention\.output\.dense\.", + "layers.$1.attn_output.", + ), + ( + r"^encoder\.layer\.(\d+)\.attention\.output\.LayerNorm\.", + "layers.$1.attn_layer_norm.", + ), + ( + r"^encoder\.layer\.(\d+)\.intermediate\.dense\.", + "layers.$1.intermediate.", + ), + ( + r"^encoder\.layer\.(\d+)\.output\.dense\.", + "layers.$1.output.", + ), + ( + r"^encoder\.layer\.(\d+)\.output\.LayerNorm\.", + "layers.$1.output_layer_norm.", + ), +]; + +/// Build from `dir/config.json` and load the checkpoint, checked — +/// `model.safetensors` preferred, `pytorch_model.bin` (the state dict this +/// model family originally shipped) as the fallback. /// NOTE: LayerNorm keys keep their `.weight`/`.bias` suffixes — the /// `PyTorchToBurnAdapter` renames those to `gamma`/`beta` itself (unlike /// RmsNorm in the decoder models, where the rename is manual — the asymmetry -/// is intentional). `pooler.*` stays unused (masked-mean pooling instead). +/// is intentional; `PytorchStore` applies that adapter internally). +/// `pooler.*` stays unused (masked-mean pooling instead). pub fn load_from_dir( dir: &Path, device: &B::Device, ) -> Result, ImportError> { let cfg_path = required_file(dir, "config.json")?; - let weights = required_file(dir, "model.safetensors")?; let cfg_bytes = std::fs::read(&cfg_path).map_err(|e| ImportError::Parse { file: cfg_path.clone(), reason: e.to_string(), @@ -139,36 +173,27 @@ pub fn load_from_dir( let mut model = build::(&config, device); let target_float = Tensor::::zeros([1], device).dtype(); - let mut store = SafetensorsStore::from_file(weights.clone()) - .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) - .allow_partial(true) - .with_key_remapping(r"^bert\.", "") // some checkpoints prefix `bert.` - .with_key_remapping(r"^embeddings\.LayerNorm\.", "embeddings.layer_norm.") - .with_key_remapping( - r"^encoder\.layer\.(\d+)\.attention\.self\.(query|key|value)\.", - "layers.$1.$2.", - ) - .with_key_remapping( - r"^encoder\.layer\.(\d+)\.attention\.output\.dense\.", - "layers.$1.attn_output.", - ) - .with_key_remapping( - r"^encoder\.layer\.(\d+)\.attention\.output\.LayerNorm\.", - "layers.$1.attn_layer_norm.", - ) - .with_key_remapping( - r"^encoder\.layer\.(\d+)\.intermediate\.dense\.", - "layers.$1.intermediate.", - ) - .with_key_remapping( - r"^encoder\.layer\.(\d+)\.output\.dense\.", - "layers.$1.output.", - ) - .with_key_remapping( - r"^encoder\.layer\.(\d+)\.output\.LayerNorm\.", - "layers.$1.output_layer_norm.", - ); - load_checked(&mut model, &mut store, &weights)?; + match weights_file(dir)? { + WeightsFile::Safetensors(weights) => { + let mut store = SafetensorsStore::from_file(weights.clone()) + .with_from_adapter(PyTorchToBurnAdapter.chain(CastFloatAdapter::new(target_float))) + .allow_partial(true); + for (pattern, replacement) in KEY_REMAPS { + store = store.with_key_remapping(*pattern, *replacement); + } + load_checked(&mut model, &mut store, &weights)?; + } + WeightsFile::PytorchBin(weights) => { + // No cast adapter on this path (PytorchStore has no adapter + // chaining) — .bin-era checkpoints are f32, which every backend + // ingests directly. + let mut store = PytorchStore::from_file(weights.clone()).allow_partial(true); + for (pattern, replacement) in KEY_REMAPS { + store = store.with_key_remapping(*pattern, *replacement); + } + load_checked(&mut model, &mut store, &weights)?; + } + } Ok(LoadedMiniLm { model, config }) } diff --git a/crates/mummu/tests/real_pytorch.rs b/crates/mummu/tests/real_pytorch.rs new file mode 100644 index 0000000..1c5aa2b --- /dev/null +++ b/crates/mummu/tests/real_pytorch.rs @@ -0,0 +1,75 @@ +//! Real-weights PyTorch state-dict import proof: fetch MiniLM's original +//! `pytorch_model.bin` from the Hub, load it through `PytorchStore`, and +//! prove the embedding matches the safetensors-loaded model on the SAME +//! sentence — identical weights through two formats must agree. Ignored by +//! default (network, ~90 MB per format); run with +//! +//! ```text +//! MUMMU_HUB_DEST=some/tmp/dir cargo test -p mummu --release --test real_pytorch -- --ignored --nocapture +//! ``` + +use std::path::PathBuf; + +use mummu::backend::Cpu; +use mummu::hub; +use mummu::models::minilm; +use tokenizers::Tokenizer; + +const REPO: &str = "sentence-transformers/all-MiniLM-L6-v2"; +const SENTENCE: &str = "Two formats, one set of weights."; + +fn embed_from(dir: &std::path::Path) -> Vec { + let device = burn::tensor::Device::::default(); + let loaded = minilm::load_from_dir::(dir, &device).expect("checked load"); + let tok = Tokenizer::from_file(dir.join("tokenizer.json")).expect("tokenizer loads"); + let enc = tok.encode(SENTENCE, true).expect("encodes"); + let mask: Vec = enc.get_attention_mask().iter().map(|&m| m as f32).collect(); + loaded + .embed_ids(enc.get_ids(), &mask, &device) + .expect("embeds") +} + +#[test] +#[ignore = "needs network (MUMMU_HUB_DEST names the download dir)"] +fn pytorch_bin_load_matches_safetensors_load() { + 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 ~180 MB of downloads"); + }; + + // Reference: the safetensors copy (the catalog dir other tests share). + let st_dir = dest.join("all-minilm-l6-v2"); + hub::fetch_model(REPO, "main", &st_dir, |_| {}).expect("safetensors fetch"); + + // Candidate: a dir holding ONLY the PyTorch state dict (+ config & + // tokenizer), so weights_file() must take the .bin path. + let pt_dir = dest.join("minilm-pt"); + for f in ["config.json", "tokenizer.json", "pytorch_model.bin"] { + hub::fetch_file(&hub::hub_file_url(REPO, "main", f), &pt_dir.join(f), |_| {}) + .expect("pt fetch"); + } + assert!( + !pt_dir.join("model.safetensors").exists(), + "candidate dir must not contain safetensors — the test would prove nothing" + ); + + let reference = embed_from(&st_dir); + let candidate = embed_from(&pt_dir); + assert_eq!(reference.len(), 384); + assert_eq!(candidate.len(), 384); + + let max_abs_diff = reference + .iter() + .zip(&candidate) + .map(|(a, b)| (a - b).abs()) + .fold(0.0_f32, f32::max); + let cosine: f32 = reference.iter().zip(&candidate).map(|(a, b)| a * b).sum(); + eprintln!( + "[real_pytorch] safetensors vs pytorch_model.bin: cosine {cosine:.8}, max |Δ| {max_abs_diff:e}" + ); + // Same f32 weights through two containers on one backend: numerically + // identical up to load-order noise — far tighter than any semantic bound. + assert!( + max_abs_diff < 1e-6, + "formats diverge: max |Δcomponent| = {max_abs_diff}" + ); +} From 50cdd3b5d2235cc3098a7f103dd174f9bc529b25 Mon Sep 17 00:00:00 2001 From: Justin Icenhour Date: Sat, 11 Jul 2026 07:43:10 -0500 Subject: [PATCH 6/6] =?UTF-8?q?feat(p6):=20device=20inventory=20=E2=80=94?= =?UTF-8?q?=20CPU=20cores/RAM=20+=20per-adapter=20max=20buffer=20size?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The P6 planner's device set fills out: DeviceInventory gains CpuInfo (logical cores via available_parallelism; total RAM via GlobalMemoryStatusEx on Windows and /proc/meminfo on Linux, one syscall at inventory time) and each GpuAdapter now records max_buffer_bytes (adapter.limits) — the hard per-tensor allocation bound placement must respect. True per-adapter VRAM capacity is not portably exposed by wgpu; per-API wgpu-hal queries stay a noted P6 follow-up. Dev box readout: 4070 Ti SUPER 4 GiB max buffer on Vulkan / 2 GiB on DX12 (SHADER_F16 true/false respectively), integrated Radeon 2 GiB, CPU 32 logical cores / 127 GiB RAM. 93 unit tests green (2 new), clippy clean. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 1 + README.md | 3 +- ROADMAP.md | 6 ++ crates/mummu/Cargo.toml | 5 ++ crates/mummu/src/backend.rs | 136 ++++++++++++++++++++++++++++++++++-- 5 files changed, 146 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4f26f44..4649f66 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3556,6 +3556,7 @@ dependencies = [ "tokenizers", "ureq 3.3.0", "wgpu", + "windows-sys 0.61.2", ] [[package]] diff --git a/README.md b/README.md index 0e372dc..af78ea0 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,8 @@ It exists because two local-first apps — **[laurelane](https://github.com/phys - **Workspace + backends** — `crates/mummu` (library) + `crates/mummu-bench` (criterion); one binary compiles both `Wgpu` (with `fusion` + `autotune`) and `burn-flex` (CPU), with a cached runtime GPU probe and a - device inventory that records per-adapter/per-API `SHADER_F16`. + device inventory that records per-adapter/per-API `SHADER_F16` and max buffer size, plus the host + CPU's cores and total RAM — the planner's (and settings UIs') device set. - **Shared blocks, generic over `B: Backend`** — cache-aware GQA attention (optional per-head q/k RMSNorm), manual RoPE, SwiGLU, and LFM2's double-gated causal short-conv with rolling decode state; unit tests prove prefill+decode ≡ full-forward for both cache kinds. diff --git a/ROADMAP.md b/ROADMAP.md index 554da34..77cf327 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -241,6 +241,12 @@ The "use all the hardware" phase — inventory the machine, then pick the precis that fits the model AND uses every device to the fullest. - [ ] **Device inventory** — enumerate every GPU (`wgpu` adapters: name, backend, VRAM) and the CPU (cores, RAM); a stable device set cached at startup, reported so the apps can show it in settings. + *(2026-07-11)* Everything but true VRAM shipped: `DeviceInventory` now carries per-adapter + `max_buffer_bytes` (the planner's hard per-tensor bound; dev box: 4 GiB Vulkan / 2 GiB DX12) and + `CpuInfo` (logical cores + total RAM — `GlobalMemoryStatusEx` on Windows, `/proc/meminfo` on + Linux; dev box: 32 cores / 127 GiB). Remaining: per-adapter VRAM capacity, which wgpu does not + expose portably — needs per-API `wgpu-hal` queries (Vulkan memory heaps / DXGI) — and a macOS RAM + sysctl. - [ ] **Precision selection** — pick a per-device dtype (f32 / **f16** / int8 / int4) that fits: f16 via `Wgpu`; drop to int8/int4 (P9) when f16 still won't fit. *(2026-07-11) The f16 backend itself is now **fully validated** (all 3 claims — see the islands item below); what remains diff --git a/crates/mummu/Cargo.toml b/crates/mummu/Cargo.toml index 4aded66..4050762 100644 --- a/crates/mummu/Cargo.toml +++ b/crates/mummu/Cargo.toml @@ -20,3 +20,8 @@ sha2 = { workspace = true } once_cell = { workspace = true } burn-flex = { workspace = true } burn-store = { workspace = true } + +[target.'cfg(windows)'.dependencies] +# One syscall (GlobalMemoryStatusEx) for the device inventory's total-RAM +# figure; the narrow feature keeps the compile cost near zero. +windows-sys = { version = "0.61", features = ["Win32_System_SystemInformation"] } diff --git a/crates/mummu/src/backend.rs b/crates/mummu/src/backend.rs index 5b7da37..e40be81 100644 --- a/crates/mummu/src/backend.rs +++ b/crates/mummu/src/backend.rs @@ -37,15 +37,42 @@ pub struct GpuAdapter { pub device_type: wgpu::DeviceType, /// Does this adapter advertise `SHADER_F16` (native f16 shader arithmetic)? pub shader_f16: bool, + /// Largest single buffer this adapter permits — a hard bound the placement + /// planner respects per tensor/shard. (True VRAM capacity is NOT exposed + /// portably by wgpu; querying it per-API via wgpu-hal is a P6 follow-up.) + pub max_buffer_bytes: u64, } -/// Every hardware GPU visible to wgpu, enumerated once per process. +/// The host CPU as a compute device (the `burn-flex` target and the P6 +/// offload pool). +#[derive(Debug, Clone)] +pub struct CpuInfo { + /// Logical cores (SMT threads) available to this process; at least 1. + pub logical_cores: usize, + /// Total physical RAM; `None` where no query is implemented yet (macOS). + pub total_ram_bytes: Option, +} + +impl Default for CpuInfo { + fn default() -> Self { + Self { + logical_cores: 1, + total_ram_bytes: None, + } + } +} + +/// Every hardware GPU visible to wgpu plus the host CPU, enumerated once per +/// process — the device set the P6 hardware planner (and app settings UIs) +/// read. #[derive(Debug, Clone, Default)] pub struct DeviceInventory { /// Hardware adapters across the primary graphics APIs. The same physical /// card appears once per API that exposes it (e.g. Vulkan AND DX12) — /// deliberate, because features like `SHADER_F16` differ per API. pub gpus: Vec, + /// The host CPU (cores + RAM). + pub cpu: CpuInfo, } impl DeviceInventory { @@ -78,11 +105,73 @@ fn enumerate(instance: &wgpu::Instance, backends: wgpu::Backends) -> Vec Option { + use windows_sys::Win32::System::SystemInformation::{GlobalMemoryStatusEx, MEMORYSTATUSEX}; + let mut status = MEMORYSTATUSEX { + dwLength: core::mem::size_of::() as u32, + dwMemoryLoad: 0, + ullTotalPhys: 0, + ullAvailPhys: 0, + ullTotalPageFile: 0, + ullAvailPageFile: 0, + ullTotalVirtual: 0, + ullAvailVirtual: 0, + ullAvailExtendedVirtual: 0, + }; + // SAFETY: `status` is a live, writable MEMORYSTATUSEX with dwLength set, + // exactly what the API contract requires. + let ok = unsafe { GlobalMemoryStatusEx(&mut status) }; + (ok != 0).then_some(status.ullTotalPhys) +} + +#[cfg(target_os = "linux")] +fn total_ram_bytes() -> Option { + let meminfo = std::fs::read_to_string("/proc/meminfo").ok()?; + let kib: u64 = meminfo + .lines() + .find(|l| l.starts_with("MemTotal:"))? + .split_whitespace() + .nth(1)? + .parse() + .ok()?; + Some(kib * 1024) +} + +#[cfg(not(any(windows, target_os = "linux")))] +fn total_ram_bytes() -> Option { + None // macOS et al.: a sysctl query is a P6 follow-up +} + +/// The host CPU: logical cores + total RAM. +fn cpu_info() -> CpuInfo { + let logical_cores = std::thread::available_parallelism() + .map(std::num::NonZero::get) + .unwrap_or(1); + let total_ram_bytes = total_ram_bytes(); + assert!( + logical_cores >= 1, + "a running process has at least one core" + ); + // Negative space: an answer below 64 MiB is a parse/API bug, not a machine. + debug_assert!( + total_ram_bytes.is_none_or(|b| b >= 64 << 20), + "implausible total RAM: {total_ram_bytes:?}" + ); + CpuInfo { + logical_cores, + total_ram_bytes, + } +} + /// The process-lifetime device inventory. Enumerated once (first call pays /// ~tens of milliseconds); every later call is a cache read. pub fn inventory() -> &'static DeviceInventory { @@ -90,7 +179,10 @@ pub fn inventory() -> &'static DeviceInventory { INVENTORY.get_or_init(|| { let instance = wgpu::Instance::default(); let gpus = enumerate(&instance, wgpu::Backends::PRIMARY); - let inv = DeviceInventory { gpus }; + let inv = DeviceInventory { + gpus, + cpu: cpu_info(), + }; // Positive space: every inventoried adapter is real hardware. debug_assert!( inv.gpus @@ -164,15 +256,51 @@ mod tests { } } + #[test] + fn cpu_inventory_reports_cores_and_plausible_ram() { + let cpu = &inventory().cpu; + assert!(cpu.logical_cores >= 1); + // Windows and Linux have a RAM query; its answer must be a real + // machine's (1 GiB ..= 64 TiB), not a unit slip. + if cfg!(any(windows, target_os = "linux")) { + let ram = cpu.total_ram_bytes.expect("RAM query exists here"); + assert!((1 << 30..=1u64 << 46).contains(&ram), "implausible: {ram}"); + } + } + + #[test] + fn adapters_report_a_usable_buffer_bound() { + // Every real adapter permits at least the WebGPU floor (256 MiB); + // a smaller answer means the limits plumbing broke. + for gpu in &inventory().gpus { + assert!( + gpu.max_buffer_bytes >= 256 << 20, + "{}: max_buffer_bytes {} below the WebGPU floor", + gpu.name, + gpu.max_buffer_bytes + ); + } + } + /// Print the inventory so the nightly log records what this machine has. #[test] fn report_inventory() { for gpu in &inventory().gpus { eprintln!( - "[mummu] {:?} / {} ({:?}): SHADER_F16 = {}", - gpu.backend, gpu.name, gpu.device_type, gpu.shader_f16 + "[mummu] {:?} / {} ({:?}): SHADER_F16 = {}, max buffer {:.1} GiB", + gpu.backend, + gpu.name, + gpu.device_type, + gpu.shader_f16, + gpu.max_buffer_bytes as f64 / f64::from(1u32 << 30), ); } + let cpu = &inventory().cpu; + eprintln!( + "[mummu] CPU: {} logical cores, RAM {:?} GiB", + cpu.logical_cores, + cpu.total_ram_bytes.map(|b| b >> 30), + ); eprintln!("[mummu] policy: {}", device_label()); } }