T5Model crashes when trained with multiple GPUs
Author: eladyt
Pull request: no

## Environment info
- `transformers` version: 4.5.1
- Platform: Linux-5.4.0-1041-azure-x86_64-with-debian-buster-sid
- Python version: 3.7.9
- PyTorch version (GPU?): 1.8.1+cu102 (True)
- Tensorflow version (GPU?): 2.4.1 (False)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes

### Who can help
@patrickvonplaten, @patil-suraj

## Information

I'm training a T5 translation model. It works on a CPU or single GPU, but when I try to run it with multiple GPUs I get the following error:
Traceback (most recent call last):
  File "train2.py", line 43, in <module>
    model.train_model(train_df, eval_data=eval_df)
  File "/home/eladyt/.local/lib/python3.7/site-packages/simpletransformers/t5/t5_model.py", line 206, in train_model
    **kwargs,
  File "/home/eladyt/.local/lib/python3.7/site-packages/simpletransformers/t5/t5_model.py", line 605, in train
    **kwargs,
  File "/home/eladyt/.local/lib/python3.7/site-packages/simpletransformers/t5/t5_model.py", line 705, in eval_model
    result = self.evaluate(eval_dataset, output_dir, verbose=verbose, silent=silent, **kwargs)
  File "/home/eladyt/.local/lib/python3.7/site-packages/simpletransformers/t5/t5_model.py", line 763, in evaluate
    outputs = model(**inputs)
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 167, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 177, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
    output.reraise()
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/_utils.py", line 429, in reraise
    raise self.exc_type(msg)
RuntimeError: Caught RuntimeError in replica 1 on device 1.
Original Traceback (most recent call last):
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 167, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 177, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
    output.reraise()
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/_utils.py", line 429, in reraise
    raise self.exc_type(msg)
RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/eladyt/.local/lib/python3.7/site-packages/transformers/models/t5/modeling_t5.py", line 1506, in forward
    return_dict=return_dict,
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/eladyt/.local/lib/python3.7/site-packages/transformers/models/t5/modeling_t5.py", line 881, in forward
    inputs_embeds = self.embed_tokens(input_ids)
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/modules/sparse.py", line 158, in forward
    self.norm_type, self.scale_grad_by_freq, self.sparse)
  File "/home/eladyt/.local/lib/python3.7/site-packages/torch/nn/functional.py", line 1916, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: Input, output and indices must be on the current device

## To reproduce

Steps to reproduce the behavior:

Here is the code (from https://towardsdatascience.com/how-to-train-an-mt5-model-for-translation-with-simple-transformers-30ba5fa66c5f):

```python

import logging
import pandas as pd
from simpletransformers.t5 import T5Model, T5Args

import torch.multiprocessing
torch.multiprocessing.set_sharing_strategy('file_system')

logging.basicConfig(level=logging.INFO)
transformers_logger = logging.getLogger("transformers")
transformers_logger.setLevel(logging.WARNING)

train_df = pd.read_csv("data2/train.tsv", sep="\t").astype(str)
eval_df = pd.read_csv("data2/eval.tsv", sep="\t").astype(str)

train_df["prefix"] = ""
eval_df["prefix"] = ""

model_args = T5Args()
model_args.max_seq_length = 25
model_args.train_batch_size = 20
model_args.eval_batch_size = 20
model_args.num_train_epochs = 20
model_args.evaluate_during_training = True
model_args.evaluate_during_training_steps = 30000
model_args.use_multiprocessing = False
model_args.fp16 = True
model_args.save_steps = -1
model_args.save_eval_checkpoints = False
model_args.no_cache = True
model_args.reprocess_input_data = True
model_args.overwrite_output_dir = True
model_args.preprocess_inputs = True
model_args.num_return_sequences = 1
model_args.n_gpu=4
model_args.is_model_parallel = True

model = T5Model("mt5", "google/mt5-base", args=model_args)

model.train_model(train_df, eval_data=eval_df)
```

## Expected behavior
A model should be generated. 

Comments:
This code seems to be using `simpletransformers`, sadly we won't dive into that. You could use the `run_translation.py` script [here](https://github.com/huggingface/transformers/tree/master/examples/seq2seq) which supports multi-gpu training. See this [doc](https://github.com/huggingface/transformers/tree/master/examples#distributed-training-and-mixed-precision) for distributed training using `Trainer`.

And with `Trainer` you could also leverage `deepspeed` to get more efficiency, see this [doc](https://huggingface.co/transformers/main_classes/trainer.html#trainer-integrations) for `deepspeed` integration 

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/master/CONTRIBUTING.md) are likely to be ignored.