Saving model with shared tensors fails on cpu but succeeds on gpu
Author: kylesayrs
Pull request: no
Labels: name: bug

### System Info

platform: linux: `ubuntu 22.04`
python version: `3.10.12`
transformers version: `4.44.2`

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python3
# example.py

import torch
import pytest
from transformers import AutoModelForCausalLM

@pytest.mark.parametrize(
    "torch_dtype,tie_word_embeddings,device_map",
    [
        (torch.float16, True,  "cpu"   ),  # passes
        (torch.float16, False, "cpu"   ),  # passes
        (torch.float32, True,  "cpu"   ),  # passes
        (torch.float32, False, "cpu"   ),  # fails
        (torch.float32, False, "cuda:0"),  # passes
    ],
)
def test_model_save(torch_dtype, tie_word_embeddings, device_map, tmp_path):
    model = AutoModelForCausalLM.from_pretrained(
        "Xenova/llama2.c-stories15M",
        torch_dtype=torch_dtype,
        tie_word_embeddings=tie_word_embeddings,
        device_map=device_map,
    )
    model.save_pretrained(tmp_path, safe_serialization=True)

    # test that the model saved correctly
    reloaded = AutoModelForCausalLM.from_pretrained(
        tmp_path,
        torch_dtype="auto",
        device_map=device_map
    )

    model_dict = model.state_dict()
    reloaded_dict = reloaded.state_dict()
    assert model_dict.keys() == reloaded_dict.keys()
    for key in model_dict:
        assert torch.equal(model_dict[key], reloaded_dict[key])
        assert model_dict[key].device == reloaded_dict[key].device
```

```bash
python3 -m pytest example.py
```

```
RuntimeError: 
  Some tensors share memory, this will lead to duplicate memory on disk and potential differences when loading them again: [{'lm_head.weight', 'model.embed_tokens.weight'}].
  A potential way to correctly save your model is to use `save_model`.
  More information at https://huggingface.co/docs/safetensors/torch_shared_tensors
```

### Expected behavior

I expect `save_pretrained` to have the same behavior, regardless of model data type, and regardless of device

Comments:
Hmm indeed, this issue shouldn't pop up. @ydshieh, if you have the bandwidth, do you mind helping @kylesayrs out?

So far no clear idea yet, but when I tried with

> "meta-llama/Llama-2-7b-hf"

all cases are passing.

Would it possible how you create the original `Xenova/llama2.c-stories15M`?

```
    "torch_dtype,tie_word_embeddings,device_map",
    [
        (torch.float16, True,  "cpu"   ),  
        (torch.float16, False, "cpu"   ),  
        (torch.float16, True,  "cuda:0"   ), 
        (torch.float16, False, "cuda:0"   ),  
        (torch.float32, True,  "cpu"   ),  
        (torch.float32, False, "cpu"   ),  
        (torch.float32, True, "cuda:0"),  
        (torch.float32, False, "cuda:0"), 
    ],
```

@ydshieh I do not know the details of how the model was created, but from the [config.json](https://huggingface.co/Xenova/llama2.c-stories15M/blob/main/config.json) it seems like it was saved with `tie_word_embeddings=True`. Other models such as `TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T` do not have tied word embeddings and pass the test cases.

Some other models which have the same `tie_word_embeddings=True` in the config such as `"unsloth/Llama-3.2-3B-Instruct"` pass, so it may be an issue with this stories model in particular.

This model is not atypical in any way that I know of, but I'll do some more investigation to see if I notice anything different

Oh, ok. I thought you are the author of `llama2.c-stories15M` 😅  sorry. I would try to see if I can figure it out next week.


(just for me)
```python
import torch
from transformers import AutoModelForCausalLM

configs = [
    (torch.float32, False, "cpu"   ),  # fails
    #(torch.float16, True,  "cpu"   ),  # passes
    #(torch.float16, False, "cpu"   ),  # passes
    #(torch.float32, True,  "cpu"   ),  # passes
    #(torch.float32, False, "cpu"   ),  # fails
    #(torch.float32, False, "cuda:0"),  # passes
]

def test_model_save(torch_dtype, tie_word_embeddings, device_map, tmp_path="./"):
    model = AutoModelForCausalLM.from_pretrained(
        "Xenova/llama2.c-stories15M",
        torch_dtype=torch_dtype,
        tie_word_embeddings=tie_word_embeddings,
        device_map=device_map,
    )
    model.save_pretrained(tmp_path, safe_serialization=True)


for config in configs:
    test_model_save(*config)
```

We will merge #33913 after #36033 is merged

Fixed in #37031

Unfortunately, there is some issues and #37031 is currently reverted 

https://github.com/huggingface/transformers/pull/37178

And I don't have yet the bandwidth to look this again

This has either been resolved or was an issue with my usage