

TITLE
[GPT2] Loss NaN after some time with FP16

COMMENTS
5

REACTIONS
+1: 1
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
## 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: 3.4.0
- Platform: Linux-4.4.0-176-generic-x86_64-with-glibc2.17
- Python version: 3.8.6
- PyTorch version (GPU?): 1.7.0+cu101 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes, HF datasets


### Who can help
 
@LysandreJik 
@sgugger

## Information

Model I am using (Bert, XLNet ...): GPT2

The problem arises when using:
* [ ] the official example scripts: 
* [x] my own modified scripts: examples/language_modeling/run_language_modeling.py

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset

## To reproduce

Steps to reproduce the behavior:

1. Run with ```--fp16 --n_ctx 2048```
2. 
```
  warnings.warn('Was asked to gather along dimension 0, but all '
[W python_anomaly_mode.cpp:104] Warning: Error detected in SoftmaxBackward. Traceback of forward call that caused the error:
  File "/usr/lib/python3.8/threading.py", line 890, in _bootstrap
    self._bootstrap_inner()
  File "/usr/lib/python3.8/threading.py", line 932, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.8/threading.py", line 870, in run
    self._target(*self._args, **self._kwargs)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/transformers/modeling_gpt2.py", line 765, in forward
    transformer_outputs = self.transformer(
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/transformers/modeling_gpt2.py", line 645, in forward
    outputs = block(
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/transformers/modeling_gpt2.py", line 285, in forward
    attn_outputs = self.attn(
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/transformers/modeling_gpt2.py", line 235, in forward
    attn_outputs = self._attn(query, key, value, attention_mask, head_mask, output_attentions)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/transformers/modeling_gpt2.py", line 176, in _attn
    w = nn.Softmax(dim=-1)(w)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/nn/modules/activation.py", line 1198, in forward
    return F.softmax(input, self.dim, _stacklevel=5)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/nn/functional.py", line 1512, in softmax
    ret = input.softmax(dim)
 (function _print_stack)
Traceback (most recent call last):
  File "examples/language-modeling/run_language_modeling.py", line 349, in <module>
    main()
  File "examples/language-modeling/run_language_modeling.py", line 313, in main
    trainer.train(model_path=model_path)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/transformers/trainer.py", line 756, in train
    tr_loss += self.training_step(model, inputs)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/transformers/trainer.py", line 1065, in training_step
    self.scaler.scale(loss).backward()
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/tensor.py", line 221, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/autograd/__init__.py", line 130, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Function 'SoftmaxBackward' returned nan values in its 0th output.
  0%|                                              | 0/19506024 [00:25<?, ?it/s]
```

## Expected behavior

Not print Nan
