LLaMA-2 runtime error after `resize_token_embeddings`
Author: gargutsav
Pull request: no

### System Info

Python: `3.8.17`
Transformers: `4.31.0`
Accelerate: `0.21.0`
GPUs: `8 x A100 (80GB)`

### Who can help?

@ArthurZucker @pacman100 

### 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

1. Load `LLaMA-13B` with the following arguments using `AutoModelForCausalLM.from_pretrained`
> device_map="auto"
> torch_dtype=torch.bfloat16 
2. Add special tokens to the tokenizer, and do `model.resize_token_embeddings(len(tokenizer))`
3. Error shows in the forward pass
```
  File "/opt/conda/envs/llm/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 824, in forward
    logits = self.lm_head(hidden_states)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:6 and cuda:7! (when checking argument for argument mat2 in method wrapper_CUDA_mm)
```

### Expected behavior

When no additional tokens are added to the tokenizer, then there is no error. So the error probably originates from the `device_map` being incorrect after `resize_token_embeddings`. 

Comments:
Hi @gargutsav

Could you provide a self-contained code snippet please? Thanks!

@ydshieh I was able to recreate it with this snippet
```python
import torch
from datasets import load_dataset
from transformers import (AutoModelForCausalLM, AutoTokenizer,
                          DataCollatorForLanguageModeling, Trainer,
                          TrainingArguments)

if __name__ == "__main__":
    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-2-13b-hf",
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        device_map="auto",
    )

    tokenizer = AutoTokenizer.from_pretrained(
        "meta-llama/Llama-2-13b-hf", use_fast=False
    )
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.add_special_tokens(
        {"additional_special_tokens": ["
", "
", "
", "
"]}
    )
    model.resize_token_embeddings(len(tokenizer))

    # Prepare dataset
    eli5 = load_dataset("eli5", split="train_asks[:5000]")
    eli5 = eli5.train_test_split(test_size=0.2)
    eli5 = eli5.flatten()

    def preprocess_function(examples):
        return tokenizer(
            [" ".join(x) for x in examples["answers.text"]],
            truncation=True,
            max_length=1024,
        )

    tokenized_eli5 = eli5.map(
        preprocess_function,
        batched=True,
        num_proc=4,
        remove_columns=eli5["train"].column_names,
    )
    block_size = 128

    def group_texts(examples):
        concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
        total_length = len(concatenated_examples[list(examples.keys())[0]])
        result = {
            k: [t[i: i + block_size] for i in range(0, total_length, block_size)]
            for k, t in concatenated_examples.items()
        }
        result["labels"] = result["input_ids"].copy()
        return result

    lm_dataset = tokenized_eli5.map(group_texts, batched=True, num_proc=4)

    training_arguments = TrainingArguments(
        output_dir="test",
        group_by_length=True,
        lr_scheduler_type="cosine",
        bf16=True,
        report_to="none",
    )
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    trainer = Trainer(
        model=model,
        args=training_arguments,
        train_dataset=lm_dataset["train"],
        eval_dataset=lm_dataset["test"],
        data_collator=data_collator,
    )
    trainer.train()
```
It works okay, if we get rid of the `additional_tokens` part
```python
tokenizer.add_special_tokens({"additional_special_tokens": ["
", "
", "
", "
"]})
model.resize_token_embeddings(len(tokenizer))
```

Hello, cc @SunMarc who knows more about `device_map`

Hi @gargutsav, I think the problem is that by resizing the token_embedding, it will modify both the tokens_embedding and the lm_head. In this process, the hooks are were dispatched previously (using device_map = 'auto') on these modules no longer exist. Thus, we have this issue. I will try to fix this behavior in a future PR !

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.

@SunMarc is this fixed in https://github.com/huggingface/transformers/pull/25596? What version of `transformers` has this update?

> @SunMarc is this fixed in #25596? What version of `transformers` has this update?

nvm see it's fixed in 4.32.0, thank you!

I am on transformers version 4.35.2  and get `ValueError: Trying to set a tensor of shape torch.Size([32000, 4096]) in "weight" (which has shape torch.Size([32001, 4096])), this look incorrect.
`

When I try to run:
```
from transformers import AutoTokenizer, AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(model_name_or_path, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

tokenizer.add_tokens(["<bla>"])
model.resize_token_embeddings(len(tokenizer))
out = model.generate(tokenizer.encode('The <bla> quick brown fox jumps over the lazy dog', return_tensors='pt'), max_length=50, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
```