Bintensors: The Ultimate Guide to Fast Model Storage for ML Developers

高效码农

9 hours ago

What is bintensors? A Complete Guide for Junior College Graduates

In this blog post, we’ll explore bintensors, a binary encoded file format designed for fast storage of models and tensors. This guide is tailored for junior college graduates and above, with a focus on clarity and practicality. We’ll cover installation, usage, file format details, performance benefits, and answer common questions. All content is derived solely from the provided source material, ensuring technical accuracy and authenticity.

Introduction to bintensors

In the realm of machine learning, efficient model storage and loading are crucial. Bintensors emerges as a novel binary file format, offering rapid zero-copy access to models and tensors. Initially developed as an exploration of the safetensors format, bintensors aims to enhance understanding of model distribution across subnets. It maintains similar properties to safetensors but introduces significant performance improvements.

Installation Methods

Using Cargo (for Rust Developers)

Rust developers can easily add bintensors to their projects via Cargo, the Rust package manager. Simply run the following command in your terminal:

cargo add bintensors

This command will download and install the latest version of bintensors, making it available for use in your Rust project.

Using Pip (for Python Developers)

Python developers can install bintensors using pip, the Python package installer. Execute the following command to install bintensors:

pip install bintensors

After running this command, bintensors will be installed in your Python environment, ready for use in your Python scripts.

From Source Code

If you prefer to install bintensors from the source code, you’ll first need to have Rust installed. You can install Rust by running:

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

Update Rust to ensure you’re using the stable version:

rustup update

Next, clone the bintensors repository from GitHub:

git clone https://github.com/GnosisFoundation/bintensors

Navigate to the Python bindings directory and install the required setuptools_rust:

cd bintensors/bindings/python
pip install setuptools_rust

Finally, install bintensors with the following command:

pip install -e .

Basic Usage Examples

Python Example: Saving and Loading Tensors

Here’s a simple Python example demonstrating how to save and load tensors using bintensors:

import torch
from bintensors import safe_open
from bintensors.torch import save_file

# Create some tensor data
tensors = {
   "weight1": torch.zeros((1024, 1024)),
   "weight2": torch.zeros((1024, 1024))
}

# Save tensors to a bintensors file
save_file(tensors, "model.bt")

# Load tensors from the file
tensors = {}
with safe_open("model.bt", framework="pt", device="cpu") as f:
   for key in f.keys():
       tensors[key] = f.get_tensor(key)

In this example, we first create two tensors and save them to a file named “model.bt”. We then use the safe_open function to open the file and load each tensor by iterating over the keys in the file.

Rust Example: Handling bintensors Files

Here’s an example of how to handle bintensors files in Rust:

use bintensors::BinTensors;
use memmap2::MmapOptions;
use std::fs::File;

let filename = "model.bt";
use std::io::Write;
let serialized = b"\x18\x00\x00\x00\x00\x00\x00\x00\x00\x01\x08weight_1\x00\x02\x02\x02\x00\x04       \x00\x00\x00\x00";
File::create(&filename).unwrap().write(serialized).unwrap();
let file = File::open(filename).unwrap();
let buffer = unsafe { MmapOptions::new().map(&file).unwrap() };
let tensors = BinTensors::deserialize(&buffer).unwrap();
let tensor = tensors
        .tensor("weight_1");
std::fs::remove_file(filename).unwrap()

This example demonstrates how to create a bintensors file, write it to disk, and then open it using memory mapping. The file is then deserialized to access the tensor data.

Understanding the bintensors File Format

The bintensors file format consists of three main sections: header size, header data, and tensor data. Let’s break down each component:

Header Size

The header size is an 8-byte little-endian unsigned 64-bit integer indicating the size of the header data. It is capped at 100MB, similar to safetensors, though this limit may be adjusted in the future.

Header Data

The header data is a dynamically serialized table encoded in a condensed binary format for efficient tensor lookups. It contains a map of string-to-string pairs. Arbitrary JSON structures are not permitted; all values must be strings. The header data deserialization decoding capacity is limited to 100MB.

Tensor Data

The tensor data is a sequence of bytes representing the layered tensor data. The buffer size can be calculated using the following formula:

$B_{M} = t_{i} \in T \sum ([j = 1 \prod n_{i} d_{i, j}] \cdot D (x_{i}))$

Where:

( B_M ) is the total buffer size for a model or model subset ( M )
( T ) is the set of tensors in the model
Each tensor ( t_i ) has a shape dimension tuple ( (d_{i,1}, d_{i,2}, \dots, d_{i,n_i}) )
( D ) is a function mapping tensor type ( x_i ) to bytes of tensor dtype sizes, which can be ( {1, 2, 4, 8} )

For example, let’s calculate the bytes required to store the embedding layer of the GPT-2 model. The model has two large tensors: the token embedder (wte) with shape ( (50,257,768) ) and the position embedder (wpe) with shape ( (1,024,768) ). Assuming all weights are stored as float32:

$B_{e mb e dd in g} = (50, 257 \times 768) \times 4 + (1, 024 \times 768) \times 4 = (50, 257 \times 768) \times 4 + (1, 024 \times 768) \times 4 = 38, 597, 376 \times 4 + 786, 432 \times 4 = 154, 389, 504 + 3, 145, 728 = 157, 535, 232 ∴ B_{e n co d er} totals 157, 535, 232 bytes or 157.54 MB$

Notes on the File Format

Duplicate keys are not allowed, although not all parsers may enforce this.
Tensor values are not validated, meaning NaN and +/-Inf may be present in the file.
Empty tensors (tensors with one dimension being 0) are permitted. They do not store data in the data buffer but retain size in the header. They are accepted as they are valid tensors in traditional tensor libraries.
The byte buffer must be fully indexed with no holes, preventing the creation of polyglot files.
Endianness is little-endian.
Order is ‘C’ or row-major.
A checksum is applied to the bytes, giving the file a unique identity. This allows distributed networks to validate distributed layer checksums.

Performance Benefits of bintensors

Faster Deserialization with bincode

Bintensors replaces the serde_json used in safetensors with the bincode library for serialization and deserialization. This change results in a significant performance boost, nearly tripling deserialization speed.

Benchmarking code can be found in the bintensors/bench/benchmark.rs file. Two separate tests were conducted per repository to compare model test serialization performance in safetensors and bintensors within the Rust-only implementation. The results highlight the substantial improvements achieved with bintensors.

To better understand the reasons behind this enhancement, a call stack analysis was performed, comparing the performance characteristics of serde_json and bincode. Flame graphs were generated to visualize execution paths and identify potential bottlenecks in the serde_json deserializer.

Performance Comparison

The bincode library is designed for fast binary serialization and deserialization, making it more efficient than serde_json for handling large-scale model metadata. The flame graphs illustrate the differences in performance between bincode and serde_json deserializers.

Advantages of bintensors

Performance Improvement

Bintensors provides a noticeable performance boost to the growing ecosystem of model storage. Its use of bincode for serialization and deserialization results in faster file loading times, especially for large models.

Prevention of DOS Attacks

Bintensors implements robust security measures to prevent DOS attacks. The header buffer is strictly limited to 100MB, preventing resource exhaustion attacks through oversized metadata. Additionally, strict address boundary validation ensures non-overlapping tensor allocations, guaranteeing that memory consumption does not exceed the actual file size during loading operations. This dual approach effectively mitigates both memory exhaustion and buffer overflow vulnerabilities.

Faster Loading Times

Compared to major machine learning formats, PyTorch offers fast file loading. However, bintensors bypasses the extra CPU copy that PyTorch typically performs by using torch.UntypedStorage.from_file. Currently, CPU loading times with bintensors are extremely fast, outperforming pickle. GPU loading times are either equivalent or faster than PyTorch. Loading tensors on CPU first with memory mapping and then transferring them to GPU also appears to be faster, similar to the behavior observed with PyTorch pickle.

Lazy Loading

In distributed settings (multi-node or multi-GPU), lazy loading is highly beneficial. For instance, using bintensors with the BLOOM model reduced loading time from 10 minutes with regular PyTorch weights to just 45 seconds across 8 GPUs. This significantly accelerates feedback loops during model development. When changing distribution strategies (e.g., Pipeline Parallelism vs Tensor Parallelism), separate weight copies are not required.

Frequently Asked Questions (FAQ)

Q1: What is the difference between bintensors and safetensors?

A1: Bintensors is a simple fork of safetensors. While it shares many properties with safetensors, such as preventing DOS attacks, fast loading, and lazy loading, the primary difference lies in the use of bincode instead of serde_json for storing metadata. This change results in a significant performance improvement in deserialization speed, nearly three times faster.

Q2: Which programming languages does bintensors support?

A2: Bintensors primarily supports Rust and Python. Rust developers can use Cargo to add bintensors as a project dependency, while Python developers can install it via pip. The installation methods for both languages are provided earlier in this guide.

Q3: How does bintensors handle large model files?

A3: Bintensors efficiently handles large model files through its compact binary format and lazy loading mechanism. The header data size is limited to 100MB, ensuring metadata does not become excessively large. Tensor data is stored in a compact binary format, reducing file size. Lazy loading allows only the required portions of the model to be loaded in distributed environments, saving memory and reducing loading time.

Q4: Is the bintensors file format secure?

A4: Yes, bintensors implements several security measures. The header buffer size is restricted to 100MB to prevent DOS attacks using oversized metadata. Strict address boundary validation ensures tensor allocations do not overlap, preventing memory consumption from exceeding the actual file size during loading. These measures effectively reduce security risks.

Q5: Can bintensors be used on different operating systems?

A5: Yes, bintensors is platform-independent. Although benchmarking was conducted on macOS, its design ensures compatibility across various operating systems, including Windows, Linux, and macOS. You can use bintensors to store and load model files on different platforms.

Q6: How can file integrity be verified?

A6: Bintensors calculates a checksum for the file’s byte data, providing each file with a unique identity. This allows distributed networks to validate the checksums of distributed layers, ensuring file integrity and consistency.

Q7: What tensor data types does bintensors support?

A7: Bintensors supports various common tensor data types, including float32, float64, int32, and int64. The function D(x_i) maps tensor types to byte sizes, typically 1, 2, 4, or 8 bytes.

Q8: Does bintensors support tensor data compression?

A8: Currently, the bintensors file format does not explicitly mention compression functionality for tensor data. It focuses on efficient storage and rapid access to tensor data. However, due to its binary format and compact header data, file sizes are generally smaller compared to text-based formats.

Q9: How to specify the device (CPU or GPU) when using bintensors in Python?

A9: When using the safe_open function in Python with bintensors, you can specify the framework (e.g., “pt” for PyTorch) and the device (e.g., “cpu” or “cuda”). For example:

with safe_open("model.bt", framework="pt", device="cpu") as f:
   for key in f.keys():
       tensors[key] = f.get_tensor(key)

Q10: Does bintensors support dynamic shapes?

A10: Bintensors allows for the storage of tensors with varying shapes. The shape dimensions of each tensor are explicitly recorded in the header data. As long as shape information is consistent during saving and loading, bintensors can support tensors with dynamic shapes.

Future Outlook

As machine learning models continue to grow in size and complexity, the need for efficient, secure, and flexible storage solutions becomes increasingly important. Bintensors, with its compact binary format, fast loading capabilities, and robust security features, presents an attractive option in this field. We anticipate seeing further development and adoption of bintensors within the machine learning community, driving innovation and improvement in model workflows.

Conclusion

This comprehensive guide has covered bintensors, including installation methods, usage examples, file format details, performance benefits, and answers to common questions. We hope this guide helps you grasp the core concepts and practical applications of bintensors. Whether you’re a Rust or Python developer, bintensors offers an efficient and secure solution for model and tensor storage.

By employing a conversational writing style and FAQ section, we aim to address potential questions and smooth the learning process. We encourage you to experiment with bintensors in your projects and explore the optimizations and improvements it can bring to your machine learning workflows.

For more detailed information, please refer to the official bintensors documentation (docs.rs) or visit its GitHub repository (GnosisFoundation/bintensors).