Implementation:LaurentMazare Tch rs Mmap Safetensors Load

Knowledge Sources	tch-rs
Domains	Deep_Learning, Memory_Optimization
Last Updated	2026-02-08 14:00 GMT

Overview

Concrete pattern for loading large model weights via memory-mapped safetensors with zero-copy tensor construction, demonstrated in the tch-rs LLaMA example.

Description

This pattern combines memmap2::MmapOptions for memory-mapped file I/O, safetensors::SafeTensors::deserialize for metadata parsing, and Tensor::from_blob for zero-copy tensor creation. The tensors are then copied into VarStore variables by directly accessing the variables_ mutex. This avoids the double memory allocation of standard VarStore::load.

Usage

Use for loading multi-gigabyte LLM weights. This requires direct access to VarStore internals (vs.variables_.lock()), which is a specialized pattern not part of the public API.

Code Reference

Source Location

Repository: tch-rs
File: examples/llama/main.rs
Lines: 363-378

Signature

// Pattern (not a single function):
let file = std::fs::File::open(&safetensors_path)?;
let content = unsafe { memmap2::MmapOptions::new().map(&file)? };
let safetensors = safetensors::SafeTensors::deserialize(&content)?;

let variables = vs.variables_.lock().unwrap();
for (name, var) in variables.named_variables.iter() {
    let view = safetensors.tensor(name)?;
    let tensor = Tensor::from_blob(view.data().as_ptr(), shape, strides, kind, device);
    var.f_copy_(&tensor)?;
}

Import

use memmap2::MmapOptions;
use safetensors::SafeTensors;
use tch::Tensor;

I/O Contract

Inputs

Name	Type	Required	Description
safetensors_path	Path	Yes	Path to .safetensors weight file
vs	&VarStore	Yes	VarStore with model parameters to populate

Outputs

Name	Type	Description
()	unit	VarStore variables populated from file via zero-copy + copy_ pattern

Usage Examples

use std::fs::File;
use memmap2::MmapOptions;
use safetensors::SafeTensors;

let file = File::open("llama.safetensors")?;
let content = unsafe { MmapOptions::new().map(&file)? };
let safetensors = SafeTensors::deserialize(&content)?;

let variables = vs.variables_.lock().unwrap();
for (name, var) in variables.named_variables.iter() {
    if let Ok(view) = safetensors.tensor(name) {
        let data = view.data();
        let tensor = Tensor::from_blob(
            data.as_ptr(), &shape, &strides, kind, tch::Device::Cpu
        );
        var.f_copy_(&tensor)?;
    }
}

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment