Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:3!
### System Info

- `transformers` version: 4.29.2
- Platform: Linux-5.4.0-137-generic-x86_64-with-glibc2.31
- Python version: 3.11.4
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1 (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
- Using distributed or parallel set-up in script?: Yes

### Who can help?

_No response_

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

I am trying to finetune an [InstructCodeT5+](https://huggingface.co/Salesforce/instructcodet5p-16b) model on some training data using a multi-GPU setup. The same code (given further below) seems to work in a single-GPU setting (when i set `CUDA_VISIBLE_DEVICES=0`):

```
CUDA_VISIBLE_DEVICES=0,1,2,3 python training.py

===================================BUG REPORT===================================
Welcome to bitsandbytes. For bug reports, please run

python -m bitsandbytes

 and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues
================================================================================
bin /root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/bitsandbytes/libbitsandbytes_cuda118.so
CUDA SETUP: CUDA runtime path found: /root/anaconda3/envs/datachat_env/lib/libcudart.so.11.0
CUDA SETUP: Highest compute capability among GPUs detected: 8.0
CUDA SETUP: Detected CUDA version 118
CUDA SETUP: Loading binary /root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/bitsandbytes/libbitsandbytes_cuda118.so...
[nltk_data] Downloading package punkt to /root/nltk_data...
[nltk_data]   Package punkt is already up-to-date!
The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function.
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:24<00:00,  4.92s/it]
/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/transformers/optimization.py:407: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning
  warnings.warn(
  0%|                                                                                                                                                                                                                                | 0/17409 [00:00<?, ?it/s]You're using a CodeGenTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.
Traceback (most recent call last):
  File "/root/Custom-LLM/training.py", line 364, in <module>
    main()
  File "/root/Custom-LLM/training.py", line 336, in main
    trainer.train()
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/transformers/trainer.py", line 1664, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/transformers/trainer.py", line 1940, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/transformers/trainer.py", line 2735, in training_step
    loss = self.compute_loss(model, inputs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/transformers/trainer.py", line 2767, in compute_loss
    outputs = model(**inputs)
              ^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/huggingface/modules/transformers_modules/instructcodet5p-16b/modeling_codet5p.py", line 904, in forward
    encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/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:0 and cuda:3! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)
```
Code for the above error is given below:

```
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import pandas as pd
import os
import torch
from peft import TaskType
from transformers import DataCollatorForSeq2Seq, Seq2SeqTrainer, Seq2SeqTrainingArguments, BitsAndBytesConfig
import evaluate
import nltk
import numpy as np
from nltk.tokenize import sent_tokenize
from datasets import Dataset, DatasetDict
import argparse
import pickle
import json
import statistics
import ast
from copy import deepcopy

device = 'cuda'

parser = argparse.ArgumentParser(description='Options')
parser.add_argument('--dataset_dir', default='data', type=str, help="folder in which the dataset is stored")
parser.add_argument('--output_dir', default="lora-instructcodet5p", type=str, help="output directory for the model")
parser.add_argument('--results_dir', default="results", type=str, help="where the results should be stored")
args = parser.parse_args()

tokenized_dataset = DatasetDict.load_from_disk(args.dataset_dir)

pad_tok = 50256
token_id="Salesforce/instructcodet5p-16b"

tokenizer = AutoTokenizer.from_pretrained(token_id)


def main():
    # huggingface hub model id
    model_id="instructcodet5p-16b"
    if not os.path.exists(model_id):
        model_id=token_id

    # load model from the hub
    model = AutoModelForSeq2SeqLM.from_pretrained(model_id,
                                              torch_dtype=torch.bfloat16,
                                              low_cpu_mem_usage=True,
                                              trust_remote_code=True, decoder_start_token_id=1, pad_token_id=pad_tok, device_map="auto").to(device)
    

    # we want to ignore tokenizer pad token in the loss
    label_pad_token_id = pad_tok
    # Data collator
    data_collator = DataCollatorForSeq2Seq(
        tokenizer,
        model=model,
        label_pad_token_id=label_pad_token_id,
        pad_to_multiple_of=8
    )

    output_dir=args.output_dir

    training_args = Seq2SeqTrainingArguments(
        output_dir=output_dir,
        per_device_train_batch_size=1,
        predict_with_generate=True,
        weight_decay=0.05,
        warmup_steps=100,
        fp16=False, 
        learning_rate=1e-3,
        num_train_epochs=3,
        logging_dir=f"{output_dir}/logs",
        logging_strategy="epoch",
        save_strategy="no",
        report_to="tensorboard",
        push_to_hub=False,
        generation_max_length=200,
        include_inputs_for_metrics = True,
        lr_scheduler_type = 'cosine'
    )

    # Create Trainer instance
    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=tokenized_dataset["train"]
    )

    # train model
    trainer.train()

if __name__ == '__main__':
    main()
```


### Expected behavior

Expected behavior is that the model should train in a multi-GPU setting without throwing any errors. The same script works in single-GPU setting but throws the above error in a multi-GPU setting
You cannot use `device_map="auto"` with the `to` method afterward as you do. The model will be split up on the GPUs already.
Also, how are you launching your training script after?<|||||>Getting the same error when i remove the `to` method. Traceback given below. Also, I launch the training script using `CUDA_VISIBLE_DEVICES=0,1,2,3 python training.py`

```
Traceback (most recent call last):
  File "/root/Custom-LLM/new_train.py", line 91, in <module>
    main()
  File "/root/Custom-LLM/new_train.py", line 88, in main
    trainer.train()
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/transformers/trainer.py", line 1645, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/transformers/trainer.py", line 1938, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/transformers/trainer.py", line 2759, in training_step
    loss = self.compute_loss(model, inputs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/transformers/trainer.py", line 2784, in compute_loss
    outputs = model(**inputs)
              ^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/huggingface/modules/transformers_modules/instructcodet5p-16b/modeling_codet5p.py", line 932, in forward
    loss = loss_fct(logits.reshape(-1, self.decoder.config.vocab_size), labels.view(-1))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/torch/nn/modules/loss.py", line 1174, in forward
    return F.cross_entropy(input, target, weight=self.weight,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/anaconda3/envs/datachat_env/lib/python3.11/site-packages/torch/nn/functional.py", line 3029, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:3 and cuda:0! (when checking argument for argument target in method wrapper_CUDA_nll_loss_forward)
  0%|                                                                                                                                   | 0/17409 [00:02<?, ?it/s]
```<|||||>Ah, the problem lies in the custom code of this model. You need to move the `labels` to the device of the logits [here](https://huggingface.co/Salesforce/instructcodet5p-16b/blob/main/modeling_codet5p.py#L930) by adding `labels = labels.to(logits.device)`. Your logits are on the last GPU but your labels are still on the first one.<|||||>I made a [PR](https://huggingface.co/Salesforce/instructcodet5p-16b/discussions/4) with this suggestion on the repo. You can check it out locally by adding the `revision="pr_4"` argument when loading the model.<|||||>Thanks a lot! That solved the issue