Manually setting `device_map` causes a RuntimeError.
Author: Songjw133
Pull request: no
Labels: name: bug

### System Info

python version: 3.11.10
transformers version: 4.46.0
torch version: 2.4.0+cu118

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

from transformers import AutoTokenizer, AutoModelForCausalLM

model_path = "my_workspace/llama3-8B"

tokenizer = AutoTokenizer.from_pretrained(model_path, padding_side="left")

device_map = {}
device_map['model.embed_tokens'] = 0
for layer_idx in range(20):  
    device_map[f'model.layers.{layer_idx}'] = 0

for layer_idx in range(20, 32):
    device_map[f'model.layers.{layer_idx}'] = 1

device_map['lm_head.weight'] = 1
device_map['model.norm.weight'] = 1

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype="auto",
    device_map=device_map
)

print(model(**tokenizer("111",return_tensors="pt").to(0)).logits.shape)
```

With `transformers==4.46.0`, the code above results in:

> `RuntimeError`: Expected all tensors to be on the same device, but found at least two devices, CPU and CUDA:0! (when checking argument for mat2 in method wrapper_CUDA_bmm)


### Expected behavior

This issue does not occur with lower versions of transformers. I tried `transformers==4.40.0`, and the code successfully outputs:

> torch.Size([1, 1, 128256])


Comments:
The `rotary_emb` is moved to `LlamaModel` in 4.46: https://github.com/huggingface/transformers/blob/1d063793318b20654ebb850f48f43e0a247ab7bb/src/transformers/models/llama/modeling_llama.py#L852
https://github.com/huggingface/transformers/blob/1d063793318b20654ebb850f48f43e0a247ab7bb/src/transformers/models/llama/modeling_llama.py#L306
https://github.com/huggingface/transformers/commit/2e113422b3504fe6de821bb9911b24273b11aa9c.
That's explain the diff between 4.40 and 4.46.

For `RuntimeError`, I think it is related to `check_support_param_buffer_assignment` function.



cc @muellerzr for accelerate/device_map and @gante because I see him in the diff :eye: 

Hi @techkang @Songjw133 👋 

Since v4.40 we have indeed moved `self.rotary_emb` to the main model -- otherwise, we were holding a copy of the fixed parameters and computing the same position embeddings once per layer, which is inefficient

Naturally, the manual device map has to be updated accordingly: you must move `model.rotary_emb` to the desired device. If you add `device_map['model.rotary_emb'] = 1` to the example above, it runs as expected :)

I think it's a bug related to accelerate. However, the bug is complex and is rarely triggered.

Specifically, if a model has two layers and one parameter and sets the `device_map = {'layer0': 0, 'layer1': 1, 'param': 1}`, then layer weights will be loaded on GPU0 and GPU1 respectively. However, the param will be on GPU0!

The reason for this phenomenon  is when `device_map` is not None, `from_pretrained` will call `dispatch_model` to put weights on correct device:  [ref](https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L4303).If the `len(set(device_map.values())) = 1`, the `dispatch_model` will directly call `model.to(device)` which will automatically handle all parameters and buffers correctly: 
[ref](https://github.com/huggingface/accelerate/blob/main/src/accelerate/big_modeling.py#L494). However, if `len(set(device_map.values())) > 1`, it goes complicated. The function will call `attach_align_device_hook_on_blocks` to attach `AlignDevicesHook` to all blocks of the given model. The whole model will be attached to the hook because the default value of the `module_name` is `""`: [ref](https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L647-L654)

It seems that it only affect blocks and will only take effect during forward according to doc string of `AlignDevicesHook`. However, the hook has a `init_hook` function and will be called during attaching: 
[ref](https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L161). **In `init_hook` function, all the parameters (and buffers if `offload` is False) are set to `execution_device` which is not the device set in `device_map`!** :[ref](https://github.com/huggingface/accelerate/blob/main/src/accelerate/hooks.py#L283). The `execution_device` for the module is actually the first device in `device_map`:[ref0](https://github.com/huggingface/accelerate/blob/ba7ab93f5e688466ea56908ea3b056fae2f9a023/src/accelerate/big_modeling.py#L368) [ref1](https://github.com/huggingface/accelerate/blob/ba7ab93f5e688466ea56908ea3b056fae2f9a023/src/accelerate/big_modeling.py#L392).


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.

same problem