[Core] Saving models with multiple shared tensor groups is not supported when model is dispatched
Author: kylesayrs
Pull request: no
Labels: name: Distributed Training / Models

name: bug

### System Info

Any transformers version, any python version, any system

### Who can help?

@Rocketknight1 @ArthurZucker 

### Reproduction

```python3
from transformers import AutoModelForCausalLM

device_map = {
    "model.embed_tokens": "cuda:0",
    "model.layers": "cpu",
    "model.norm": "cpu",
    "model.rotary_emb": "cpu",
    "lm_head": "cuda:0"
}
model = AutoModelForCausalLM.from_pretrained("nm-testing/llama2.c-stories15M", device_map=device_map)

transform_a = torch.nn.Linear(1, 1, bias=False)
transform_a._dynamic_tied_weights_keys = ["weight"]
transform_b = torch.nn.Linear(1, 1, bias=False)
transform_b._dynamic_tied_weights_keys = ["weight"]

model.model.layers[0].self_attn.q_proj.register_module("transform", transform_a)
model.model.layers[1].self_attn.q_proj.register_module("transform", transform_a)
model.model.layers[2].self_attn.q_proj.register_module("transform", transform_b)
model.model.layers[3].self_attn.q_proj.register_module("transform", transform_b)

model.save_pretrained("tmp")
```

### Expected behavior

Downstream user [LLM Compressor](https://github.com/vllm-project/llm-compressor) would like to save models which have multiple shared tensor groups. For example, these two modules share two sets of shared tensors.
```python3
assert module_a.weight is module_b.weight
assert module_a.bias is module_b.bias
assert module_a._dynamic_tied_weights_keys = ["weight", "bias"]
assert module_b._dynamic_tied_weights_keys = ["weight", "bias"]
```

However, attempting to save a model with modules like these raises an error only if `hf_device_map` is present.
```
RuntimeError:
Some tensors share memory, this will lead to duplicate memory on disk and potential differences when loading them again: [...]
```

The problematic code seems to lie here
```python3
if hasattr(self, "hf_device_map"):
    # if the model has offloaded parameters, we must check using find_tied_parameters()
    tied_params = find_tied_parameters(self)
    if tied_params:
        tied_names = tied_params[0]
        shared_ptrs = {
            ptr: names for ptr, names in ptrs.items() if any(name in tied_names for name in names)
        }
    else:
        shared_ptrs = {}
else:
    shared_ptrs = {ptr: names for ptr, names in ptrs.items() if len(names) > 1}
```

As you can see, the line `tied_names = tied_params[0]` assumes that there is only one shared tensor group. This means that only the first shared tensor group is used to populate `shared_ptrs`, which means that only the first group is deleted. Other groups are not deleted which results in the downstream error.

This assumption was likely made because the most common case is `lm_head.weight == embed_tokens.weight`. However, there's no reason why multiple groups can't be supported, as it's provably possible in the non-offloaded case.

Comments:
cc @matej-svejda and @SunMarc ! 

Thanks for the report and giving a try to fix this issue ! Just note that we need to use `find_tied_parameters` to find tied parameters because the tensors are on the `meta` device. We can't rely on `id(tensor)` to find shared tensors.