Possible fix to make AMP work with DDP in the trainer
closes https://github.com/huggingface/transformers/issues/4657

Using multiple GPUs in PyTorch (with DistrubutedDataParallel) speeds up performance drastically. To get even more speed out of it, the example scripts often - if not always - allow the use of apex for automatic mixed precision. This is great. However, some issues can arise, particularly the infamous "illegal memory access" error which seems to have an easy, one-line solution. 

Currently, we assume that by using things like `.to(args.device)` we solve any and all issues of where our data or model should go. However, the author of AMP @mcarilli seems to [suggest](https://github.com/NVIDIA/apex/issues/319#issuecomment-503372924) that it is recommended to always set the current process' default device, too, to ensure no further issues. This suggestion also [helped](https://github.com/huggingface/transformers/issues/4657#issuecomment-637703146) the aforementioned issue. Therefore it seems a good idea to also implement this. In fact, some examples such as Hans already do this.

https://github.com/huggingface/transformers/blob/b231a413f5d58592bb4d98304c3d3b668c5d4a42/examples/adversarial/test_hans.py#L518

To avoid DRY issues, I would suspect that the trainer_args file is the best place to do this _only once_ but other suggestions are welcome (this is different from what I suggested in the linked issue, though I think `trainer_args` is the better place). I am not sure which examples do not use trainer_args, but those would need to be checked and updated as well. If anyone can give a quick rundown of which examples do NOT use the new trainer, I can have a look quickly. Otherwise I'll have to go over the examples another time.

**Side note**: I am not sure how and if this works with TPUs so to be sure that this only involves CUDA devices, I first check whether the device is a CUDA device.
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/4728?src=pr&el=h1) Report
> Merging [#4728](https://codecov.io/gh/huggingface/transformers/pull/4728?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/b231a413f5d58592bb4d98304c3d3b668c5d4a42&el=desc) will **decrease** coverage by `1.65%`.
> The diff coverage is `33.33%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/4728/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/4728?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #4728      +/-   ##
==========================================
- Coverage   77.27%   75.62%   -1.66%     
==========================================
  Files         128      128              
  Lines       20980    20982       +2     
==========================================
- Hits        16213    15868     -345     
- Misses       4767     5114     +347     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/4728?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/training\_args.py](https://codecov.io/gh/huggingface/transformers/pull/4728/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmluZ19hcmdzLnB5) | `75.82% <33.33%> (-0.59%)` | :arrow_down: |
| [src/transformers/modeling\_longformer.py](https://codecov.io/gh/huggingface/transformers/pull/4728/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19sb25nZm9ybWVyLnB5) | `17.51% <0.00%> (-75.92%)` | :arrow_down: |
| [src/transformers/configuration\_longformer.py](https://codecov.io/gh/huggingface/transformers/pull/4728/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX2xvbmdmb3JtZXIucHk=) | `75.00% <0.00%> (-25.00%)` | :arrow_down: |
| [src/transformers/modeling\_t5.py](https://codecov.io/gh/huggingface/transformers/pull/4728/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190NS5weQ==) | `77.13% <0.00%> (-6.37%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/4728/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl91dGlscy5weQ==) | `88.63% <0.00%> (-0.33%)` | :arrow_down: |
| [src/transformers/modeling\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/4728/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19iZXJ0LnB5) | `86.75% <0.00%> (-0.19%)` | :arrow_down: |
| [src/transformers/modeling\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/4728/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ191dGlscy5weQ==) | `90.15% <0.00%> (-0.13%)` | :arrow_down: |
| [src/transformers/file\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/4728/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9maWxlX3V0aWxzLnB5) | `74.27% <0.00%> (+0.41%)` | :arrow_up: |
| [src/transformers/modeling\_openai.py](https://codecov.io/gh/huggingface/transformers/pull/4728/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19vcGVuYWkucHk=) | `81.72% <0.00%> (+1.37%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/4728?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/4728?src=pr&el=footer). Last update [b231a41...8af2fd8](https://codecov.io/gh/huggingface/transformers/pull/4728?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).
<|||||>LGTM, thanks for the detailed write-up and research @BramVanroy 

I think this was here before but I removed it when refactoring, assuming that it was redundant – looks like it wasn't really:)

As for which examples do NOT use the new trainer, you should refer to the table at https://github.com/huggingface/transformers/tree/master/examples – we expect all of them to use Trainer/TFTrainer eventually.

Thank you!