Device errors when loading in 8 bit 
Author: cassianlewis
Pull request: no

### System Info

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

- `transformers` version: 4.31.0
- Platform: Linux-5.10.178-162.673.amzn2.x86_64-x86_64-with-glibc2.26
- Python version: 3.10.10
- 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.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes (4 GPUs)
- Using distributed or parallel set-up in script?: 

### Who can help?

@younesbelkada 
@sgugger 
@mue

### 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

This error occurs when trying to split a quantised `t5-large` model (or any t5 model for that matter) across 4 GPUs using a custom device map (which works when it is not quantised)!

Steps to reproduce:

1. 
```
from transformers import AutoTokenizer, DataCollatorWithPadding, TrainingArguments, Trainer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
from peft import get_peft_config, get_peft_model, PromptTuningInit, PromptTuningConfig, TaskType, PeftType
from torch.utils.data import TensorDataset, DataLoader,Dataset
from accelerate import dispatch_model, infer_auto_device_map, init_empty_weights
from accelerate.utils import get_balanced_memory

model_name = "t5-large"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name, cache_dir = 'models', load_in_8bit=True) 
```

2.
```
max_memory = get_balanced_memory(
    model,
    max_memory=None,
    no_split_module_classes=["T5Block"],
    dtype='float16',
    low_zero=False,
)
```
max_memory: 
`{0: 263982848, 1: 263982848, 2: 263982848, 3: 13860929536, 'cpu': 189321494528}`

3.
```
device_map = infer_auto_device_map(
    model,
    max_memory=max_memory,
    no_split_module_classes=["T5Block"],
    dtype='float16'
)
```
I won't show the entire device_map, just the important part:
```
{'shared': 0,
 'decoder.embed_tokens': 0,
 'encoder.embed_tokens': 0,
 'lm_head': 0,
 'encoder.block.0': 0,
 'encoder.block.1': 0,
 'encoder.block.2': 0,
 'encoder.block.3': 0,
 'encoder.block.4': 0,
 'encoder.block.5': 0,
 'encoder.block.6': 0,
 'encoder.block.7': 0,
 'encoder.block.8': 0,
 'encoder.block.9': 0,
 'encoder.block.10': 1,
 'encoder.block.11': 1,
 'encoder.block.12': 1,
```

4.

```
model = dispatch_model(model, device_map=device_map)
for i in model.named_parameters():
    print(f"{i[0]} -> {i[1].device}")
```

Again, just the pertinent part:

```
encoder.block.10.layer.0.SelfAttention.q.weight -> cuda:0
encoder.block.10.layer.0.SelfAttention.k.weight -> cuda:0
encoder.block.10.layer.0.SelfAttention.v.weight -> cuda:0
encoder.block.10.layer.0.SelfAttention.o.weight -> cuda:0
encoder.block.10.layer.0.layer_norm.weight -> cuda:0
encoder.block.10.layer.1.DenseReluDense.wi.weight -> cuda:0
encoder.block.10.layer.1.DenseReluDense.wo.weight -> cuda:0
encoder.block.10.layer.1.layer_norm.weight -> cuda:0
encoder.block.11.layer.0.SelfAttention.q.weight -> cuda:1
encoder.block.11.layer.0.SelfAttention.k.weight -> cuda:1
encoder.block.11.layer.0.SelfAttention.v.weight -> cuda:1
encoder.block.11.layer.0.SelfAttention.o.weight -> cuda:1
encoder.block.11.layer.0.layer_norm.weight -> cuda:1
encoder.block.11.layer.1.DenseReluDense.wi.weight -> cuda:1
encoder.block.11.layer.1.DenseReluDense.wo.weight -> cuda:1
encoder.block.11.layer.1.layer_norm.weight -> cuda:1
```

5.

```
batch = tokenizer("Hello World", return_tensors="pt")
model(**batch, decoder_input_ids = batch['input_ids'])
```


### Expected behavior

Error:

```
File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/transformers/models/t5/modeling_t5.py:260, in T5LayerNorm.forward(self, hidden_states)
    257 if self.weight.dtype in [torch.float16, torch.bfloat16]:
    258     hidden_states = hidden_states.to(self.weight.dtype)
--> 260 return self.weight * hidden_states

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!
```

Note that repeating this with `load_in_8bit = False` works normally.
    
Thanks!

Comments:
You cannot re-dispatch a model that was loaded in 8bit. You need to pass along your `max_memory` or `device_map` to the call to `from_pretrained`.

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.

The explanation above makes sense to me! Feel free to re-open the issue if you think that doesn't answer your quesiton