Repository: huggingface/transformers

Issue:
[s2s] only save metrics.json from rank zero
<!-- This line specifies which issue to close after the pull request is merged. -->
Fixes #{issue number}

Pytorch 1.6 DataParallel
I get an error similar to that of [#4189](https://github.com/huggingface/transformers/issues/4189) and [#3936](https://github.com/huggingface/transformers/issues/3936), when using DataParallel with GPT2.  
The issue should be resolved in newer versions of transformers, but I still get an error. 

## Environment info
     
- `transformers` version: 3.3.1
- Platform: Linux-3.10.0-1062.18.1.el7.x86_64-x86_64-with-glibc2.10
- Python version: 3.8.5
- PyTorch version (GPU?): 1.6.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes

CUDA compilation tools, release 10.2, V10.2.89

### Who can help
albert, bert, GPT2, XLM: @LysandreJik
TextGeneration: @TevenLeScao

## Information

Model I am using: GPT2

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behaviour:
1. Run forward on GPT2 with 2 GPUs


### Small example

```python
import torch
from torch.nn import DataParallel
from transformers import GPT2Tokenizer, GPT2LMHeadModel

device = "cuda:0"

# Get model
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")
model = DataParallel(model, device_ids=list(range(torch.cuda.device_count())))
model.to(device=device)

# Run forward
inputs = tokenizer(["This is an example"], return_tensors="pt")
outputs = model(
    input_ids=inputs["input_ids"].to(device),
    attention_mask=inputs["attention_mask"].to(device),
    labels=inputs["input_ids"].to(device),
)

print(f"outputs: {outputs}")
print("Success.")
```

### Output

```
Traceback (most recent call last):
  File "minimum_example.py", line 15, in <module>
    outputs = model(
  File "/home/user/.conda/envs/main/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/user/.conda/envs/main/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 155, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/home/user/.conda/envs/main/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 165, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/home/user/.conda/envs/main/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 85, in parallel_apply
    output.reraise()
  File "/home/user/.conda/envs/main/lib/python3.8/site-packages/torch/_utils.py", line 395, in reraise
    raise self.exc_type(msg)
StopIteration: Caught StopIteration in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/user/.conda/envs/main/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 60, in _worker
    output = module(*input, **kwargs)
  File "/home/user/.conda/envs/main/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/user/.conda/envs/main/lib/python3.8/site-packages/transformers/modeling_gpt2.py", line 752, in forward
    transformer_outputs = self.transformer(
  File "/home/user/.conda/envs/main/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/user/.conda/envs/main/lib/python3.8/site-packages/transformers/modeling_gpt2.py", line 587, in forward
    attention_mask = attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
StopIteration

```

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

I simply expect to get an output and no error :)

<!-- A clear and concise description of what you would expect to happen. -->



Discussion:

Resolving patch:
diff --git a/src/transformers/modeling_gpt2.py b/src/transformers/modeling_gpt2.py
--- a/src/transformers/modeling_gpt2.py
+++ b/src/transformers/modeling_gpt2.py
@@ -584,7 +584,7 @@ def forward(
             # positions we want to attend and -10000.0 for masked positions.
             # Since we are adding it to the raw scores before the softmax, this is
             # effectively the same as removing these entirely.
-            attention_mask = attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
+            attention_mask = attention_mask.to(dtype=self.dtype)  # fp16 compatibility
             attention_mask = (1.0 - attention_mask) * -10000.0
 
         # If a 2D ou 3D attention mask is provided for the cross-attention


Test patch: