Problem saving torch nn.GRU

[hslr4]

System Info

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

• transformers version: 4.43.4
• Platform: Linux-6.1.0-37-amd64-x86_64-with-glibc2.35
• Python version: 3.10.14
• Huggingface_hub version: 0.27.1
• Safetensors version: 0.6.2
• Accelerate version: not installed
• Accelerate config: not found
• PyTorch version (GPU?): 2.8.0+cu128 (True)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?: no
• Using GPU in script?: yes
• GPU type: NVIDIA RTX 4000 Ada Generation

Information

• The official example scripts
• My own modified scripts

Reproduction

I found the problem to be reported already here: huggingface/accelerate#3101 (comment)
However their solution safe_serialization=False seems not to work anymore and also does only partially solve the problem.

Basically the following error occurs when I try to save a custom model containing a torch.nn.GRU with model.save_pretrained('test'):

Traceback (most recent call last):
  File "/path/main.py", line 156, in <module>
    model.save_pretrained('./test')
  File "/opt/miniconda3/envs/customlipsync/lib/python3.10/site-packages/huggingface_hub/hub_mixin.py", line 408, in save_pretrained
    self._save_pretrained(save_directory)
  File "/opt/miniconda3/envs/customlipsync/lib/python3.10/site-packages/huggingface_hub/hub_mixin.py", line 755, in _save_pretrained
    save_model_as_safetensor(model_to_save, str(save_directory / constants.SAFETENSORS_SINGLE_FILE))
  File "/opt/miniconda3/envs/customlipsync/lib/python3.10/site-packages/safetensors/torch.py", line 169, in save_model
    to_removes = _remove_duplicate_names(state_dict)
  File "/opt/miniconda3/envs/customlipsync/lib/python3.10/site-packages/safetensors/torch.py", line 113, in _remove_duplicate_names
    raise RuntimeError(
RuntimeError: Error while trying to find names to remove to save state dict, but found no suitable name to keep for saving amongst: {'gru.weight_ih_l0'}. None is covering the entire storage.Refusing to save/load the model since you could be storing much more memory than needed. Please refer to https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an issue.

Adding a check if there is only one shared tensor, I found here: https://huggingface.co/spaces/safetensors/convert/blob/main/convert.py#L54 to the torch implementation here (I guess)

safetensors/bindings/python/py_src/safetensors/torch.py

Line 112 in fa22263

if not complete_names:

seems to solve the problem.

Expected behavior

The model is saved without an error.

[Narsil]

Can you add a reproducer without accelerate ?

This seems to be working fine:

#!/usr/bin/env python3

from safetensors.torch import save_file
import torch.nn as nn


class ReferenceEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.gru = nn.GRU(
            input_size=12,
            hidden_size=12,
            batch_first=True,
        )


# Test
model = ReferenceEncoder()
save_file(model.state_dict(), "out.safetensors")

[Silverster98]

Can you add a reproducer without accelerate ?

This seems to be working fine:

#!/usr/bin/env python3

from safetensors.torch import save_file
import torch.nn as nn

class ReferenceEncoder(nn.Module):
def init(self):
super().init()
self.gru = nn.GRU(
input_size=12,
hidden_size=12,
batch_first=True,
)

Test

model = ReferenceEncoder()
save_file(model.state_dict(), "out.safetensors")

@Narsil hi. I notice that the flatten_parameters() only works on GPU (see pytorch doc). Once you add model = model.cuda(), the error will occur.

Code:

from safetensors.torch import save_file, load_model

import torch.nn as nn
class ReferenceEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.gru = nn.GRU(
            input_size=12,
            hidden_size=12,
            batch_first=True,
        )

# Test
model = ReferenceEncoder().cuda()
save_file(model.state_dict(), "out.safetensors")
load_model(model, "out.safetensors")

Output:

Traceback (most recent call last):
  File "/home/wangzan/Projects/see-inter/a.py", line 29, in <module>
    load_model(model, "out.safetensors")
  File "/home/wangzan/miniconda3/envs/inter/lib/python3.12/site-packages/safetensors/torch.py", line 221, in load_model
    to_removes = _remove_duplicate_names(
                 ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/wangzan/miniconda3/envs/inter/lib/python3.12/site-packages/safetensors/torch.py", line 113, in _remove_duplicate_names
    raise RuntimeError(
RuntimeError: Error while trying to find names to remove to save state dict, but found no suitable name to keep for saving amongst: {'gru.weight_ih_l0'}. None is covering the entire storage.Refusing to save/load the model since you could be storing much more memory than needed. Please refer to https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an issue.

[Narsil]

Hmm. I cannot reproduce on the GPU I have access to (Hopper). Also torch 2.8, python 3.10...

Have you tried doing {k: t.contiguous() for k, t in state_dict.items()} ?

This should work for saving, but since I cannot inspect how the tensors are laid out, it might cause issues when reloading although I don't really see why it should make any difference.
Honestly, I'm not sure what torch does there, but I'm guessing they are allocating the entire storage at once, and slicing in the tensors (I'm not sure what's the benefit but hey).

[Silverster98]

I just created a Colab notebook and ran the code using a T4 GPU; the error can be reproduced. Please check the code on Colab.

I also try to do {k: t.contiguous() for k, t in state_dict.items()} before saving state_dict, but the error still occur.

[Narsil]

Here is the cleaned up version:

from safetensors.torch import save_file, load_model, load_file
import torch.nn as nn
import torch

class ReferenceEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.gru = nn.GRU(
            input_size=12,
            hidden_size=12,
            batch_first=True,
        )
# Test 1
model = ReferenceEncoder().cuda()
state_dict = model.state_dict()
save_file(state_dict, "out.safetensors")
new_state = load_file("out.safetensors")
model.load_state_dict(new_state)

load_model should only really be used with save_model (and it doesn't work here specifically because it's the pytorch model GRU that isn't laid out normally.)

So load_file/save_file works great here.

The issue with load_model is that it's slightly too strict into what it expects to save you from and the error could be avoided.

[Silverster98]

Thanks! It works!

But, I also use the accelerate for training models. I noticed that accelerate is using save_file for saving and using load_model for loading. Is there any workaround?

I will also create an issue in accelerate repo to mention this.

[Narsil]

For completeness, it can be solved in safetensors by modifying the code here:

safetensors/bindings/python/py_src/safetensors/torch.py

Line 112 in fa22263

if not complete_names:

To allow for "shared tensors" which don't have a complete overlap. We could simply check that there is only 1 such tensor and assume that it's because it's part of a disjoint set.
The issue is that I need to spend time confirming if that assumption can actually be discarded here, or propagate more into load_model to make sure the line isn't hit.