voxtral_micro/models/
weights.rs

1//! Weight loading from SafeTensors files.
2//!
3//! Loads pre-trained Voxtral weights into Burn model modules.
4
5use anyhow::{Context, Result};
6use burn::module::{Param, ParamId};
7use burn::nn::Linear;
8use burn::tensor::backend::Backend;
9use burn::tensor::{Tensor, TensorData};
10use safetensors::SafeTensors;
11use std::path::Path;
12use std::sync::Arc;
13
14/// Backing storage for SafeTensors bytes — either heap-allocated or memory-mapped.
15enum BytesBacking {
16    Owned(Arc<Vec<u8>>),
17    Mapped(memmap2::Mmap),
18}
19
20impl AsRef<[u8]> for BytesBacking {
21    fn as_ref(&self) -> &[u8] {
22        match self {
23            BytesBacking::Owned(v) => v,
24            BytesBacking::Mapped(m) => m,
25        }
26    }
27}
28
29/// Load a tensor from SafeTensors by name.
30pub fn load_tensor<B: Backend, const D: usize>(
31    safetensors: &SafeTensors,
32    name: &str,
33    device: &B::Device,
34) -> Result<Tensor<B, D>> {
35    let tensor_view = safetensors
36        .tensor(name)
37        .with_context(|| format!("Tensor '{}' not found", name))?;
38
39    let shape: Vec<usize> = tensor_view.shape().to_vec();
40    let dtype = tensor_view.dtype();
41
42    // Convert to f32
43    let data: Vec<f32> = match dtype {
44        safetensors::Dtype::F32 => {
45            let bytes = tensor_view.data();
46            bytes
47                .chunks_exact(4)
48                .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
49                .collect()
50        }
51        safetensors::Dtype::F16 => {
52            let bytes = tensor_view.data();
53            bytes
54                .chunks_exact(2)
55                .map(|b| {
56                    let bits = u16::from_le_bytes([b[0], b[1]]);
57                    half::f16::from_bits(bits).to_f32()
58                })
59                .collect()
60        }
61        safetensors::Dtype::BF16 => {
62            let bytes = tensor_view.data();
63            bytes
64                .chunks_exact(2)
65                .map(|b| {
66                    let bits = u16::from_le_bytes([b[0], b[1]]);
67                    half::bf16::from_bits(bits).to_f32()
68                })
69                .collect()
70        }
71        _ => anyhow::bail!("Unsupported dtype: {:?}", dtype),
72    };
73
74    // Create tensor data with shape
75    let tensor_data = TensorData::new(data, shape);
76
77    // Create Burn tensor
78    let tensor: Tensor<B, D> = Tensor::from_data(tensor_data, device);
79    Ok(tensor)
80}
81
82/// Load a single named tensor directly from raw safetensors bytes.
83///
84/// This bypasses `SafeTensors::deserialize` to avoid a `usize` overflow
85/// in the crate's validation on wasm32: for tensors with > 268M elements,
86/// the intermediate `nelements * bitsize_in_bits` exceeds `u32::MAX`.
87///
88/// The byte-level size is fine — only the bits calculation overflows.
89pub fn load_tensor_raw<B: Backend, const D: usize>(
90    bytes: &[u8],
91    name: &str,
92    device: &B::Device,
93) -> Result<Tensor<B, D>> {
94    use anyhow::bail;
95
96    if bytes.len() < 8 {
97        bail!("Safetensors data too short for header length");
98    }
99
100    let header_size = u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize;
101    if 8 + header_size > bytes.len() {
102        bail!(
103            "Safetensors header size {} exceeds file length {}",
104            header_size,
105            bytes.len()
106        );
107    }
108
109    let header: serde_json::Value = serde_json::from_slice(&bytes[8..8 + header_size])
110        .context("Failed to parse safetensors header JSON")?;
111
112    let info = header
113        .get(name)
114        .with_context(|| format!("Tensor '{}' not found in safetensors header", name))?;
115
116    let dtype_str = info["dtype"]
117        .as_str()
118        .context("Missing dtype in tensor info")?;
119    let shape: Vec<usize> = info["shape"]
120        .as_array()
121        .context("Missing shape in tensor info")?
122        .iter()
123        .map(|v| v.as_u64().unwrap() as usize)
124        .collect();
125    let start = info["data_offsets"][0]
126        .as_u64()
127        .context("Missing data_offsets[0]")? as usize;
128    let end = info["data_offsets"][1]
129        .as_u64()
130        .context("Missing data_offsets[1]")? as usize;
131
132    let data_start = 8 + header_size;
133    if data_start + end > bytes.len() {
134        bail!(
135            "Tensor '{}' data range [{}, {}) exceeds file length {}",
136            name,
137            data_start + start,
138            data_start + end,
139            bytes.len()
140        );
141    }
142    let tensor_bytes = &bytes[data_start + start..data_start + end];
143
144    let data: Vec<f32> = match dtype_str {
145        "F32" => tensor_bytes
146            .chunks_exact(4)
147            .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
148            .collect(),
149        "F16" => tensor_bytes
150            .chunks_exact(2)
151            .map(|b| {
152                let bits = u16::from_le_bytes([b[0], b[1]]);
153                half::f16::from_bits(bits).to_f32()
154            })
155            .collect(),
156        "BF16" => tensor_bytes
157            .chunks_exact(2)
158            .map(|b| {
159                let bits = u16::from_le_bytes([b[0], b[1]]);
160                half::bf16::from_bits(bits).to_f32()
161            })
162            .collect(),
163        _ => bail!("Unsupported dtype: {}", dtype_str),
164    };
165
166    let tensor_data = TensorData::new(data, shape);
167    Ok(Tensor::from_data(tensor_data, device))
168}
169
170/// Owning wrapper for SafeTensors that keeps bytes alive without leaking.
171///
172/// This struct owns the backing storage (heap bytes or memory-mapped file)
173/// and provides safe access to the SafeTensors view.
174/// The backing is freed when this struct is dropped.
175pub struct OwnedSafeTensors {
176    _backing: BytesBacking,
177    // SAFETY: safetensors borrows from _backing which we keep alive.
178    // We use 'static here but the actual lifetime is tied to _backing.
179    safetensors: SafeTensors<'static>,
180}
181
182impl OwnedSafeTensors {
183    /// Load SafeTensors from a file path using memory-mapping.
184    ///
185    /// The OS pages in data on demand — no multi-GB heap allocation for the
186    /// raw file bytes. This dramatically reduces peak memory when loading
187    /// large models (e.g. 8.9 GB safetensors → ~17.8 GB peak instead of ~25 GB).
188    pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
189        let file = std::fs::File::open(path.as_ref())
190            .with_context(|| format!("Failed to open: {}", path.as_ref().display()))?;
191        let mmap = unsafe { memmap2::Mmap::map(&file) }
192            .with_context(|| format!("Failed to mmap: {}", path.as_ref().display()))?;
193        let backing = BytesBacking::Mapped(mmap);
194
195        // SAFETY: We're creating a SafeTensors that borrows from `backing`.
196        // We store both in the same struct, and _backing is never moved or dropped
197        // while safetensors exists. The mmap stays valid for the struct's lifetime.
198        let safetensors = unsafe {
199            let static_ref: &'static [u8] = std::mem::transmute(backing.as_ref());
200            SafeTensors::deserialize(static_ref).context("Failed to deserialize SafeTensors")?
201        };
202
203        Ok(Self {
204            _backing: backing,
205            safetensors,
206        })
207    }
208
209    /// Create from raw bytes (heap-allocated).
210    pub fn from_bytes(bytes: Vec<u8>) -> Result<Self> {
211        let backing = BytesBacking::Owned(Arc::new(bytes));
212
213        // SAFETY: We're creating a SafeTensors that borrows from `backing`.
214        // We store both in the same struct, and _backing is never moved or dropped
215        // while safetensors exists. The Arc ensures the bytes live long enough.
216        let safetensors = unsafe {
217            let static_ref: &'static [u8] = std::mem::transmute(backing.as_ref());
218            SafeTensors::deserialize(static_ref).context("Failed to deserialize SafeTensors")?
219        };
220
221        Ok(Self {
222            _backing: backing,
223            safetensors,
224        })
225    }
226
227    /// Get a reference to the SafeTensors.
228    pub fn tensors(&self) -> &SafeTensors<'_> {
229        &self.safetensors
230    }
231}
232
233// Implement Deref for convenient access
234impl std::ops::Deref for OwnedSafeTensors {
235    type Target = SafeTensors<'static>;
236
237    fn deref(&self) -> &Self::Target {
238        &self.safetensors
239    }
240}
241
242/// Load SafeTensors file from disk (returns owning wrapper).
243pub fn load_safetensors<P: AsRef<Path>>(path: P) -> Result<OwnedSafeTensors> {
244    OwnedSafeTensors::from_file(path)
245}
246
247/// Weight name prefixes for different model components.
248pub mod prefixes {
249    /// Audio encoder prefix.
250    pub const ENCODER: &str = "mm_streams_embeddings.embedding_module.whisper_encoder";
251    /// LLM decoder prefix.
252    pub const DECODER: &str = "layers";
253    /// Token embeddings prefix.
254    pub const TOK_EMBEDDINGS: &str = "mm_streams_embeddings.embedding_module.tok_embeddings.weight";
255    /// Adapter prefix.
256    pub const ADAPTER: &str = "mm_streams_embeddings.embedding_module.audio_language_projection";
257    /// Final norm prefix.
258    pub const FINAL_NORM: &str = "norm.weight";
259}
260
261/// List all tensor names in a SafeTensors file.
262pub fn list_tensors<'a>(safetensors: &'a SafeTensors<'a>) -> Vec<&'a str> {
263    safetensors.names().into_iter().collect()
264}
265
266/// Filter tensor names by prefix.
267pub fn filter_tensors<'a>(safetensors: &'a SafeTensors<'a>, prefix: &str) -> Vec<&'a str> {
268    safetensors
269        .names()
270        .into_iter()
271        .filter(|name| name.starts_with(prefix))
272        .collect()
273}
274
275/// Create a Linear layer from weight tensors.
276///
277/// Note: PyTorch stores Linear weights as [out_features, in_features],
278/// but Burn expects [in_features, out_features]. This function handles
279/// the transpose automatically.
280pub fn linear_from_weights<B: Backend>(
281    weight: Tensor<B, 2>,
282    bias: Option<Tensor<B, 1>>,
283) -> Linear<B> {
284    // PyTorch Linear: weight [out, in], Burn Linear: weight [in, out]
285    // Need to transpose
286    let weight = weight.transpose();
287
288    Linear {
289        weight: Param::initialized(ParamId::new(), weight),
290        bias: bias.map(|b| Param::initialized(ParamId::new(), b)),
291    }
292}
293
294/// Load a Linear layer from SafeTensors.
295///
296/// # Arguments
297/// * `safetensors` - SafeTensors file
298/// * `weight_name` - Name of the weight tensor
299/// * `bias_name` - Optional name of the bias tensor
300/// * `device` - Device to load tensors on
301pub fn load_linear<B: Backend>(
302    safetensors: &SafeTensors,
303    weight_name: &str,
304    bias_name: Option<&str>,
305    device: &B::Device,
306) -> Result<Linear<B>> {
307    let weight: Tensor<B, 2> = load_tensor(safetensors, weight_name, device)?;
308    let bias = if let Some(name) = bias_name {
309        // Check if bias exists
310        if safetensors.tensor(name).is_ok() {
311            Some(load_tensor::<B, 1>(safetensors, name, device)?)
312        } else {
313            None
314        }
315    } else {
316        None
317    };
318
319    Ok(linear_from_weights(weight, bias))
320}
321
322/// Encoder layer weight names.
323pub fn encoder_layer_weight_names(layer_idx: usize) -> EncoderLayerWeightNames {
324    let prefix = format!("{}.transformer.layers.{}", prefixes::ENCODER, layer_idx);
325
326    EncoderLayerWeightNames {
327        attention_norm: format!("{}.attention_norm.weight", prefix),
328        wq_weight: format!("{}.attention.wq.weight", prefix),
329        wq_bias: format!("{}.attention.wq.bias", prefix),
330        wk_weight: format!("{}.attention.wk.weight", prefix),
331        wv_weight: format!("{}.attention.wv.weight", prefix),
332        wv_bias: format!("{}.attention.wv.bias", prefix),
333        wo_weight: format!("{}.attention.wo.weight", prefix),
334        wo_bias: format!("{}.attention.wo.bias", prefix),
335        ffn_norm: format!("{}.ffn_norm.weight", prefix),
336        w1_weight: format!("{}.feed_forward.w1.weight", prefix),
337        w2_weight: format!("{}.feed_forward.w2.weight", prefix),
338        w2_bias: format!("{}.feed_forward.w2.bias", prefix),
339        w3_weight: format!("{}.feed_forward.w3.weight", prefix),
340    }
341}
342
343/// Weight names for an encoder layer.
344pub struct EncoderLayerWeightNames {
345    pub attention_norm: String,
346    pub wq_weight: String,
347    pub wq_bias: String,
348    pub wk_weight: String,
349    pub wv_weight: String,
350    pub wv_bias: String,
351    pub wo_weight: String,
352    pub wo_bias: String,
353    pub ffn_norm: String,
354    pub w1_weight: String,
355    pub w2_weight: String,
356    pub w2_bias: String,
357    pub w3_weight: String,
358}
359
360/// Decoder layer weight names.
361pub fn decoder_layer_weight_names(layer_idx: usize) -> DecoderLayerWeightNames {
362    let prefix = format!("{}.{}", prefixes::DECODER, layer_idx);
363
364    DecoderLayerWeightNames {
365        // ADA RMSNorm conditioning (t-embed projection)
366        ada_norm_down: format!("{}.ada_rms_norm_t_cond.0.weight", prefix),
367        ada_norm_up: format!("{}.ada_rms_norm_t_cond.2.weight", prefix),
368        attention_norm: format!("{}.attention_norm.weight", prefix),
369        wq_weight: format!("{}.attention.wq.weight", prefix),
370        wk_weight: format!("{}.attention.wk.weight", prefix),
371        wv_weight: format!("{}.attention.wv.weight", prefix),
372        wo_weight: format!("{}.attention.wo.weight", prefix),
373        ffn_norm: format!("{}.ffn_norm.weight", prefix),
374        w1_weight: format!("{}.feed_forward.w1.weight", prefix),
375        w2_weight: format!("{}.feed_forward.w2.weight", prefix),
376        w3_weight: format!("{}.feed_forward.w3.weight", prefix),
377    }
378}
379
380/// Weight names for a decoder layer.
381pub struct DecoderLayerWeightNames {
382    pub ada_norm_down: String,
383    pub ada_norm_up: String,
384    pub attention_norm: String,
385    pub wq_weight: String,
386    pub wk_weight: String,
387    pub wv_weight: String,
388    pub wo_weight: String,
389    pub ffn_norm: String,
390    pub w1_weight: String,
391    pub w2_weight: String,
392    pub w3_weight: String,
393}
394
395/// Conv downsampler weight names.
396pub fn conv_weight_names() -> ConvWeightNames {
397    ConvWeightNames {
398        conv1_weight: format!("{}.conv_layers.0.conv.weight", prefixes::ENCODER),
399        conv1_bias: format!("{}.conv_layers.0.conv.bias", prefixes::ENCODER),
400        conv2_weight: format!("{}.conv_layers.1.conv.weight", prefixes::ENCODER),
401        conv2_bias: format!("{}.conv_layers.1.conv.bias", prefixes::ENCODER),
402    }
403}
404
405/// Weight names for the conv downsampler.
406pub struct ConvWeightNames {
407    pub conv1_weight: String,
408    pub conv1_bias: String,
409    pub conv2_weight: String,
410    pub conv2_bias: String,
411}
412
413/// Adapter weight names.
414pub fn adapter_weight_names() -> AdapterWeightNames {
415    AdapterWeightNames {
416        linear1_weight: format!("{}.0.weight", prefixes::ADAPTER),
417        linear2_weight: format!("{}.2.weight", prefixes::ADAPTER),
418    }
419}
420
421/// Weight names for the audio-language adapter.
422pub struct AdapterWeightNames {
423    pub linear1_weight: String,
424    pub linear2_weight: String,
425}
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430    use std::path::PathBuf;
431    use std::sync::OnceLock;
432
433    fn model_path() -> PathBuf {
434        PathBuf::from("models/voxtral/consolidated.safetensors")
435    }
436
437    /// Shared SafeTensors loader for tests - loads model once and reuses.
438    /// This prevents OOM from parallel tests each loading ~8GB.
439    static SHARED_SAFETENSORS: OnceLock<OwnedSafeTensors> = OnceLock::new();
440
441    fn get_shared_safetensors() -> Option<&'static OwnedSafeTensors> {
442        let path = model_path();
443        if !path.exists() {
444            return None;
445        }
446        Some(
447            SHARED_SAFETENSORS.get_or_init(|| {
448                load_safetensors(&path).expect("Failed to load shared safetensors")
449            }),
450        )
451    }
452
453    #[test]
454    fn test_load_safetensors_exists() {
455        let Some(safetensors) = get_shared_safetensors() else {
456            println!("Skipping: model not downloaded. Run: ./scripts/download_model.py");
457            return;
458        };
459
460        let names = list_tensors(safetensors);
461
462        println!("Found {} tensors", names.len());
463        assert!(!names.is_empty(), "Should have tensors");
464
465        // Check expected tensor names
466        assert!(names.contains(&"norm.weight"), "Should have final norm");
467    }
468
469    #[test]
470    fn test_filter_encoder_tensors() {
471        let Some(safetensors) = get_shared_safetensors() else {
472            println!("Skipping: model not downloaded");
473            return;
474        };
475
476        let encoder_tensors = filter_tensors(safetensors, prefixes::ENCODER);
477
478        println!("Found {} encoder tensors", encoder_tensors.len());
479        assert!(!encoder_tensors.is_empty(), "Should have encoder tensors");
480    }
481
482    #[test]
483    fn test_load_tensor() {
484        use burn::backend::Wgpu;
485
486        let Some(safetensors) = get_shared_safetensors() else {
487            println!("Skipping: model not downloaded");
488            return;
489        };
490
491        let device = Default::default();
492
493        // Load final norm weight
494        let norm_weight: Tensor<Wgpu, 1> =
495            load_tensor(safetensors, prefixes::FINAL_NORM, &device).unwrap();
496
497        println!("Norm weight shape: {:?}", norm_weight.dims());
498        assert_eq!(norm_weight.dims(), [3072], "Should be [3072]");
499    }
500
501    #[test]
502    fn test_load_encoder_attention_weight() {
503        use burn::backend::Wgpu;
504
505        let Some(safetensors) = get_shared_safetensors() else {
506            println!("Skipping: model not downloaded");
507            return;
508        };
509
510        let device = Default::default();
511
512        // Load encoder layer 0 attention wq weight
513        let names = encoder_layer_weight_names(0);
514        let wq: Tensor<Wgpu, 2> = load_tensor(safetensors, &names.wq_weight, &device).unwrap();
515
516        println!("Encoder wq shape: {:?}", wq.dims());
517        // Should be [n_heads * head_dim, d_model] = [32 * 64, 1280] = [2048, 1280]
518        assert_eq!(wq.dims(), [2048, 1280]);
519    }
520
521    #[test]
522    fn test_load_linear() {
523        use burn::backend::Wgpu;
524
525        let Some(safetensors) = get_shared_safetensors() else {
526            println!("Skipping: model not downloaded");
527            return;
528        };
529
530        let device = Default::default();
531
532        // Load adapter linear1
533        let names = adapter_weight_names();
534        let linear: Linear<Wgpu> =
535            load_linear(safetensors, &names.linear1_weight, None, &device).unwrap();
536
537        // Adapter linear1: [5120, 3072] in PyTorch -> [3072, 5120] after transpose
538        // Actually, looking at the docs, it's [out, in] -> [in, out]
539        // So if PyTorch has [3072, 5120], Burn needs [5120, 3072]
540        let dims = linear.weight.dims();
541        println!("Adapter linear1 weight shape: {:?}", dims);
542        // PyTorch [3072, 5120] -> Burn [5120, 3072]
543        assert_eq!(dims[0], 5120, "d_input");
544        assert_eq!(dims[1], 3072, "d_output");
545    }
546
547    #[test]
548    fn test_encoder_layer_weight_names() {
549        let names = encoder_layer_weight_names(5);
550        assert!(names.wq_weight.contains(".5."));
551        assert!(names.wq_weight.ends_with("attention.wq.weight"));
552    }
553
554    #[test]
555    fn test_decoder_layer_weight_names() {
556        let names = decoder_layer_weight_names(10);
557        assert!(names.wq_weight.contains(".10."));
558        assert!(names.wq_weight.ends_with("attention.wq.weight"));
559    }
560}
voxtral_micro/models/weights.rs

voxtral_micro/models/
weights.rs
