zsfm_burn/
gguf_weights.rs

1//! Direct GGUF weight loading for Burn inference — no candle.
2//!
3//! Reads F32/F16/BF16/Q8_0 tensors with the in-workspace `zsfm_gguf`
4//! reader, decodes everything to F32 host bytes, and builds Burn tensors.
5
6use std::collections::HashMap;
7use std::io::BufReader;
8use std::path::Path;
9
10use anyhow::{Context, Result};
11use burn::tensor::{Device, Tensor, TensorData};
12use zsfm_checkpoint::cast::{decode_to_f32, f32_to_bytes, SrcDtype};
13use zsfm_gguf::{GGMLType, GGUFFile, GGUFMetaValue};
14
15/// All F32-decoded tensors of one GGUF file, keyed by tensor name.
16pub struct GgufWeights {
17    tensors: HashMap<String, (Vec<u64>, Vec<u8>)>,
18    meta: Vec<(String, GGUFMetaValue)>,
19}
20
21fn device() -> Device {
22    Device::flex()
23}
24
25/// Open `path` and decode every tensor to F32.
26pub fn open_gguf_weights(path: &Path) -> Result<GgufWeights> {
27    let file = std::fs::File::open(path).with_context(|| format!("open {}", path.display()))?;
28    let mut reader = BufReader::with_capacity(zsfm_gguf::READ_BUF_CAPACITY, file);
29    let gguf = GGUFFile::read(&mut reader).context("parse GGUF header")?;
30    let mut tensors = HashMap::with_capacity(gguf.tensors.len());
31    for info in &gguf.tensors {
32        // Stored dims are GGUF order (reversed row-major) — flip back.
33        let shape: Vec<u64> = info.shape.iter().rev().copied().collect();
34        let bytes = match info.dtype {
35            GGMLType::F32 | GGMLType::F16 | GGMLType::BF16 => {
36                gguf.tensor_bytes(&mut reader, info)?
37            }
38            GGMLType::Q8_0 => {
39                let vals = gguf.tensor_f32(&mut reader, info)?;
40                f32_to_bytes(&vals)
41            }
42        };
43        // Validate element count before accepting.
44        let n: usize = shape.iter().product::<u64>() as usize;
45        let src = match info.dtype {
46            GGMLType::F32 => SrcDtype::F32,
47            GGMLType::F16 => SrcDtype::F16,
48            GGMLType::BF16 => SrcDtype::BF16,
49            GGMLType::Q8_0 => SrcDtype::F32,
50        };
51        let expect = n * src.bytes_per_elem();
52        anyhow::ensure!(
53            bytes.len() == expect,
54            "tensor {}: {} bytes, want {expect}",
55            info.name,
56            bytes.len()
57        );
58        tensors.insert(info.name.clone(), (shape, bytes));
59    }
60    Ok(GgufWeights {
61        tensors,
62        meta: gguf.metadata,
63    })
64}
65
66impl GgufWeights {
67    fn get(&self, name: &str) -> Result<&(Vec<u64>, Vec<u8>)> {
68        self.tensors
69            .get(name)
70            .with_context(|| format!("tensor '{name}' not found in GGUF"))
71    }
72
73    fn f32_vec(&self, name: &str) -> Result<Vec<f32>> {
74        let (shape, bytes) = self.get(name)?;
75        let n: usize = shape.iter().product::<u64>() as usize;
76        let mut vals = decode_to_f32(bytes, SrcDtype::F32)?;
77        anyhow::ensure!(
78            vals.len() == n,
79            "tensor {name}: decoded {} values, want {n}",
80            vals.len()
81        );
82        // decode_to_f32 returns exactly n values; keep the binding for clarity.
83        let _ = n;
84        Ok(std::mem::take(&mut vals))
85    }
86
87    /// U32 metadata value (accepts u32/u64/i32/i64 spellings).
88    pub fn meta_u32(&self, key: &str) -> Option<u32> {
89        use GGUFMetaValue::*;
90        self.meta
91            .iter()
92            .find(|(k, _)| k == key)
93            .and_then(|(_, v)| match v {
94                Uint32(x) => Some(*x),
95                Uint64(x) => u32::try_from(*x).ok(),
96                Int32(x) => u32::try_from(*x).ok(),
97                Int64(x) => u32::try_from(*x).ok(),
98                _ => None,
99            })
100    }
101
102    /// Load a 2D `[rows, cols]` F32 tensor.
103    pub fn load_2d(&self, name: &str) -> Result<Tensor<2>> {
104        let (shape, _) = self.get(name)?;
105        anyhow::ensure!(shape.len() == 2, "tensor {name}: want 2D, got {shape:?}");
106        let data = self.f32_vec(name)?;
107        Ok(Tensor::<2>::from_data(
108            TensorData::new(data, [shape[0] as usize, shape[1] as usize]),
109            &device(),
110        ))
111    }
112
113    /// Load a 2D weight with the `load_weight` transpose guard: PyTorch stores
114    /// `(d_out, d_in)`; if the file holds it transposed (the Q8_0 transpose
115    /// trick some converters apply), flip it back.
116    pub fn load_weight_2d(&self, name: &str, expected_d_out: usize) -> Result<Tensor<2>> {
117        let (shape, _) = self.get(name)?;
118        anyhow::ensure!(shape.len() == 2, "tensor {name}: want 2D, got {shape:?}");
119        let (r, c) = (shape[0] as usize, shape[1] as usize);
120        let data = self.f32_vec(name)?;
121        if r == expected_d_out {
122            return Ok(Tensor::<2>::from_data(
123                TensorData::new(data, [r, c]),
124                &device(),
125            ));
126        }
127        anyhow::ensure!(
128            c == expected_d_out,
129            "tensor {name}: dim(0) is {r}, want {expected_d_out} (transpose fallback needs {c})"
130        );
131        let mut t = vec![0.0f32; data.len()];
132        for i in 0..r {
133            for j in 0..c {
134                t[j * r + i] = data[i * c + j];
135            }
136        }
137        Ok(Tensor::<2>::from_data(
138            TensorData::new(t, [c, r]),
139            &device(),
140        ))
141    }
142
143    /// Load a 1D `[n]` F32 tensor. Like the retired candle bridge
144    /// (`burn_weight_1d`), this flattens any stored shape: real checkpoints
145    /// sometimes store row vectors as `[1, n]`.
146    pub fn load_1d(&self, name: &str) -> Result<Tensor<1>> {
147        let (shape, _) = self.get(name)?;
148        let data = self.f32_vec(name)?;
149        let n: usize = shape.iter().product::<u64>() as usize;
150        anyhow::ensure!(
151            data.len() == n,
152            "tensor {name}: decoded {} values, want {n}",
153            data.len()
154        );
155        Ok(Tensor::<1>::from_data(
156            TensorData::new(data, [n]),
157            &device(),
158        ))
159    }
160
161    /// Load a 3D `[d0, d1, d2]` F32 tensor.
162    pub fn load_3d(&self, name: &str) -> Result<Tensor<3>> {
163        let (shape, _) = self.get(name)?;
164        anyhow::ensure!(shape.len() == 3, "tensor {name}: want 3D, got {shape:?}");
165        let data = self.f32_vec(name)?;
166        Ok(Tensor::<3>::from_data(
167            TensorData::new(
168                data,
169                [shape[0] as usize, shape[1] as usize, shape[2] as usize],
170            ),
171            &device(),
172        ))
173    }
174
175    /// Raw F32 values of any tensor (for embedding tables, rope vectors, …).
176    pub fn load_flat(&self, name: &str) -> Result<Vec<f32>> {
177        self.f32_vec(name)
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    #[test]
186    fn missing_tensor_errors() {
187        let w = GgufWeights {
188            tensors: HashMap::new(),
189            meta: Vec::new(),
190        };
191        assert!(w.load_2d("nope").is_err());
192    }
193}
zsfm_burn/gguf_weights.rs

zsfm_burn/
gguf_weights.rs
