diff --git a/CHANGELOG.md b/CHANGELOG.md index 7ce8011..e498db5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,14 @@ All notable changes to this project are documented in this file. ### Added +- **MoE extraction on mmap-backed GGUF ([RM-1462](https://linear.app/rpd-34/issue/RM-1462)):** + [`GgufTensorAccess`] shares tensor-directory and payload access between + [`GgufLayout`] and [`GgufLayoutMmap`]. [`list_experts_from`] / + [`extract_expert_from`] (and `list_experts_mmap` / `extract_expert_mmap` + when the `mmap` feature is enabled) reuse the same stacked and per-expert + parsing as the existing owned-layout helpers. Expert payloads are still + copied into owned `Vec` buffers on return. Integration tests assert + owned/mmap parity on synthetic fixtures. - **Packed dequant for Q4_0 / Q4_1 / Q4_K / Q3_K / Q2_K (#125):** `dequantize_q4_0`, `dequantize_q4_1`, `dequantize_q2_k`, `dequantize_q3_k`, `dequantize_q4_k`, wired into `dequantize_packed` and `packed_row_size` diff --git a/src/analysis/moe/extract.rs b/src/analysis/moe/extract.rs index 9e01d05..dfe2733 100644 --- a/src/analysis/moe/extract.rs +++ b/src/analysis/moe/extract.rs @@ -1,6 +1,6 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 -//! Per-expert weight extraction from a parsed [`GgufLayout`]. +//! Per-expert weight extraction from a parsed GGUF layout (owned or mmap). //! //! Two on-disk conventions are supported: //! @@ -12,29 +12,50 @@ //! `blk.{b}.ffn_{role}.{e}.weight`. Whole-tensor payload is cloned. //! //! No arithmetic is performed — only byte-level slicing plus a `Vec` -//! clone so the returned [`MoeExpertWeights`] is self-owning. +//! clone so the returned [`MoeExpertWeights`] is self-owning. Payloads +//! are always copied into `RawTensor::bytes` even when the layout borrows +//! from an mmap mapping (callers get an owned buffer; see copy policy on +//! [`extract_expert_from`]). use std::collections::BTreeSet; use super::expert::{MoeExpertWeights, RawTensor}; use crate::error::{ParserError, Result}; -use crate::gguf::{GgufLayout, Tensor}; +use crate::gguf::{GgufLayout, GgufTensorAccess, Tensor}; + +/// Enumerate every `(block, expert)` pair in an owned [`GgufLayout`]. +/// +/// Convenience wrapper around [`list_experts_from`]. +pub fn list_experts(layout: &GgufLayout) -> Vec<(usize, usize)> { + list_experts_from(layout) +} /// Enumerate every `(block, expert)` pair present in the checkpoint. /// /// A pair is reported if *any* of the three projections (`gate`, `up`, /// `down`) is discoverable for it under either the stacked or the /// per-expert naming scheme. -pub fn list_experts(layout: &GgufLayout) -> Vec<(usize, usize)> { +/// +/// Works for any [`GgufTensorAccess`] implementation (owned buffer or +/// `mmap` mapping). +pub fn list_experts_from(layout: &L) -> Vec<(usize, usize)> { let mut pairs: BTreeSet<(usize, usize)> = BTreeSet::new(); - for tensor in layout.tensors.values() { + for tensor in layout.gguf_tensors().values() { record_expert_pairs(tensor, &mut pairs); } pairs.into_iter().collect() } +#[cfg(feature = "mmap")] +/// Enumerate MoE `(block, expert)` pairs from a memory-mapped layout. +/// +/// Identical to [`list_experts_from`] on [`crate::GgufLayoutMmap`]. +pub fn list_experts_mmap(layout: &crate::GgufLayoutMmap) -> Vec<(usize, usize)> { + list_experts_from(layout) +} + fn record_expert_pairs(tensor: &Tensor, pairs: &mut BTreeSet<(usize, usize)>) { if let Some((block, _role)) = parse_stacked_name(&tensor.name) { if let Some(n_experts) = stacked_expert_count(tensor) { @@ -50,14 +71,31 @@ fn record_expert_pairs(tensor: &Tensor, pairs: &mut BTreeSet<(usize, usize)>) { } } +/// Extract the raw weights for a single `(block, expert)` pair from an +/// owned [`GgufLayout`]. +/// +/// Convenience wrapper around [`extract_expert_from`]. +pub fn extract_expert( + layout: &GgufLayout, + block: usize, + expert: usize, +) -> Result { + extract_expert_from(layout, block, expert) +} + /// Extract the raw weights for a single `(block, expert)` pair. /// /// Returns [`MoeExpertWeights`] whose `gate / up / down` fields are /// `None` when that projection was not found in the checkpoint. If /// *none* of the three projections can be located, returns /// [`ParserError::MissingTensor`]. -pub fn extract_expert( - layout: &GgufLayout, +/// +/// **Copy policy:** tensor payloads are always cloned into owned +/// `Vec` in the returned [`RawTensor`] values, including stacked +/// slices and mmap-backed sources. The layout is only borrowed for +/// directory lookup and byte views during the call. +pub fn extract_expert_from( + layout: &L, block: usize, expert: usize, ) -> Result { @@ -68,7 +106,7 @@ pub fn extract_expert( if gate.is_none() && up.is_none() && down.is_none() { return Err(ParserError::MissingTensor { name: format!("blk.{block}.ffn_(gate|up|down)[_exps].weight"), - path: layout.path.clone(), + path: layout.gguf_path().to_owned(), }); } @@ -81,15 +119,27 @@ pub fn extract_expert( }) } -fn extract_role( - layout: &GgufLayout, +#[cfg(feature = "mmap")] +/// Extract one MoE expert from a memory-mapped layout. +/// +/// Identical to [`extract_expert_from`] on [`crate::GgufLayoutMmap`]. +pub fn extract_expert_mmap( + layout: &crate::GgufLayoutMmap, + block: usize, + expert: usize, +) -> Result { + extract_expert_from(layout, block, expert) +} + +fn extract_role( + layout: &L, block: usize, expert: usize, role: &str, ) -> Result> { // Prefer the stacked convention first. let stacked_name = format!("blk.{block}.ffn_{role}_exps.weight"); - if let Some(tensor) = layout.tensors.get(&stacked_name) { + if let Some(tensor) = layout.gguf_tensors().get(&stacked_name) { return Ok(Some(slice_stacked_expert(layout, block, tensor, expert)?)); } @@ -100,7 +150,7 @@ fn extract_role( format!("blk.{block}.ffn_{role}.{expert}.weight"), format!("blk.{block}.ffn_{role}_{expert}.weight"), ] { - if let Some(tensor) = layout.tensors.get(&candidate) { + if let Some(tensor) = layout.gguf_tensors().get(&candidate) { let bytes = layout.tensor_bytes(tensor)?.to_vec(); return Ok(Some(RawTensor { source_name: tensor.name.clone(), @@ -123,8 +173,8 @@ fn extract_role( /// dimension in the `dims` vector. The per-expert chunk consists of /// the first `n_elements / n_experts` elements per expert, laid out /// contiguously in the tensor buffer. -fn slice_stacked_expert( - layout: &GgufLayout, +fn slice_stacked_expert( + layout: &L, block: usize, tensor: &Tensor, expert: usize, @@ -144,14 +194,14 @@ fn slice_stacked_expert( }) } -fn validate_stacked_expert_count( - layout: &GgufLayout, +fn validate_stacked_expert_count( + layout: &L, block: usize, tensor: &Tensor, expert: usize, ) -> Result { let n_experts = stacked_expert_count(tensor).ok_or_else(|| ParserError::InvalidLayout { - path: layout.path.clone(), + path: layout.gguf_path().to_owned(), reason: format!( "stacked tensor '{}' has rank {} (dims={:?}); expected rank >= 2", tensor.name, @@ -162,7 +212,7 @@ fn validate_stacked_expert_count( if n_experts == 0 { return Err(ParserError::InvalidLayout { - path: layout.path.clone(), + path: layout.gguf_path().to_owned(), reason: format!("stacked tensor '{}' has zero experts", tensor.name), }); } @@ -178,8 +228,8 @@ fn validate_stacked_expert_count( Ok(n_experts) } -fn stacked_slice_range( - layout: &GgufLayout, +fn stacked_slice_range( + layout: &L, tensor: &Tensor, expert: usize, n_experts: usize, @@ -187,7 +237,7 @@ fn stacked_slice_range( ) -> Result<(usize, usize)> { if !tensor.byte_len.is_multiple_of(n_experts) { return Err(ParserError::InvalidLayout { - path: layout.path.clone(), + path: layout.gguf_path().to_owned(), reason: format!( "stacked tensor '{}' byte length {} is not divisible by n_experts={n_experts}", tensor.name, tensor.byte_len, @@ -199,18 +249,18 @@ fn stacked_slice_range( let start = expert .checked_mul(stride) .ok_or_else(|| ParserError::InvalidLayout { - path: layout.path.clone(), + path: layout.gguf_path().to_owned(), reason: format!("stacked stride overflow for tensor '{}'", tensor.name), })?; let end = start .checked_add(stride) .ok_or_else(|| ParserError::InvalidLayout { - path: layout.path.clone(), + path: layout.gguf_path().to_owned(), reason: format!("stacked end overflow for tensor '{}'", tensor.name), })?; if end > buffer_len { return Err(ParserError::InvalidLayout { - path: layout.path.clone(), + path: layout.gguf_path().to_owned(), reason: format!( "stacked slice range [{start}..{end}] for tensor '{}' exceeds buffer len {buffer_len}", tensor.name, diff --git a/src/analysis/moe/mod.rs b/src/analysis/moe/mod.rs index 6b3e8e1..0105409 100644 --- a/src/analysis/moe/mod.rs +++ b/src/analysis/moe/mod.rs @@ -2,16 +2,24 @@ //! Mixture-of-Experts weight extraction (GGUF-specific). //! -//! Locates MoE expert tensors inside a parsed [`GgufLayout`](crate::gguf::GgufLayout) -//! and returns raw per-expert byte buffers. Performs no neural-network -//! math — callers receive `Vec` + shape/dtype metadata and are -//! responsible for any downstream compute. +//! Locates MoE expert tensors inside a parsed GGUF layout and returns raw +//! per-expert byte buffers. Performs no neural-network math — callers receive +//! `Vec` + shape/dtype metadata and are responsible for any downstream +//! compute. +//! +//! Owned checkpoints use [`list_experts`] / [`extract_expert`] on +//! [`GgufLayout`](crate::gguf::GgufLayout). Mmap-backed checkpoints use +//! the same helpers via [`list_experts_from`] / [`extract_expert_from`] +//! (or the `*_mmap` aliases when the `mmap` feature is enabled) on +//! [`GgufLayoutMmap`](crate::gguf::GgufLayoutMmap). Both paths share +//! [`GgufTensorAccess`](crate::gguf::GgufTensorAccess) so parsing cannot +//! drift between layout types. //! //! # Scope //! //! This is a model-analysis specialization, not part of the core -//! checkpoint contract: [`extract_expert`] and [`list_experts`] take a -//! concrete `&GgufLayout`, and [`RawTensor`] carries the GGUF-specific +//! checkpoint contract: entry points take a concrete GGUF layout type or +//! [`GgufTensorAccess`], and [`RawTensor`] carries the GGUF-specific //! [`DType`](crate::gguf::DType) / `ggml_type` fields of its source //! tensor. It does not accept [`Checkpoint`](crate::checkpoint::Checkpoint) //! — a format-independent MoE/topology layer is planned v0.4 work. @@ -32,4 +40,6 @@ mod expert; mod extract; pub use expert::{MoeExpertWeights, RawTensor}; -pub use extract::{extract_expert, list_experts}; +pub use extract::{extract_expert, extract_expert_from, list_experts, list_experts_from}; +#[cfg(feature = "mmap")] +pub use extract::{extract_expert_mmap, list_experts_mmap}; diff --git a/src/checkpoint/open.rs b/src/checkpoint/open.rs index e51417b..a1be482 100644 --- a/src/checkpoint/open.rs +++ b/src/checkpoint/open.rs @@ -56,7 +56,10 @@ impl AnyCheckpoint { } } - /// Memory-mapped GGUF layout. + /// Memory-mapped GGUF layout for + /// [`crate::analysis::moe::list_experts_from`] / + /// [`crate::analysis::moe::extract_expert_from`] (or the `*_mmap` + /// aliases). #[cfg(feature = "mmap")] pub fn as_gguf_mmap(&self) -> Option<&crate::GgufLayoutMmap> { match self { diff --git a/src/gguf/access.rs b/src/gguf/access.rs new file mode 100644 index 0000000..12b9810 --- /dev/null +++ b/src/gguf/access.rs @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +//! Shared tensor-directory + payload access for owned and mmap GGUF layouts. + +use std::collections::HashMap; + +use super::layout::GgufLayout; +use super::tensor::Tensor; +use crate::error::Result; + +/// Tensor directory and payload access shared by [`GgufLayout`] and +/// [`GgufLayoutMmap`](super::map::GgufLayoutMmap). +/// +/// MoE helpers ([`crate::analysis::moe::list_experts_from`], +/// [`crate::analysis::moe::extract_expert_from`]) and other analysis code +/// use this trait so owned and mmap-backed checkpoints cannot drift. +pub trait GgufTensorAccess { + /// Source path label (filesystem path or synthetic label). + fn gguf_path(&self) -> &str; + + /// Parsed tensor directory (`name` -> [`Tensor`] metadata). + fn gguf_tensors(&self) -> &HashMap; + + /// Borrowed raw payload bytes for `tensor`. + fn tensor_bytes<'a>(&'a self, tensor: &Tensor) -> Result<&'a [u8]>; +} + +impl GgufTensorAccess for GgufLayout { + fn gguf_path(&self) -> &str { + &self.path + } + + fn gguf_tensors(&self) -> &HashMap { + &self.tensors + } + + fn tensor_bytes<'a>(&'a self, tensor: &Tensor) -> Result<&'a [u8]> { + GgufLayout::tensor_bytes(self, tensor) + } +} + +#[cfg(feature = "mmap")] +impl GgufTensorAccess for super::map::GgufLayoutMmap { + fn gguf_path(&self) -> &str { + &self.path + } + + fn gguf_tensors(&self) -> &HashMap { + &self.tensors + } + + fn tensor_bytes<'a>(&'a self, tensor: &Tensor) -> Result<&'a [u8]> { + super::map::GgufLayoutMmap::tensor_bytes(self, tensor) + } +} diff --git a/src/gguf/map.rs b/src/gguf/map.rs index 5993982..73a41d6 100644 --- a/src/gguf/map.rs +++ b/src/gguf/map.rs @@ -22,7 +22,10 @@ use crate::error::{ParserError, Result}; /// GGUF layout backed by a read-only file mapping instead of an owned `Vec`. /// /// Header, KV metadata, and the tensor directory are identical to -/// [`GgufLayout`]. Tensor payloads are borrowed from the mapping. +/// [`GgufLayout`]. Tensor payloads are borrowed from the mapping. MoE +/// expert listing and extraction use [`super::access::GgufTensorAccess`] +/// via [`crate::analysis::moe::list_experts_from`] / +/// [`crate::analysis::moe::extract_expert_from`]. pub struct GgufLayoutMmap { /// Path the checkpoint was mapped from (kept for error messages). pub path: String, diff --git a/src/gguf/mod.rs b/src/gguf/mod.rs index 956ffbb..3b656af 100644 --- a/src/gguf/mod.rs +++ b/src/gguf/mod.rs @@ -7,6 +7,7 @@ //! byte buffer. Packed K-quant dequant helpers are re-exported from this module. //! Optional mmap-backed loading is behind the `mmap` feature. +mod access; mod cursor; mod dequant; mod layout; @@ -18,6 +19,7 @@ mod tensor; use std::fs; use std::path::Path; +pub use access::GgufTensorAccess; pub use dequant::{ dequantize_iq3_m, dequantize_packed, dequantize_q2_k, dequantize_q3_k, dequantize_q4_0, dequantize_q4_1, dequantize_q4_k, dequantize_q5_k, dequantize_q6_k, dequantize_q8_0, diff --git a/src/lib.rs b/src/lib.rs index 722cc89..339f33e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -69,7 +69,10 @@ //! ``` //! //! MoE expert analysis (GGUF-specific) uses the canonical -//! [`analysis::moe`] path: +//! [`analysis::moe`] path. Owned loads use [`load_gguf`]; with the +//! `mmap` feature, [`load_gguf_mmap`] plus [`analysis::moe::list_experts_from`] +//! / [`analysis::moe::extract_expert_from`] avoid copying the file into a +//! `Vec` (payloads are still cloned into owned expert buffers). //! //! ```no_run //! use engram_parser::analysis::moe::{extract_expert, list_experts}; @@ -160,6 +163,7 @@ pub use gguf::{ GGUF_VALUE_TYPE_UINT64, GgufLayout, GgufMetadata, + GgufTensorAccess, ParseLimits, Tensor, bf16_bits_to_f32, @@ -189,4 +193,9 @@ pub use gguf::{ }; // MoE analysis result types and extractors, re-exported for // compatibility; the canonical path is `analysis::moe`. -pub use analysis::moe::{MoeExpertWeights, RawTensor, extract_expert, list_experts}; +pub use analysis::moe::{ + MoeExpertWeights, RawTensor, extract_expert, extract_expert_from, list_experts, + list_experts_from, +}; +#[cfg(feature = "mmap")] +pub use analysis::moe::{extract_expert_mmap, list_experts_mmap}; diff --git a/tests/mmap_gguf.rs b/tests/mmap_gguf.rs index ec13b53..f9d6b98 100644 --- a/tests/mmap_gguf.rs +++ b/tests/mmap_gguf.rs @@ -13,7 +13,8 @@ use common::*; use engram_parser::{ DType, ParserError, dequantize_iq3_m, dequantize_packed, dequantize_q5_k, dequantize_q6_k, - dequantize_q8_0, load_gguf, load_gguf_mmap, os_page_size, parse_bytes, + dequantize_q8_0, extract_expert, extract_expert_from, extract_expert_mmap, list_experts, + list_experts_from, list_experts_mmap, load_gguf, load_gguf_mmap, os_page_size, parse_bytes, }; use std::fs; #[cfg(unix)] @@ -417,3 +418,94 @@ impl Drop for TempGuard { let _ = fs::remove_file(&self.0); } } + +#[test] +fn mmap_moe_extraction_matches_owned_stacked_and_per_expert() { + let inner = 4usize; + let outer = 2usize; + let n_experts = 3usize; + let per_expert = inner * outer; + + let mut gate = vec![0.0f32; n_experts * per_expert]; + for e in 0..n_experts { + let base = e * per_expert; + for slot in gate.iter_mut().skip(base).take(per_expert) { + *slot = (e as f32) + 0.1; + } + } + let stacked_dims = vec![inner, outer, n_experts]; + let stacked_kv = [ + ("general.alignment", KvValue::U32(ALIGNMENT)), + ("general.architecture", KvValue::Str("olmoe")), + ]; + let stacked_tensors = [TensorSpec { + name: "blk.0.ffn_gate_exps.weight", + dims: stacked_dims, + ggml_type: GGML_F32, + payload: f32_vec_to_le_bytes(&gate), + }]; + let stacked_bytes = build_gguf(&stacked_kv, &stacked_tensors); + let stacked_path = write_temp_gguf(&stacked_bytes); + let _stacked_guard = TempGuard(stacked_path.clone()); + + let owned_stacked = parse_bytes(stacked_bytes, "mem://stacked".into()).expect("parse stacked"); + let mapped_stacked = load_gguf_mmap(&stacked_path).expect("mmap stacked"); + + assert_eq!( + list_experts(&owned_stacked), + list_experts_from(&mapped_stacked) + ); + assert_eq!( + list_experts(&owned_stacked), + list_experts_mmap(&mapped_stacked) + ); + for e in 0..n_experts { + let owned_w = extract_expert(&owned_stacked, 0, e).expect("owned stacked"); + let mapped_w = extract_expert_from(&mapped_stacked, 0, e).expect("trait stacked"); + let mapped_alias = extract_expert_mmap(&mapped_stacked, 0, e).expect("mmap alias stacked"); + assert_eq!( + owned_w.gate.as_ref().map(|g| g.bytes.clone()), + mapped_w.gate.as_ref().map(|g| g.bytes.clone()) + ); + assert_eq!( + mapped_w.gate.as_ref().map(|g| g.bytes.clone()), + mapped_alias.gate.as_ref().map(|g| g.bytes.clone()) + ); + assert!(mapped_w.gate.as_ref().is_some_and(|g| g.stacked_slice)); + } + + let per_gate = f32_vec_to_le_bytes(&[1.0, 2.0, 3.0, 4.0]); + let per_kv = stacked_kv; + let per_tensors = [ + TensorSpec { + name: "blk.0.ffn_gate.0.weight", + dims: vec![4], + ggml_type: GGML_F32, + payload: per_gate.clone(), + }, + TensorSpec { + name: "blk.0.ffn_gate_1.weight", + dims: vec![4], + ggml_type: GGML_F32, + payload: per_gate, + }, + ]; + let per_bytes = build_gguf(&per_kv, &per_tensors); + let per_path = write_temp_gguf(&per_bytes); + let _per_guard = TempGuard(per_path.clone()); + + let owned_per = parse_bytes(per_bytes, "mem://per".into()).expect("parse per"); + let mapped_per = load_gguf_mmap(&per_path).expect("mmap per"); + + assert_eq!(list_experts(&owned_per), vec![(0, 0), (0, 1)]); + assert_eq!(list_experts_from(&mapped_per), list_experts(&owned_per)); + for e in 0..2 { + let owned_w = extract_expert(&owned_per, 0, e).expect("owned per"); + let mapped_w = extract_expert_mmap(&mapped_per, 0, e).expect("mmap per"); + assert_eq!( + owned_w.gate.as_ref().map(|g| g.bytes.clone()), + mapped_w.gate.as_ref().map(|g| g.bytes.clone()) + ); + assert!(!mapped_w.gate.as_ref().is_some_and(|g| g.stacked_slice)); + } +}