InstantTensor is an ultra-fast, distributed Safetensors loader designed to maximize I/O throughput when moving model weights from Safetensors files to GPU memory.
| Model | GPU | Backend | Load Time (s) | Throughput (GB/s) | Speedup |
|---|---|---|---|---|---|
| Qwen3-30B-A3B | 1*H200 | Safetensors | 57.4 | 1.1 | 1x |
| Qwen3-30B-A3B | 1*H200 | InstantTensor | 1.77 | 35 | 32.4x |
| DeepSeek-R1 | 8*H200 | Safetensors | 160 | 4.3 | 1x |
| DeepSeek-R1 | 8*H200 | InstantTensor | 15.3 | 45 | 10.5x |
See Benchmark for full benchmarks.
from instanttensor import safe_open
tensors = {}
with safe_open("model.safetensors", framework="pt", device=0) as f:
for name, tensor in f.tensors():
tensors[name] = tensorYielded tensors own their memory by default (copy=True). For zero-copy
streaming into preallocated storage, see Zero-copy mode.
See Usage for multi-file and distributed usage.
- Fast weight loading
- Direct I/O: Avoid the slow page cache allocation on cold start. Friendly for large models and tight memory budgets.
- Tuned I/O size and concurrency: Maximize hardware throughput.
- Pipelining and prefetching: Parallelize and overlap the various stages of transmission.
- Distributed loading
- Use
torch.distributed(NCCL) to speed up loading under any parallelism policy (TP/PP/EP/CP/DP).
- Use
- Multiple I/O backends
- Supports multiple backends: GPUDirect Storage, Legacy Storage, and Memory-based Storage.
InstantTensor is recommended if any of the following conditions are met:
- High storage bandwidth (>= 5 GB/s).
- Unable to keep the model cached in host memory, for example:
- Limited free memory for model caching (for example, when most memory is used for KV cache offloading in LLM serving).
- Infrequent model loading, where Linux page cache is less effective.
- Model switching, where multiple models cannot be cached in memory simultaneously.
- The model is heavily sharded (for example, TP=8), resulting in small, non-contiguous I/O per GPU.
- Loading from
tmpfs.
- GPU platforms: CUDA, ROCm
- Framework: PyTorch
pip install instanttensorgit clone https://github.com/scitix/InstantTensor.git
cd InstantTensor
./checkout_submodules.sh
pip install .
# For a debug build, set "DEBUG=1" before "pip"Passing a list of files allows the backend to plan reads and provides higher throughput than making multiple calls to load single files:
from instanttensor import safe_open
files = ["model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"]
tensors = {}
with safe_open(files, framework="pt", device=0) as f:
for name, tensor in f.tensors():
tensors[name] = tensorInstantTensor can use a torch.distributed NCCL process group to coordinate loading and achieve higher throughput compared to running safe_open independently on each GPU.
import torch
import torch.distributed as dist
from instanttensor import safe_open
dist.init_process_group(backend="nccl")
process_group = dist.GroupMember.WORLD
files = ["model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"]
tensors = {}
with safe_open(files, framework="pt", device=torch.cuda.current_device(), process_group=process_group) as f:
for name, tensor in f.tensors():
tensors[name] = tensorNOTE: You can also load weights using a subgroup created via
dist.new_group, which allows multiple subgroups to load weights independently. For example, if you have TP=8 and PP=2 (i.e., two TP groups), you can create two subgroups and load weights independently on each TP group. In cross-node (multi-machine) scenarios, loading using per-node subgroups can sometimes be faster than loading on the world group. However, for most cases, the world group is a good default choice.
For regular disk files, InstantTensor defaults to Direct I/O to prioritize cold load performance. This is usually the right choice when a model is loaded once for a long-running workload.
If the same model is loaded repeatedly within a short period, Buffered I/O can
be faster because later reads may benefit from the page cache. Enable it with
the backend=BackendPolicy.BUFFERED argument to safe_open, or, when
backend=None, by setting the INSTANTTENSOR_BACKEND=BUFFERED environment
variable.
Pass copy=False to skip the per-tensor clone and yield views into the
internal ring buffer:
with safe_open(files, framework="pt", device=0, copy=False) as f:
for name, tensor in f.tensors():
model_param[name].copy_(tensor)Two rules:
- Consume each tensor before the next is yielded —
list(f.tensors())and similar patterns silently corrupt data whenbuffer_size < total_tensor_size. - Do not keep references past the
withblock — the buffer is freed on exit.
A UserWarning fires when copy=False and buffer_size < total_tensor_size.
Both attributes are public on the safe_open object.
See tests/test.py for a full benchmark harness (TP/PP grouping, checksums, etc.).
InstantTensor selects an I/O backend automatically by default. You can provide
one or more backend candidates with the backend argument to safe_open, or
with the INSTANTTENSOR_BACKEND environment variable when backend=None.
InstantTensor tries the candidates in order and uses the first backend that is
supported by the file system and available on the current system.
Supported backend values are Backend.AIO, Backend.AIO_BUFFERED,
Backend.URING, Backend.URING_BUFFERED, Backend.CUFILE, and
Backend.MMAP. The backend argument accepts a single Backend or a list of
Backend/BackendPolicy values:
from instanttensor import Backend, BackendPolicy, safe_open
safe_open("model.safetensors", framework="pt", device=0, backend=Backend.URING)
safe_open("model.safetensors", framework="pt", device=0, backend=[Backend.URING, Backend.AIO])
safe_open("model.safetensors", framework="pt", device=0, backend=BackendPolicy.BUFFERED)BackendPolicy.BUFFERED expands to [Backend.URING_BUFFERED, Backend.AIO_BUFFERED, Backend.MMAP]. This is a good choice when you want Buffered I/O.
INSTANTTENSOR_BACKEND accepts comma-separated backend or policy names:
INSTANTTENSOR_BACKEND=URING,AIO
INSTANTTENSOR_BACKEND=BUFFEREDBackends are used in different file-system and I/O scenarios:
- In-memory file systems (available backends:
MMAP,URING_BUFFERED,AIO_BUFFERED): when model files are stored on tmpfs or ramfs,MMAPprovides the best compatibility and performance for this case. The other backends are usually slower for in-memory files. - Regular file systems: InstantTensor can use either Direct I/O or Buffered
I/O.
- Direct I/O (available backends:
AIO,URING,CUFILE) is best when a model is loaded once for a long-running workload. It avoids page-cache cold-start effects and reduces page-cache pollution. When choosing manually,URINGmay be faster on newer platforms.AIOhas the broadest platform compatibility.CUFILErequires GPUDirect Storage support, and its higher throughput can be offset by cuFile initialization overhead. - Buffered I/O (available backends:
AIO_BUFFERED,URING_BUFFERED,MMAP) is best when the same model is loaded repeatedly within a short period. Later reads can benefit from the page cache, though the first read is usually slower than Direct I/O.URING_BUFFEREDis preferred on platforms with io_uring support;AIO_BUFFEREDprovides a more compatible option, whileMMAPis available but usually not preferred.
- Direct I/O (available backends:
If no backend is specified, InstantTensor tries [URING, AIO] for regular disk
files. For tmpfs/ramfs files, it uses MMAP. If none of the requested
candidates can be used, InstantTensor raises an error listing why each candidate
was rejected.
Thanks to the AI Systems and Optimization team at ScitiX AI and the Wenfei Wu Lab at Peking University.