Adding `mlmf` library as a dependency

Run this command in a terminal, in your project's directory:

cargo add mlmf

To add it manually, edit your project's Cargo.toml file and add to the [dependencies] section:

mlmf = "0.2.0"

The mlmf library will be automatically available globally. Read the mlmf library documentation.

Back to the crate overview.

Readme

MLMF - Machine Learning Model Files

MLMF (Machine Learning Model Files) is a comprehensive Rust crate for working with ML model files. MLMF provides loading, saving, conversion, and dynamic mapping capabilities for transformer models across multiple formats including SafeTensors, GGUF, ONNX, PyTorch, and AWQ. It eliminates code duplication and provides a unified, efficient API for model file operations.

Features

🏗️ Architecture Detection: Automatically detects model architecture (LLaMA, GPT-2, GPT-NeoX) from tensor names
📦 Multiple Formats: Comprehensive support for SafeTensors, GGUF, ONNX, PyTorch, and AWQ formats
🗺️ Name Mapping: Intelligent tensor name mapping between HuggingFace and custom formats
💾 Memory Efficient: Memory-mapped loading for large models (30GB+)
🔧 Device Management: Automatic CUDA detection with CPU fallback
📊 Progress Reporting: Optional progress callbacks for long-running operations
🛡️ Type Safety: Comprehensive error handling with detailed context
🔄 Model Conversion: Direct format conversion with batch processing and progress tracking

Quick Start

Add mlmf to your Cargo.toml:

[dependencies]
mlmf = "0.1"

Basic Usage

use mlmf::{LoadOptions, loader};
use candle_core::{Device, DType};

// Load a LLaMA model from SafeTensors
let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
let options = LoadOptions {
    device: device.clone(),
    dtype: DType::F16,
    use_mmap: true,
    validate_cuda: false,
    progress: Some(mlmf::progress::default_progress()),
};

let loaded_model = loader::load_safetensors("./models/llama-7b", options)?;

// Access components
let var_builder = loaded_model.var_builder;
let config = loaded_model.config;
let name_mapper = loaded_model.name_mapper;

// Use name mapper to convert HF names to your format
if let Some(mapped_name) = name_mapper.map_name("model.layers.0.self_attn.q_proj.weight") {
    println!("Mapped name: {}", mapped_name);
}

Architecture Detection

use mlmf::name_mapping::{TensorNameMapper, Architecture};

let tensor_names = vec![
    "model.embed_tokens.weight".to_string(),
    "model.layers.0.self_attn.q_proj.weight".to_string(),
    "model.norm.weight".to_string(),
];

let mapper = TensorNameMapper::from_tensor_names(&tensor_names)?;
assert_eq!(mapper.architecture(), Architecture::LLaMA);

Model Conversion

use mlmf::conversion::{convert_model, ConversionFormat, ConversionOptions};
use std::path::Path;

// Convert from SafeTensors to ONNX
let options = ConversionOptions::default();
let result = convert_model(
    Path::new("model.safetensors"),
    Path::new("model.onnx"),
    ConversionFormat::ONNX,
    options,
)?;

println!("Conversion completed in {:.2}s", result.duration.as_secs_f64());

Architecture

MLMF provides a modular architecture with the following components:

loader: High-level loading API
conversion: Direct model format conversion with batch processing
name_mapping: Architecture detection and tensor name mapping
config: HuggingFace config parsing with field aliases
formats: Format-specific loaders and exporters (SafeTensors, GGUF, ONNX, PyTorch, AWQ)
validation: CUDA validation and dtype checking
progress: Progress reporting utilities

Supported Models

LLaMA Family: LLaMA 2/3, TinyLlama, Qwen, Mistral
GPT Family: GPT-2, GPT-J
GPT-NeoX Family: GPT-NeoX, Pythia, StableLM

Examples

See the examples/ directory for complete working examples:

load_llama.rs - Loading LLaMA models from SafeTensors
load_gpt2.rs - Loading GPT-2 models
load_gguf.rs - Loading quantized GGUF models

Performance

MLMF is optimized for performance:

Memory-mapped loading: Loads 70B models (130GB) in ~10 seconds
Architecture detection: Typically completes in <100ms
Zero-copy: Direct tensor access without unnecessary copying
Incremental builds: Changes compile in <10 seconds

Contributing

Contributions are welcome! Please see CONTRIBUTING.md for guidelines.

License

This project is licensed under either of:

Apache License, Version 2.0 (LICENSE-APACHE)
MIT License (LICENSE-MIT)

at your option.

Adding mlmf library as a dependency

Adding `mlmf` library as a dependency