Storage alignment is not preserved

[AnthonyBarbier]

System Info

Version: 0.5.3

Information

• The official example scripts
• My own modified scripts

Reproduction

Some accelerators require the torch.Tensor storage to be 4 bytes aligned but because of the metadata at the start of the safetensors file and because tensors are serialised back to back without padding it happens that some fp32 values for example are only 2 bytes aligned (Which causes our HW to throw an exception).

from diffusers import StableDiffusionPipeline
import os

model_id = "stabilityai/stable-diffusion-2"
pipe = StableDiffusionPipeline.from_pretrained(
            model_id,
            use_safetensors=int(os.environ.get("USE_SAFETENSORS"))!= 0,
)
print(f"{pipe.text_encoder.text_model.encoder.layers[0].self_attn.q_proj.bias.data_ptr()/ 4.0}")
print(f"{pipe.text_encoder.text_model.encoder.layers[0].self_attn.q_proj.weight.data_ptr()/ 4.0}")

Note: because of a bug in diffusers you'll need to switch between two cache folders between runs or the model won't switch between safetensors / not-safetensors.

HF_HUB_CACHE=/tmp/safetensors USE_SAFETENSORS=1 TRANSFORMERS_VERBOSITY=debug python3 bug.py

Output:

loading weights file /tmp/safetensors/models--stabilityai--stable-diffusion-2/snapshots/1e128c8891e52218b74cde8f26dbfc701cb99d79/text_encoder/model.safetensors                                                                                                         
...
64712809483414.5
64712809484438.5

HF_HUB_CACHE=/tmp/no_safetensors USE_SAFETENSORS=0 TRANSFORMERS_VERBOSITY=debug python3 bug.py

Output:

loading weights file /tmp/no_safetensors/models--stabilityai--stable-diffusion-2/snapshots/1e128c8891e52218b74cde8f26dbfc701cb99d79/text_encoder/pytorch_model.bin
...
65466141378016.0
65466129834112.0

Expected behavior

All tensors alignments should match the original PyTorch tensor alignment (4 bytes).

[Narsil]

Some accelerators require the torch.Tensor storage to be 4 bytes aligned but because of the metadata at the start of the safetensors file and because tensors are serialised back to back without padding it happens that some fp32 values for example are only 2 bytes aligned (Which causes our HW to throw an exception).

Which accelerators ? Is this documented anywhere, it could be a torch bug no ?

Torch storage is now untyped, so I don't think it provides any alignement guarantees, so it should be up to the tensor creation to ensure byte-alignement, no ?

[AnthonyBarbier]

In this case it's Graphcore's IPU, but that's likely to be true for any accelerator with a unified memory architecture.

Some CPU SIMD technologies might be able to cope with it, but there will be a performance implication (e.g Arm NEON)

The default c10::alloc_cpu will use 64 bits alignment

Torch storage is untyped but it stores only one tensor per allocation, so it will align to the alloc_cpu()'s alignment, and if a tensor is a view it will at least align to the element size.

safetensors is a special case because it abuses the tensor storage to store unrelated metadata of arbitrary size at the beginning and then storing multiple tensors, potentially of different dtypes inside the same allocation.

I think the safest option would be to modify the safetensors writer to pad / align the beginning of each tensor to 64 bits to match the alignment of torch.save() / torch.load()? (This way you don't introduce unnecessary copies on load, and you can still just mmap the whole file)

[AnthonyBarbier]

Note: technically the link I've posted above requires 64 bytes alignment (For AVX512), but in our case 64 bits is enough, so up to you how compatible you want to be with other accelerators 😅

[Narsil]

I think the safest option would be to modify the safetensors writer to pad / align the beginning of each tensor to 64 bits to match the alignment of torch.save() / torch.load()? (This way you don't introduce unnecessary copies on load, and you can still just mmap the whole file)

This is not safetensors's job to do that, it's up to torch internals to figure out alignement with bytes. Torch's storage always comes first, then tensors are always views onto a storage, so if some accelerators requires specific alignment, then it's up to the tensor allocation to reallocate with proper alignement, as the alignment requirement might depend on the accelerator. And the reason is simple, alignment is not just a hardware property, it can also be important for the kernels themselves and how they are written.

Also note that typed storages are considered legacy: https://docs.pytorch.org/docs/stable/storage.html#legacy-typed-storage.

safetensors now does enforce alignement during writing to help get actual zero-copy loads, and does alignment when loading the tensor directly using it's own mmap. But old files are there, they are totally valid files