Run Language Modeling on 8 TPU cores doesn't seem to terminate
Author: jcblaisecruz02
Pull request: no

# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...): DistilGPT2 & GPT2

Language I am using the model on (English, Chinese ...): English

The problem arises when using:
* [x] the official example scripts: (give details below)
* [ ] 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)

I'm trying to test ```run_language_modeling.py``` on DistilGPT2 using all 8 TPU cores. Running on 1 core executes fine, but when I attempt to run on all 8 cores, it finishes finetuning then gets stuck on "Training completed. Do not forget to share your model on huggingface.co/models =)" and doesn't terminate.

When I check the output directory, I only see two files: config.json and pytorch_model.bin. There should be seven files in the output directory.

I'm running this on a Colab TPU Notebook.

## To reproduce

Steps to reproduce the behavior:
```
VERSION = "nightly"  #@param ["1.5" , "20200325", "nightly"]
!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
!python pytorch-xla-env-setup.py --version $VERSION

!wget https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-v1.zip
!unzip wikitext-2-v1.zip && rm wikitext-2-v1.zip
!pip install transformers
!git clone https://github.com/huggingface/transformers.git

!python transformers/examples/xla_spawn.py --num_cores 8 \
	transformers/examples/language-modeling/run_language_modeling.py \
	--output_dir=output \
        --model_type=distilgpt2 \
        --model_name_or_path=distilgpt2 \
        --train_data_file=wikitext-2/wiki.train.tokens \
        --do_train \
        --overwrite_output_dir 

```

<!-- 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

Script terminates and 7 files are in the output folder:
* config.json  
* pytorch_model.bin	     
* tokenizer_config.json  
* vocab.json
* merges.txt   
* special_tokens_map.json  
* training_args.bin

<!-- A clear and concise description of what you would expect to happen. -->

## 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: 2.9.1
- Platform: Linux-4.19.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.6.0a0+83df3be (False)
- Tensorflow version (GPU?): 2.2.0 (False)
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: yes


Comments:
@jcblaisecruz02 
Yes, there's a bug in version 2.91 which hangs the trainer. It's been fixed in master branch. Install from master branch for TPU training.

See this pull request #4339

Is it possible to run `run_language_modeling.py` on more than 8 cores when using pytorch and `xls_spawn`?
And what about tensorflow?