dataset not being sent to device when using Trainer (distributed)
Author: KennethEnevoldsen
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 4.1.1
- Platform: Linux-4.15.0-99-generic-x86_64-with-glibc2.27
- Python version: 3.8.0
- PyTorch version (GPU?): 1.7.1 (True)
- Tensorflow version (GPU?): 2.4.0 (False)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: distrubuted

### Who can help
 Trainer: @sgugger
 Text Generation /t5: @patrickvonplaten
 examples/seq2seq: @patil-suraj


## Information
The intention is to train a t5 model (preferably as last as possible) in a distributed setting using the HF Trainer. However, when setting the model_parallel to True the training breaks.

related issues might be:
https://github.com/huggingface/transformers/issues/9229
https://github.com/huggingface/transformers/issues/6821

However, do note that the script works perfectly fine training on multiple GPU in a non distributed fashion (setting model_parallel to False).

## To reproduce

I have created a minimal script for reproducing the behavior:

```
import transformers
from transformers import (
    MT5ForConditionalGeneration,
    MT5Model,
    Trainer,
    TrainingArguments,
    T5Tokenizer
)
```

<details>
  <summary> Click to see (minimal) dataset creation </summary>


```
import datasets

# making minimal test for example sake
def make_test_dataset(tokenizer="google/mt5-small"):
    if isinstance(tokenizer, str):
        tokenizer = T5Tokenizer.from_pretrained(tokenizer)

    ds = datasets.load_dataset("dane")

    def __tokenizer_input(batch):
        return tokenizer(batch['text'],
                         padding="max_length",
                         max_length=256,  # actual max is 235
                         truncation=True)

    def __tokenizer_output(batch):
        tok = tokenizer(batch['text'],
                        padding="max_length",
                        max_length=256,
                        truncation=True)
        tok["labels"] = tok.pop("input_ids")
        return tok

    # filter out empty strings (bug reported and fixed)
    ds = ds.filter(lambda batch: bool(batch["text"]))

    # tokenize both datasets (eos:   
 is added by tokenizer)
    ds = ds.map(__tokenizer_input, batched=True, batch_size=len(ds))
    ds = ds.map(__tokenizer_output, batched=True, batch_size=len(ds))

    return ds

dataset = make_test_dataset()
dataset.set_format(type='torch',
                   columns=['input_ids',
                            'attention_mask',
                            'labels'])
```

</details>


```
model = MT5ForConditionalGeneration.from_pretrained("google/mt5-small") #using a small model for example
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=1,
    logging_dir='./logs',
    evaluation_strategy="epoch",
    model_parallel=True # work fine when set to False
)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"]
)
trainer.remove_callback(transformers.integrations.WandbCallback) # removing wandb for conv.
trainer.train()
```


Results:

```
RuntimeError: Input, output and indices must be on the current device
```

<details>
  <summary> Click to see full traceback </summary>


```
  RuntimeError                              Traceback (most recent call last)
  ~/github/EDP-Efficient-Danish-Preprocessing/tmp.py in 
       64 
       65 trainer.remove_callback(transformers.integrations.WandbCallback)
  ---> 66 trainer.train()
  
  ~/.Envs/EDP/lib/python3.8/site-packages/transformers/trainer.py in train(self, model_path, trial)
      797                         tr_loss += self.training_step(model, inputs)
      798                 else:
  --> 799                     tr_loss += self.training_step(model, inputs)
      800                 self._total_flos += self.floating_point_ops(inputs)
      801 
  
  ~/.Envs/EDP/lib/python3.8/site-packages/transformers/trainer.py in training_step(self, model, inputs)
     1137                 loss = self.compute_loss(model, inputs)
     1138         else:
  -> 1139             loss = self.compute_loss(model, inputs)
     1140 
     1141         if self.args.n_gpu > 1:
  
  ~/.Envs/EDP/lib/python3.8/site-packages/transformers/trainer.py in compute_loss(self, model, inputs)
     1161         Subclass and override for custom behavior.
     1162         """
  -> 1163         outputs = model(**inputs)
     1164         # Save past state if it exists
     1165         # TODO: this needs to be fixed and made cleaner later.
  
  ~/.Envs/EDP/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
      725             result = self._slow_forward(*input, **kwargs)
      726         else:
  --> 727             result = self.forward(*input, **kwargs)
      728         for hook in itertools.chain(
      729                 _global_forward_hooks.values(),
  
  ~/.Envs/EDP/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py in forward(self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, encoder_outputs, past_key_values, head_mask, inputs_embeds, decoder_inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict)
     1422         if encoder_outputs is None:
     1423             # Convert encoder inputs in embeddings if needed
  -> 1424             encoder_outputs = self.encoder(
     1425                 input_ids=input_ids,
     1426                 attention_mask=attention_mask,
  
  ~/.Envs/EDP/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
      725             result = self._slow_forward(*input, **kwargs)
      726         else:
  --> 727             result = self.forward(*input, **kwargs)
      728         for hook in itertools.chain(
      729                 _global_forward_hooks.values(),
  
  ~/.Envs/EDP/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py in forward(self, input_ids, attention_mask, encoder_hidden_states, encoder_attention_mask, inputs_embeds, head_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict)
      858         if inputs_embeds is None:
      859             assert self.embed_tokens is not None, "You have to initialize the model with valid token embeddings"
  --> 860             inputs_embeds = self.embed_tokens(input_ids)
      861 
      862         batch_size, seq_length = input_shape
  
  ~/.Envs/EDP/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
      725             result = self._slow_forward(*input, **kwargs)
      726         else:
  --> 727             result = self.forward(*input, **kwargs)
      728         for hook in itertools.chain(
      729                 _global_forward_hooks.values(),
  
  ~/.Envs/EDP/lib/python3.8/site-packages/torch/nn/modules/sparse.py in forward(self, input)
      122 
      123     def forward(self, input: Tensor) -> Tensor:
  --> 124         return F.embedding(
      125             input, self.weight, self.padding_idx, self.max_norm,
      126             self.norm_type, self.scale_grad_by_freq, self.sparse)
  
  ~/.Envs/EDP/lib/python3.8/site-packages/torch/nn/functional.py in embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)
     1850         # remove once script supports set_grad_enabled
     1851         _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)
  -> 1852     return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
     1853 
     1854 
  
  RuntimeError: Input, output and indices must be on the current device
```
</details>

## Expected behavior
A training model

PS: I personally found that the model_parallel name was slightly confusing. I assume a more fitting name would be model_distributed (but this is a minor thing)

Thanks for great work


Comments:
The `model_parallel` argument has nothing to do with training in a parallel fashion (and is going to be deleted very soon since you're not the first user its name confuses). To use parallel training with:
- PyTorch DataParallel, there is nothing to do, the Trainer does it automatically
- PyTorch DistributedDataParallel, you should launch your script with the `python -m torch.distributed.launch` command (see the examples).

Glad to hear about the name.

I was aware that DataParallel was used by default.

By examples you must refer to: https://huggingface.co/transformers/examples.html
but it doesn't seem to provide any information on `torch.distributed.launch` am I missing something?

will try to run with flag, but would love if I could find some documentation on this

Thanks for the quick response


you can find the commands/docs to launch distributed training in the examples [readme](https://github.com/huggingface/transformers/tree/master/examples#distributed-training-and-mixed-precision)

Thanks, I will try this out as soon as our GPU's are available again!

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.