ferrotensor_loader/
lib.rs

1//! Model loading for Ferrotensor: Hugging Face `config.json` parsing
2//! ([`Lfm2Config`]), memory-mapped safetensors weights (single file or
3//! sharded, F32/F16/BF16 read as F32, int-quantized weights dequantized on
4//! the GPU) through [`WeightsContext`], and a [`Tokenizer`] / [`Sampler`]
5//! pair for text in and out.
6
7#![warn(missing_docs)]
8
9use memmap2::MmapOptions;
10use safetensors::SafeTensors;
11use serde::Deserialize;
12use std::collections::HashMap;
13use std::fs::File;
14use std::path::Path;
15use std::sync::Arc;
16use thiserror::Error;
17
18pub mod tokenizer;
19pub use tokenizer::{argmax, Sampler, Tokenizer};
20
21use ferrotensor_backend::UnifiedTensor;
22
23/// Errors returned while loading model configs, weights and tensors.
24#[derive(Error, Debug)]
25pub enum LoaderError {
26    /// Underlying I/O failure.
27    #[error("IO Error: {0}")]
28    Io(#[from] std::io::Error),
29    /// Failure parsing a safetensors file.
30    #[error("Safetensors Error: {0}")]
31    Safetensors(#[from] safetensors::SafeTensorError),
32    /// Failure parsing a JSON file.
33    #[error("JSON Error: {0}")]
34    Json(#[from] serde_json::Error),
35    /// A requested tensor name is not present in any loaded file.
36    #[error("Weight not found: {0}")]
37    MissingWeight(String),
38    /// The tensor uses a dtype this loader cannot convert to f32.
39    #[error("Unsupported dtype for tensor {name}: {dtype}")]
40    UnsupportedDtype {
41        /// Tensor name.
42        name: String,
43        /// Dtype name as reported by safetensors.
44        dtype: String,
45    },
46    /// The tensor byte length does not match its shape and dtype.
47    #[error("Invalid tensor byte length for {name}: expected {expected} bytes, got {actual}")]
48    InvalidTensorLength {
49        /// Tensor name.
50        name: String,
51        /// Byte length implied by shape and dtype.
52        expected: usize,
53        /// Byte length found in the file.
54        actual: usize,
55    },
56}
57
58/// Quantization settings from the model config.
59#[derive(Debug, Deserialize, Clone)]
60pub struct QuantizationConfig {
61    /// Number of weights that share one scale and bias.
62    pub group_size: usize,
63    /// Bits per quantized weight.
64    pub bits: usize,
65    /// Quantization scheme name, e.g. `"affine"`.
66    pub mode: String,
67}
68
69/// Model hyperparameters read from `config.json` (Llama and LFM2 layouts).
70#[derive(Debug, Deserialize, Clone)]
71pub struct Lfm2Config {
72    /// Width of the residual stream.
73    pub hidden_size: usize,
74    /// Width of the feed-forward (MLP) hidden layer.
75    pub intermediate_size: usize,
76    /// Number of attention heads.
77    pub num_attention_heads: usize,
78    /// Number of transformer layers.
79    pub num_hidden_layers: usize,
80    /// Number of tokens in the vocabulary.
81    pub vocab_size: usize,
82    /// Epsilon added inside RMSNorm; defaults to 1e-5.
83    #[serde(default = "default_norm_eps", alias = "rms_norm_eps")]
84    pub norm_eps: f32,
85    /// Hidden width used by LFM2 blocks, if set.
86    #[serde(default)]
87    pub block_dim: Option<usize>,
88    /// Feed-forward width used by LFM2 blocks, if set.
89    #[serde(default)]
90    pub block_ff_dim: Option<usize>,
91    /// Per-layer block type names (LFM2); empty when not given.
92    #[serde(default)]
93    pub layer_types: Vec<String>,
94    /// Number of key/value heads for grouped-query attention, if set.
95    #[serde(default)]
96    pub num_key_value_heads: Option<usize>,
97    /// RoPE base frequency, if set.
98    #[serde(default)]
99    pub rope_theta: Option<f64>,
100    /// Quantization settings, present for quantized checkpoints.
101    #[serde(default)]
102    pub quantization: Option<QuantizationConfig>,
103    /// Beginning-of-sequence token id, if set.
104    #[serde(default)]
105    pub bos_token_id: Option<u32>,
106    /// End-of-sequence token id, if set.
107    #[serde(default)]
108    pub eos_token_id: Option<u32>,
109    /// Attention dropout probability, if set.
110    #[serde(default)]
111    pub attention_dropout: Option<f32>,
112    /// Hugging Face `model_type`, e.g. `"llama"` or `"lfm2"`.
113    #[serde(default)]
114    pub model_type: Option<String>,
115    /// Whether `lm_head` shares the token-embedding matrix.
116    #[serde(default)]
117    pub tie_word_embeddings: Option<bool>,
118    /// Per-head dimension when it is not `hidden_size / num_attention_heads`.
119    #[serde(default)]
120    pub head_dim: Option<usize>,
121}
122
123fn default_norm_eps() -> f32 {
124    1e-5
125}
126
127// Backwards compat alias
128/// Alias kept for older code that used the Llama-specific config name.
129pub type LlamaConfig = Lfm2Config;
130
131/// Reads model config files and raw safetensors files.
132pub struct ModelLoader {
133    // Map weight name -> (mmap_index, tensor_view) logic
134    // For MVP, simplified to single file or manual lookup
135}
136
137impl ModelLoader {
138    /// Reads and parses a `config.json` file.
139    pub fn load_config<P: AsRef<Path>>(path: P) -> Result<LlamaConfig, LoaderError> {
140        let file = File::open(path)?;
141        let config: LlamaConfig = serde_json::from_reader(file)?;
142        Ok(config)
143    }
144
145    // Helper to load a single safetensors file
146    /// Memory-maps a single safetensors file.
147    pub fn load_safetensors<P: AsRef<Path>>(path: P) -> Result<memmap2::Mmap, LoaderError> {
148        let file = File::open(path)?;
149        let mmap = unsafe { MmapOptions::new().map(&file)? };
150        Ok(mmap)
151    }
152}
153
154// A wrapper to hold the loaded SafeTensors views
155// Since SafeTensors borrows from Mmap, we need a struct that holds the Mmap owner
156// and provides access. This is tricky with self-referential structs in Rust.
157//
158// Solution: "WeightsContext" that holds the Mmaps, and we pass `&WeightsContext`
159// to a function that extracts views on demand.
160
161/// Memory-mapped safetensors files and an index from tensor name to file.
162pub struct WeightsContext {
163    mmaps: HashMap<String, memmap2::Mmap>, // filename -> mmap
164    // We can also parse the index.json if it exists to know which file has which weight
165    weight_map: HashMap<String, String>, // weight_name -> filename
166}
167
168#[derive(Debug, Deserialize)]
169struct SafetensorIndex {
170    weight_map: HashMap<String, String>,
171}
172
173/// Shape and dtype of a stored tensor.
174#[derive(Debug, Clone)]
175pub struct TensorMetadata {
176    /// Tensor dimensions.
177    pub shape: Vec<usize>,
178    /// Stored element type.
179    pub dtype: safetensors::Dtype,
180}
181
182impl Default for WeightsContext {
183    fn default() -> Self {
184        Self::new()
185    }
186}
187
188impl WeightsContext {
189    /// Creates an empty context with no files loaded.
190    pub fn new() -> Self {
191        Self {
192            mmaps: HashMap::new(),
193            weight_map: HashMap::new(),
194        }
195    }
196
197    /// Maps one safetensors file and indexes its tensor names.
198    pub fn load_file<P: AsRef<Path>>(&mut self, path: P) -> Result<(), LoaderError> {
199        let path = path.as_ref();
200        let filename = path
201            .file_name()
202            .map(|n| n.to_string_lossy().into_owned())
203            .ok_or_else(|| {
204                LoaderError::Io(std::io::Error::new(
205                    std::io::ErrorKind::InvalidInput,
206                    format!("{} is not a file path", path.display()),
207                ))
208            })?;
209        let mmap = ModelLoader::load_safetensors(path)?;
210
211        // Pre-scan to populate weight_map
212        let tensors = SafeTensors::deserialize(&mmap)?;
213        for name in tensors.names() {
214            self.weight_map.insert(name.to_string(), filename.clone());
215        }
216
217        self.mmaps.insert(filename, mmap);
218        Ok(())
219    }
220
221    /// Loads a model directory: sharded files via the index, else `model.safetensors`, else `consolidated.safetensors`.
222    pub fn load_model_dir<P: AsRef<Path>>(&mut self, model_dir: P) -> Result<(), LoaderError> {
223        let model_dir = model_dir.as_ref();
224        let index_path = model_dir.join("model.safetensors.index.json");
225        if index_path.exists() {
226            let file = File::open(&index_path)?;
227            let index: SafetensorIndex = serde_json::from_reader(file)?;
228            let mut shard_files: Vec<String> = index.weight_map.values().cloned().collect();
229            shard_files.sort();
230            shard_files.dedup();
231
232            for shard in shard_files {
233                self.load_file(model_dir.join(&shard))?;
234            }
235            return Ok(());
236        }
237
238        let primary = model_dir.join("model.safetensors");
239        if primary.exists() {
240            return self.load_file(primary);
241        }
242
243        let consolidated = model_dir.join("consolidated.safetensors");
244        if consolidated.exists() {
245            return self.load_file(consolidated);
246        }
247
248        Err(LoaderError::Io(std::io::Error::new(
249            std::io::ErrorKind::NotFound,
250            "no model.safetensors, consolidated.safetensors, or model.safetensors.index.json found",
251        )))
252    }
253
254    /// Returns the shape and dtype name of a tensor, or `None` if it is missing.
255    pub fn get_tensor_info(&self, name: &str) -> Option<(Vec<usize>, String)> {
256        // Shape, Dtype(str)
257        if let Some(filename) = self.weight_map.get(name) {
258            let mmap = self.mmaps.get(filename)?;
259            let st = SafeTensors::deserialize(mmap).ok()?;
260            let view = st.tensor(name).ok()?;
261            return Some((view.shape().to_vec(), format!("{:?}", view.dtype())));
262        }
263        None
264    }
265
266    /// Returns the shape and dtype of a tensor, or an error if it is missing.
267    pub fn get_tensor_metadata(&self, name: &str) -> Result<TensorMetadata, LoaderError> {
268        if let Some(filename) = self.weight_map.get(name) {
269            let mmap = self
270                .mmaps
271                .get(filename)
272                .ok_or_else(|| LoaderError::MissingWeight(name.to_string()))?;
273            let st = SafeTensors::deserialize(mmap)?;
274            let view = st.tensor(name)?;
275            return Ok(TensorMetadata {
276                shape: view.shape().to_vec(),
277                dtype: view.dtype(),
278            });
279        }
280        Err(LoaderError::MissingWeight(name.to_string()))
281    }
282
283    // We can't return TensorView because it borrows from local mmap variable inside function
284    // We must pass a closure or return raw data copy.
285    // For our usage: We want to upload to GPU immediately.
286    /// Calls `callback` with the raw bytes and shape of a tensor.
287    pub fn read_bytes<F>(&self, name: &str, callback: F) -> Result<(), LoaderError>
288    where
289        F: FnOnce(&[u8], &[usize]),
290    {
291        if let Some(filename) = self.weight_map.get(name) {
292            let mmap = self
293                .mmaps
294                .get(filename)
295                .ok_or_else(|| LoaderError::MissingWeight(name.to_string()))?;
296            let st = SafeTensors::deserialize(mmap)?;
297            let view = st.tensor(name)?;
298            callback(view.data(), view.shape());
299            Ok(())
300        } else {
301            Err(LoaderError::MissingWeight(name.to_string()))
302        }
303    }
304
305    /// Reads an F32, F16 or BF16 tensor as `f32` values plus its shape.
306    pub fn read_f32_tensor(&self, name: &str) -> Result<(Vec<f32>, Vec<usize>), LoaderError> {
307        let meta = self.get_tensor_metadata(name)?;
308        let numel: usize = meta.shape.iter().product();
309        let mut out = Vec::with_capacity(numel);
310        let mut actual_len = 0usize;
311
312        self.read_bytes(name, |bytes, shape| {
313            debug_assert_eq!(shape, meta.shape.as_slice());
314            actual_len = bytes.len();
315            match meta.dtype {
316                safetensors::Dtype::F32 => {
317                    if bytes.len() != numel * 4 {
318                        out = Vec::new();
319                        return;
320                    }
321                    out.extend(
322                        bytes.as_chunks::<4>().0.iter().map(|chunk| {
323                            f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])
324                        }),
325                    );
326                }
327                safetensors::Dtype::F16 => {
328                    if bytes.len() != numel * 2 {
329                        out = Vec::new();
330                        return;
331                    }
332                    out.extend(bytes.as_chunks::<2>().0.iter().map(|chunk| {
333                        half::f16::from_bits(u16::from_le_bytes([chunk[0], chunk[1]])).to_f32()
334                    }));
335                }
336                safetensors::Dtype::BF16 => {
337                    if bytes.len() != numel * 2 {
338                        out = Vec::new();
339                        return;
340                    }
341                    out.extend(bytes.as_chunks::<2>().0.iter().map(|chunk| {
342                        half::bf16::from_bits(u16::from_le_bytes([chunk[0], chunk[1]])).to_f32()
343                    }));
344                }
345                _ => out = vec![],
346            }
347        })?;
348
349        let expected_len = match meta.dtype {
350            safetensors::Dtype::F32 => numel * 4,
351            safetensors::Dtype::F16 | safetensors::Dtype::BF16 => numel * 2,
352            _ => {
353                return Err(LoaderError::UnsupportedDtype {
354                    name: name.to_string(),
355                    dtype: format!("{:?}", meta.dtype),
356                })
357            }
358        };
359
360        if out.len() != numel {
361            let actual = match meta.dtype {
362                safetensors::Dtype::F32 => actual_len,
363                safetensors::Dtype::F16 | safetensors::Dtype::BF16 => actual_len,
364                _ => 0,
365            };
366            return Err(LoaderError::InvalidTensorLength {
367                name: name.to_string(),
368                expected: expected_len,
369                actual,
370            });
371        }
372
373        Ok((out, meta.shape))
374    }
375
376    /// Reads `<name>.weight`, `.scales` and `.biases` and uploads them as one quantized GPU tensor.
377    pub fn upload_quantized(
378        &self,
379        name: &str,
380        ctx: &Arc<ferrotensor_backend::WgpuContext>,
381        engine: &ferrotensor_backend::ComputeEngine,
382    ) -> Result<UnifiedTensor, LoaderError> {
383        let weight_name = format!("{}.weight", name);
384        let scales_name = format!("{}.scales", name);
385        let biases_name = format!("{}.biases", name);
386
387        if !self.weight_map.contains_key(&weight_name) {
388            return Err(LoaderError::MissingWeight(weight_name));
389        }
390
391        // 1. Get Host Data
392        let mut q_data: Vec<u8> = Vec::new();
393        let mut q_shape: Vec<usize> = Vec::new();
394        self.read_bytes(&weight_name, |data, shape| {
395            q_data = data.to_vec();
396            q_shape = shape.to_vec();
397        })?;
398
399        let mut s_data: Vec<u8> = Vec::new();
400        self.read_bytes(&scales_name, |data, _| s_data = data.to_vec())?;
401
402        let mut b_data: Vec<u8> = Vec::new();
403        self.read_bytes(&biases_name, |data, _| b_data = data.to_vec())?;
404
405        // 2. Upload to GPU buffers using UnifiedTensor (or raw create using backend's re-exported wgpu)
406        // For simplicity, we'll convert to f32 slices where appropriate
407        // Here we assume scales and biases are f32 already in safetensors
408        let s_floats: Vec<f32> = bytemuck::cast_slice(&s_data).to_vec();
409        let b_floats: Vec<f32> = bytemuck::cast_slice(&b_data).to_vec();
410
411        let s_tensor = UnifiedTensor::new(ctx, &s_floats, &[s_floats.len()]);
412        let b_tensor = UnifiedTensor::new(ctx, &b_floats, &[b_floats.len()]);
413
414        // For quantized weights, we need a u8 buffer. The dequant kernel reads u32 (4 u8s packed).
415        // For now, let's just wrap it in a storage buffer.
416        let q_buf = ctx
417            .device
418            .create_buffer(&ferrotensor_backend::wgpu::BufferDescriptor {
419                label: Some("Q Weight"),
420                size: q_data.len() as u64,
421                usage: ferrotensor_backend::wgpu::BufferUsages::STORAGE
422                    | ferrotensor_backend::wgpu::BufferUsages::COPY_DST,
423                mapped_at_creation: false,
424            });
425        ctx.queue.write_buffer(&q_buf, 0, &q_data);
426
427        // 3. Output Tensor (FP32)
428        let out_tensor = UnifiedTensor::empty(ctx, &q_shape);
429
430        // 4. Dispatch Dequant
431        engine.affine_dequantize(
432            ctx,
433            &q_buf,
434            s_tensor.buffer(),
435            b_tensor.buffer(),
436            &out_tensor,
437        );
438
439        Ok(out_tensor)
440    }
441}
ferrotensor_loader/lib.rs

ferrotensor_loader/
lib.rs
