Converting a PyTorch Checkpoint to Safetensors

A .bin checkpoint is a pickle, and loading one executes whatever the pickle says to execute. Safetensors is the same numbers behind a length-prefixed JSON header that cannot describe anything but tensors. The conversion is mechanical; the two things worth care are verification and tied weights.

Why the format change is worth the trouble

The safetensors layout is deliberately boring, and the boredom is the feature. Per the safetensors specification, a file starts with eight bytes holding a little-endian unsigned 64-bit integer — the length of the header — followed by that many bytes of UTF-8 JSON, followed by a flat byte buffer. The header maps each tensor name to its dtype, its shape and a two-element data_offsets pair giving start and one-past-end positions in the buffer.

Three consequences follow. A reader can list the tensors, their shapes and their sizes without touching the data, which is how the Hub shows you a model’s tensor list without downloading it. The buffer can be memory-mapped, so a loader pages in only what it needs rather than materialising the whole checkpoint. And there is no code path from file contents to execution — the format has no expressive power beyond “here is a tensor”, whereas torch.load on an untrusted pickle is arbitrary code execution unless you are careful.

The conversion

Two routes. If the checkpoint belongs to a Transformers model, load and re-save it — that path also rewrites the config and handles sharding for you:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("./old-checkpoint", dtype="auto")
model.save_pretrained("./new-checkpoint", safe_serialization=True)

If it is a bare state dict, or you want to avoid instantiating the model at all, go through the tensors directly. Note weights_only=True: it restricts the unpickler to tensors and primitives, which is the difference between reading a file and running it.

import torch
from safetensors.torch import save_file

sd = torch.load("pytorch_model.bin", map_location="cpu", weights_only=True)

# safetensors stores a flat name -> tensor map; contiguous, no shared storage
sd = {k: v.contiguous() for k, v in sd.items()}

save_file(sd, "model.safetensors", metadata={"format": "pt"})

The metadata argument is a free-form string-to-string map. The format key is the convention the Hugging Face ecosystem reads to know which framework’s tensor conventions apply; omit it and some loaders will complain.

Verifying tensor by tensor

Do not skip this, and do not verify by file size. The point of the conversion is that you now trust the new file, and that trust should rest on a comparison rather than on the absence of an exception.

1. Load the original with weights_only=True.
2. Load the new file with safetensors.torch.load_file.
3. Assert the key sets are identical — a missing key here is the tied-weight problem in the next section.
4. For each key, assert dtype and shape match, then assert the raw bytes match.

import torch
from safetensors.torch import load_file

a = torch.load("pytorch_model.bin", map_location="cpu", weights_only=True)
b = load_file("model.safetensors")

assert a.keys() == b.keys(), sorted(set(a) ^ set(b))
for k in a:
    x, y = a[k], b[k]
    assert x.dtype == y.dtype and x.shape == y.shape, k
    assert x.contiguous().view(torch.uint8).equal(y.contiguous().view(torch.uint8)), k
print(f"{len(a)} tensors identical")

Comparing the byte views rather than using torch.allclose is intentional. This conversion should be exact — nothing is being rounded or cast — so a tolerance-based check would hide a real bug. Reserve allclose for conversions that genuinely change precision, such as an ONNX export, where a small numerical difference is expected and the tolerance is the thing being asserted.

The shared-memory error

The one failure everybody hits:

RuntimeError: Some tensors share memory, this will lead to duplicate memory
on disk and potential differences when loading them again

PyTorch allows two entries in a state dict to point at one storage. Language models do this routinely: many tie the input embedding and the output projection, so lm_head.weight and model.embed_tokens.weight are the same memory under two names. Safetensors has no way to express that — every tensor owns a disjoint slice of the buffer — so it refuses rather than silently writing the data twice and letting the two copies drift apart on a later fine-tune.

There are two correct fixes and one wrong one. Drop the duplicate name and let the model re-tie it on load, which is what save_pretrained does using the model’s own record of which weights are tied. Or, if the two really are meant to be independent from now on, clone() one of them so they occupy separate storage. The wrong fix is falling back to safe_serialization=False, which gets you a pickle again and abandons the reason you started. Hugging Face’s note on shared tensors sets out both correct paths.

What safetensors will not carry

The format holds named tensors and a string-to-string metadata map. That is the complete list, and a .bin file can hold considerably more, so some conversions lose things that were never tensors.

• Optimizer state. A training checkpoint often bundles optimizer moments, the learning-rate schedule, the epoch counter and an RNG state under keys alongside the weights. The tensors among those convert; the Python objects do not. If you need to resume training, keep the original file — a converted inference checkpoint cannot resume anything.
• A pickled model object. Some .bin files contain a whole nn.Module rather than a state dict, because somebody called torch.save(model) instead of torch.save(model.state_dict()). Loading that with weights_only=True will refuse, correctly. You need the class definition available, must load it without the safety flag from a source you trust, and should then extract .state_dict() before converting.
• Non-tensor entries. Integers, strings and config dictionaries stored as state-dict values have nowhere to go. Move them into config.json, or into the metadata map as strings — the map accepts only strings, so anything structured needs encoding.
• Unsupported dtypes. Complex tensors and some newer low-precision types have no safetensors dtype code. Cast or drop them deliberately rather than discovering the gap at load time.

None of these is a reason not to convert. They are a reason to convert the inference weights and keep the training checkpoint, rather than treating the conversion as a replacement for the original file.

Sharded checkpoints

Anything beyond a few gigabytes arrives as several files plus an index. The index is model.safetensors.index.json, and it holds a weight_map from tensor name to the file containing it, plus a metadata object whose total_size is the checkpoint’s size in bytes. Loaders read the index, then open only the shards they need.

For conversion this means the unit of work is the shard: convert each pytorch_model-00001-of-00004.bin to the matching model-00001-of-00004.safetensors, then rewrite the index with the new filenames. Do not merge shards while converting unless you have checked the sizes — the arithmetic behind shard counts explains why the shard boundaries are where they are, and a single 40 GB file is awkward for reasons that have nothing to do with the format.