Error when "device_map='auto'" meets "load_state_dict"
Author: shutttttdown
Pull request: no

### System Info

- `transformers` version: 4.32.1
- Platform: Linux-5.4.0-164-generic-x86_64-with-glibc2.31
- Python version: 3.10.12
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.1
- Accelerate version: 0.21.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): 2.13.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed

### Who can help?

@stevhliu @MKhalusova @sgugger 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Hi everone,

I conducted an experiment on a machine with two A100 GPUs, and I encountered an issue. After simplifying, the scenario:

I set device_map='auto' and loaded two models, model1 and model2, both of which have the same structure. Then, I created a copy of model2, referred to as model2_copied, and attempted to load all its parameters from model1 using load_state_dict.

Upon comparing model2_copied after loading the parameters with model1, I found that their parameters were indeed identical. However, when I actually performed inference, I noticed that the inference results of model2_copied remained the same as model2, rather than matching my expectation of being the same as model1. 

My code is as below:
```python
from transformers import AutoModel, AutoTokenizer, LlamaForCausalLM, AutoModelForCausalLM, pipeline
import torch
import copy

model1_path = '...'
model1 = AutoModelForCausalLM.from_pretrained(model1_path, trust_remote_code=True, device_map='auto')

model2_path = '...'
model2 = AutoModelForCausalLM.from_pretrained(model2_path, trust_remote_code=True, device_map='auto')
model2_copied = copy.deepcopy(model2)

tokenizer = AutoTokenizer.from_pretrained(model2_path, use_fast=False, trust_remote_code=True)

# copy parameters from model1 to model2_copied
model2_copied.load_state_dict(model1.state_dict())
model2_copied.lm_head.load_state_dict(model1.lm_head.state_dict())

def compare_model_param(model1, model2):
    if str(model1) == str(model2):
        print("The model structures are identical.")
    else:
        print("The model structures are not identical.")

    state_dict1 = model1.state_dict()
    state_dict2 = model2.state_dict()
    param_names1 = set(state_dict1.keys())
    param_names2 = set(state_dict2.keys())

    Flag = False
    if param_names1 == param_names2:
        for param_name in param_names1:
            if not torch.equal(state_dict1[param_name], state_dict2[param_name]):
                Flag = True
                print(f"parameter {param_name} is not identical.")
    
    if Flag == False:
        print(f"All parameters are identical.")

compare_model_param(model1, model2_copied)
# The model structures are identical.
# All parameters are identical.

inp = tokenizer('Hello', return_tensors='pt', return_token_type_ids=False)
print(model1(**inp).logits)
print('*'*20)
print(model2(**inp).logits)
print('*'*20)
print(model2_copied(**inp).logits)

'''
tensor([[[11.2086, 11.2117, 25.7801,  ..., 12.2967, 11.0015, 11.6876]]],
       grad_fn=<ToCopyBackward0>)
********************
tensor([[[35.8617, 35.9438, 53.7408,  ..., 34.0522, 34.9409, 33.9223]]],
       grad_fn=<ToCopyBackward0>)
********************
tensor([[[35.8617, 35.9438, 53.7408,  ..., 34.0522, 34.9409, 33.9223]]],
       grad_fn=<ToCopyBackward0>)
'''
# Find that the output of model2_copied is the same as model2 rather than model1
```

However, when I loaded them without device_map='auto' and instead directly placed them on a GPU using .cuda(), the aforementioned issue did not exist, and the inference results of model2_copied matched those of model1.




### Expected behavior

See above,

I'm not sure if I made a mistake somewhere or if there was a misunderstanding. Thanks all.

Comments:
cc @SunMarc when you're back from leave!

Sorry for the late reply @shutttttdown, this has been fixed in in the version v0.24 (october 24) of accelerate through this [PR](https://github.com/huggingface/accelerate/pull/1971). LMK if it works on your end ! Thanks also for the great reproducer. The issue was that the hooks were not properly copied. Hence, they were still referencing the old forward (model_2) instead of the forward of `model2_copied`. 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.