When I use Bloom, I get a error which is Caught RuntimeError in replica 0 on device 0.
Author: han508
Pull request: no

### System Info

- `transformers` version: 4.28.1
- Platform: Linux-5.4.0-148-generic-x86_64-with-glibc2.27
- Python version: 3.9.16
- Huggingface_hub version: 0.13.4
- Safetensors version: not installed
- PyTorch version (GPU?): 1.13.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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>
- GPU Tesla T4 *4

### Who can help?

@sgugger ,@ArthurZucker
Hello, when I use Bloom model, the following problem occurs, but when I use RedPajamam model or other models, this kind of error does not occur

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

My code:
from transformers import AutoTokenizer, AutoModelForCausalLM
import os
import torch
from datasets import load_dataset

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
tokenizer = AutoTokenizer.from_pretrained("Bigscience/bloom-560m",cache_dir='./cache/')
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
model = AutoModelForCausalLM.from_pretrained("Bigscience/bloom-560m",device_map='auto',cache_dir='./cache/',torch_dtype=torch.float16)
model.resize_token_embeddings(len(tokenizer))
model.gradient_checkpointing_enable()
model.config.use_cache = False

# dataset = load_dataset("BelleGroup/train_1M_CN")
# from datasets import load_dataset
# dataset = load_dataset("json", data_files="./data/alpaca_data_zh_51k.json")
dataset = load_dataset("json", data_files="/home/new_store/Llama/data/alpaca_gpt4_data_zh.json")
dataset = dataset.filter(lambda x: x["output"]!=None)
dataset= dataset.filter(lambda x: x["instruction"] !=None)
dataset= dataset.filter(lambda x: x["input"] !=None)
eval_dataset = load_dataset("json", data_files="split.json")
eval_dataset = eval_dataset.filter(lambda x: x["output"]!=None)
eval_dataset = eval_dataset.filter(lambda x: x["input"] !=None)
eval_dataset = eval_dataset.filter(lambda x: x["instruction"]!=None)


def preprocess_function(sample):
    l = "<##human>:"
    for i in range(len(sample['instruction'])):
        if sample['input'][i]!='':
            sample['instruction'][i]=sample['instruction'][i]+'[PAD]'+sample['input'][i]
            # print(sample['input'][i])
    output = ['<##bot>:'+i for i in sample['output']]
    model_inputs = tokenizer(sample['instruction'], truncation=True,padding=True,max_length=256)
    labels = tokenizer(output,  truncation=True, padding=True,max_length=256)
    model_inputs["labels"] = labels["input_ids"]
    # print(model_inputs)
    return model_inputs
input_data = dataset['train'].map(preprocess_function,batched=True,remove_columns=['instruction','input','output'])
eval_data = eval_dataset['train'].map(preprocess_function,batched=True,remove_columns=['instruction','input','output'])


from transformers import TrainingArguments, Trainer, DataCollatorForLanguageModeling
trainArgs = TrainingArguments(
  output_dir= '../ckps_bloom_1M',
    do_train=True,
    # per_device_train_batch_size=1,
    auto_find_batch_size=True,
    gradient_accumulation_steps=4,
    evaluation_strategy="steps",
    save_strategy="steps",
    save_steps=500,
    eval_steps=500,
    logging_steps=100,
    warmup_steps=100,
    num_train_epochs=2,
    learning_rate=2e-5,
    #fp16=True,  
    # bf16=True, 
    load_best_model_at_end=True,
    #deepspeed= './zero.json',
    report_to="wandb"
)
trainer = Trainer(
    model=model,
    args=trainArgs,
    train_dataset=input_data,
    eval_dataset=eval_data,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.train()

Error:
![Uploading image.png…]()
RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/han/anaconda3/envs/llama/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py", line 64, 
in _worker
    output = module(*input, **kwargs)
  File "/home/han/anaconda3/envs/llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in 
_call_impl
    return forward_call(*input, **kwargs)
  File "/home/han/anaconda3/envs/llama/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/han/anaconda3/envs/llama/lib/python3.9/site-packages/transformers/models/bloom/modeling_bloom.py", 
line 913, in forward
    transformer_outputs = self.transformer(
  File "/home/han/anaconda3/envs/llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in 
_call_impl
    return forward_call(*input, **kwargs)
  File "/home/han/anaconda3/envs/llama/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/han/anaconda3/envs/llama/lib/python3.9/site-packages/transformers/models/bloom/modeling_bloom.py", 
line 730, in forward
    inputs_embeds = self.word_embeddings(input_ids)
  File "/home/han/anaconda3/envs/llama/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in 
_call_impl
    return forward_call(*input, **kwargs)
  File "/home/han/anaconda3/envs/llama/lib/python3.9/site-packages/torch/nn/modules/sparse.py", line 160, in 
forward
    return F.embedding(
  File "/home/han/anaconda3/envs/llama/lib/python3.9/site-packages/torch/nn/functional.py", line 2210, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
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 index in method wrapper__index_select)
### Expected behavior

I hope the error can be fixed

Comments:
I don't think `device_map="auto"` is compatible with gradient checkpointing.

Thanks for the response. I've cancelled the gradient checkpointing, but the problem still exists

Since we don't have access to your data files, it's going to be pretty hard to reproduce the issue. Could you:
1. format your code so we can copy/paste it
2. use a dataset from the Hub instead so we can replicate
Thanks!

Thank you again for your patient reply. I have modified the data set on Hub and formatted the code. The code is as follows.
```python
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from datasets import load_dataset
# tokenizer = AutoTokenizer.from_pretrained("Bigscience/bloom-560m",cache_dir='./cache/')
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m",cache_dir='./cache/')
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
tokenizer.padding_side='right'

model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-70m",device_map='balanced',cache_dir='./cache/',torch_dtype=torch.float16)
# model = AutoModelForCausalLM.from_pretrained("Bigscience/bloom-560m",device_map='balanced',cache_dir='./cache/',torch_dtype=torch.float16)
model.resize_token_embeddings(len(tokenizer))
model.config.use_cache = True

dataset = load_dataset("c-s-ale/alpaca-gpt4-data-zh")

def preprocess_function(sample):
    for i in range(len(sample['instruction'])):
        sample['instruction'][i]=sample['instruction'][i]+'[PAD]'+sample['input'][i]
    output = ['<bot>:'+i for i in sample['output']]
    model_inputs = tokenizer(sample['instruction'], truncation=True,padding=True,max_length=100,return_tensors="pt")
    labels = tokenizer(output,  truncation=True, padding=True,max_length=100,return_tensors="pt")
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

input_data = dataset['train'].map(preprocess_function,batched=True,remove_columns=['instruction','input','output'])



from transformers import TrainingArguments, Trainer, DataCollatorForLanguageModeling
trainArgs = TrainingArguments(
  output_dir= './ckps_bloom',
    do_train=True,
    auto_find_batch_size=True,
    gradient_accumulation_steps=4,
    evaluation_strategy="steps",
    save_strategy="steps",
    save_steps=10,
    eval_steps=10,
    logging_steps=10,
    warmup_steps=100,
    num_train_epochs=2,
    learning_rate=2e-5,
    fp16=True,  
    load_best_model_at_end=True,
    push_to_hub=False,
    report_to="wandb"
)


trainer = Trainer(
    model=model,
    args=trainArgs,
    train_dataset=input_data,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.train()
```


Thanks for sharing. Note that in any case, you won't be able to train your model in float16 (you will get an error in the line of "Attempting to unscale FP16 gradients.". Training in float16 does not converge so the Trainer does not support it. You will need to remove the line `torch_dtype=torch.float16` when loading your model.

For the Pythia model, something weird is happening with `device_map="auto"` since the model is so tiny: it is all placed on GPU-1 (in my case) and then the Trainer tries to move it to GPU-0. Will fix this but a simple workaround in the meantime is to `place_model_on_device=False` in your training arguments.

Thank you again for your reply. I removed the line torch_dtype =torch.float16 and set place_model_on_device=False, but the problem still exists. This problem existed regardless of the size of the model, and I also tried to use bloom 3b from the Hub.
When I removed device_map='auto', the program worked, but only on one GPU.
The code is as follows.
```python
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from datasets import load_dataset
tokenizer = AutoTokenizer.from_pretrained("Bigscience/bloom-3b",cache_dir='./cache/')
# tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m",cache_dir='./cache/')
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
tokenizer.padding_side='right'

# model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-70m",device_map='balanced',cache_dir='./cache/')
model = AutoModelForCausalLM.from_pretrained("Bigscience/bloom-3b",device_map='auto',cache_dir='./cache/')
model.resize_token_embeddings(len(tokenizer))
model.config.use_cache = True

dataset = load_dataset("c-s-ale/alpaca-gpt4-data-zh")

def preprocess_function(sample):
    for i in range(len(sample['instruction'])):
        sample['instruction'][i]=sample['instruction'][i]+'[PAD]'+sample['input'][i]
    output = ['<bot>:'+i for i in sample['output']]
    model_inputs = tokenizer(sample['instruction'], truncation=True,padding=True,max_length=100,return_tensors="pt")
    labels = tokenizer(output,  truncation=True, padding=True,max_length=100,return_tensors="pt")
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

input_data = dataset['train'].map(preprocess_function,batched=True,remove_columns=['instruction','input','output'])



from transformers import TrainingArguments, Trainer, DataCollatorForLanguageModeling
trainArgs = TrainingArguments(
  output_dir= './ckps_bloom',
    do_train=True,
    auto_find_batch_size=True,
    gradient_accumulation_steps=4,
    evaluation_strategy="steps",
    save_strategy="steps",
    save_steps=10,
    eval_steps=10,
    logging_steps=10,
    warmup_steps=100,
    num_train_epochs=2,
    learning_rate=2e-5,
    fp16=True,  
    load_best_model_at_end=True,
    push_to_hub=False,
    report_to="wandb",
)
TrainingArguments.place_model_on_device=False

trainer = Trainer(
    model=model,
    args=trainArgs,
    train_dataset=input_data,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.train()
```

GPU memory allocation is also weird. In the past, it was pretty even, but now it looks like this.

<img width="402" alt="image" src="https://github.com/huggingface/transformers/assets/69674181/ed1b7c87-cacc-42d5-a327-1b8f3962b0fc">

error:
<img width="1207" alt="image" src="https://github.com/huggingface/transformers/assets/69674181/26494845-b675-499e-961f-af02d15ce264">


I didn't understand what you meant before. I'm sorry, but it has been solved now.

I have the same problem but the solution is still not clear for me, can you specify @han508  ? Thanks,