Reference implementation for the hash calculation

[Nerogar]

According to the spec, the hash should be calculated as "A hash of all tensor content (ie excluding the header section), with 0x prefix, all lowercase, and no byte-separator symbols."

I'm assuming this means we need to calculate the hash of the data as it is saved in the safetensors file itself. For validation of an existing file this is fine, but how do we calculate that hash before saving a model?

Do we need to save the model, read it to calculate the hash, then save it again with the hash included? The safetensors library provides a save function to save into a byte array. But I assume this creates a copy of the entire model in RAM, which is not ideal for big models.

An efficient reference implementation would be nice.

[mcmonkey4eva]

That is an excellent question, that I do not currently have an answer to.

cc @sayakpaul @Narsil can either of you help with this? How do we calculate the hash within the safetensors library?

[kohya-ss]

I am running into the same problem. This implementation uses less memory and returns the same value as far as I have checked, but I am not sure if I can rely on the order of the state dict.

def load_bytes_in_safetensors(tensors):
    bytes = safetensors.torch.save(tensors)
    b = BytesIO(bytes)

    b.seek(0)
    header = b.read(8)
    n = int.from_bytes(header, "little")

    offset = n + 8
    b.seek(offset)

    return b.read()


def precalculate_safetensors_hashes(state_dict):
    # calculate each tensor one by one to reduce memory usage
    hash_sha256 = hashlib.sha256()
    for tensor in state_dict.values():
        single_tensor_sd = {"tensor": tensor}
        bytes_for_tensor = load_bytes_in_safetensors(single_tensor_sd)
        hash_sha256.update(bytes_for_tensor)

    return f"0x{hash_sha256.hexdigest()}"

[Nerogar]

I also thought about doing something very similar. But I think some kind of official support would be a lot better. I can think of two ways:

1. Add a new function to the safetensors library that handles everything itself and just returns a hash. Parameters could maybe just be the state dict and some hash algorithm identifier.
2. Add a new function to the safetensors library that returns a "stream" for bytes, instead of a full byte array. This would be more flexible for other applications. I'm not an expert python developer, so I don't know what the idiomatic way to do this would be.

Of course, both solutions would require an update to the safetensors library itself. (Or maybe this is already supported and I just couldn't find it in the documentation)

[Narsil]

How do we calculate the hash within the safetensors library?

There is no way, nor a standard proposed hash for checking the files.
You can however rely on the fact there's a deterministic order in which the tensors are written to disk.

DTYPE by decreasing order (F64 - U64 - U32 ....) and then regular alphabetic order: https://github.com/huggingface/safetensors/blob/main/safetensors/src/tensor.rs#L173-L175

Please note that this is a current implementation detail, and that the actual order on the files upon load doesn't matter.
And I cannot guarantee nothing will ever change. (The layout of the file can enable tricks to speed up loading via zero-copy by modifying alignment) The current choice seems pretty stable to me at the time, but the field is moving very fast and new dtypes might change things around.

There's an issue open in safetensors for a checkum: huggingface/safetensors#220

As I say in the related issue, I would differentiate checksums and hashes :

• checksum fast, and efficient hash to make sure bits are what they are supposed to be, it's an internal/hardware integrity check
• hash: Cryptographic guarantee that a local file is the same file as the person offering it to me. Then the hash MUST be sent via a different channel, otherwise the man in the middle can do whatever it wants.

Ofc it's pretty much the same algorithm but very different intents and therefore usually you can make different choices (cryptographic resistance vs speed of execution).

[liuliu]

Please note that this is a current implementation detail, and that the actual order on the files upon load doesn't matter.

Haha, this helps because my implementation of safetensors from Draw Things app (export model) doesn't really do this ordering :)

[Narsil]

@liuliu Ordering is not a requirement, it's an implementation detail. This ordering only helps to do actual zero-copy.

The reason is that 4x u8 is a valid f32 if and only if the first byte has and adress divisible by 4 (it's called aligned).
Otherwise you have to copy your bytes to such a location before casting to f32.

Therefore if we 8-align the first f64 tensors, and go by decreasing Dtype order, we're sure all tensors will be aligned with their respective dtype without ever checking !

[jenuk]

Since hashes have a fixed length, one could save the weights with yet unknown hash and a dummy value (of the required length) in the hash field. Then after saving, the hash can be calculated (implementation independent) and only that part of the header needs to be changed. The length of everything stays the same, so no need to move any bits around, which should in theory result in a fast edit.

[jenuk]

This seems to works well. You just need to add a placeholder when saving the file originally (example in initialize). Then you can add the hash (add_hash) without overhead (e.g. ~10s to calculate hash, ~1ms to save new value on my laptop for SDXL).

import hashlib
import json
from time import perf_counter
from typing import Optional

# A placeholder for the hash value, can be any ascii string that is unique
# in the safetensors header. Needs to be the exact length of the hash.
HASH_PLACEHOLDER = "SHA256-PLACEHOLDER"
HASH_PLACEHOLDER = HASH_PLACEHOLDER + "-" * (64 - len(HASH_PLACEHOLDER))


def initialize(filename: str, out: Optional[str] = None):
    """
    Example to add hash placeholder. Should be done , when all other
    metadata is also saved. You most likely don't want to use this function
    for speed reasons.
    """
    out = out if out is not None else filename

    # load the model appropriately for your use case
    from safetensors.torch import load_file

    data = load_file(filename)

    # load/generate metadata
    with open(filename, "rb") as in_file:
        N_byte = in_file.read(8)
        N = int.from_bytes(N_byte, "little")
        header = in_file.read(N)
    parsed_header = json.loads(header.decode("utf8"))
    if "__metadata__" in parsed_header:
        metadata = parsed_header["__metadata__"]
    else:
        metadata = dict()

    # Add hash placeholder to metadata
    metadata["modelspec.hash_sha256"] = "0x" + HASH_PLACEHOLDER
    # This assume that the order of tensors later-on is unknown

    # save safetensors with any appropriate library
    from safetensors.torch import save_file

    save_file(data, out, metadata)


def add_hash(filename: str):
    """Replace `HASH_PLACEHOLDER` with actual tensor-data hash"""
    t0 = perf_counter()

    # get header and calculate hash
    with open(filename, "rb") as file:
        N_byte = file.read(8)
        N = int.from_bytes(N_byte, "little")
        header = file.read(N)
        data_hash = hashlib.sha256()
        # 1MB buffers
        while (block := file.read(1024 * 1024)) != b"":
            data_hash.update(block)
    digest = data_hash.hexdigest()
    t1 = perf_counter()

    # position of placeholder
    pos = header.find((HASH_PLACEHOLDER).encode("utf8"))

    # make sure unique match was found
    assert pos >= 0, "Did not find a hash placeholder"
    pos2 = header.find((HASH_PLACEHOLDER).encode("utf8"), pos + 1)
    assert pos2 == -1, f"Found multiple hash placeholders, {pos2}"

    # replace only the placeholder
    with open(filename, "r+b") as file:
        file.seek(8 + pos)
        file.write(digest.encode("utf8"))
    t2 = perf_counter()
    print(f"Calculating hash: {t1 - t0:.2g}s\nUpdating hash: {t2 - t1:.2g}s")

[Nerogar]

One problem with that solution is, that the file still needs to be read in after the save. For reference, here is my solution, but it relies on an internal function:

import safetensors.torch as safetensors
import hashlib

def __calculate_safetensors_hash(state_dict: dict[str, Tensor]) -> str:
    sha256_hash = hashlib.sha256()

    ordered_state_dict = OrderedDict(sorted(state_dict.items()))
    for key, tensor in ordered_state_dict.items():
        data = safetensors._tobytes(tensor, key)
        sha256_hash.update(data)

    return f"0x{sha256_hash.hexdigest()}"