Error with running run_language_modeling.py on GCP TPU
Author: aldazero
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: 3.3.1
- Platform: GCP
- Python version: 3.6.10 
- PyTorch version (GPU?): 1.6.0
- Tensorflow version (GPU?): - 
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: Yes/NO

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.
 
 albert, bert, GPT2, XLM: @LysandreJik 
 tokenizers: @mfuntowicz
 Trainer: @sgugger
 Speed and Memory Benchmarks: @patrickvonplaten
 Model Cards: @julien-c
 Translation: @sshleifer
 Summarization: @sshleifer
 TextGeneration: @TevenLeScao 
 examples/distillation: @VictorSanh
 nlp datasets: [different repo](https://github.com/huggingface/nlp)
 rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
 Text Generation: @TevenLeScao
 blenderbot: @mariamabarham
 Bart: @sshleifer
 Marian: @sshleifer
 T5: @patrickvonplaten
 Longformer/Reformer: @patrickvonplaten
 TransfoXL/XLNet: @TevenLeScao 
 examples/seq2seq: @sshleifer
 examples/bert-loses-patience: @JetRunner
 tensorflow: @jplu
 examples/token-classification: @stefan-it
 documentation: @sgugger
 -->
@sgugger @LysandreJik
## Information

Model I am using (Bert, XLNet ...):

The problem arises when using:
* [ run_language_modeling.py ] 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)
* [fine tuning BERT on Wikitext103 ] my own task or dataset: (give details below)

## To reproduce

I am trying to fine-tune Bert on the Wikitext-103 dataset by running the example code run_language_modeling.py on Google cloud TPU v3-8 using xla_spawn.py launcher. I tried both num_cores =1 and num_cores > 1, but neither worked properly. 

Steps to reproduce the behavior:

1. python xla_spawn.py --num_cores 8 \
 run_language_modeling.py \
--model_name_or_path=bert-base-uncased \
--do_train \
--train_data_file\
--do_eval \
--eval_data_file\
--mlm\
--per_device_train_batch_size=4


The full output is long and you can find it [here](https://gofile.io/d/tjovN1). Here is the begining of error message:

Iteration:   0%|          | 1/28026 [00:00<3:41:44,  2.11it/s][A2020-10-09 16:05:45.623409: W
  2449 tensorflow/core/distributed_runtime/rpc/grpc_remote_master.cc:160] RPC failed with status = "Unavailable: Socket closed" and grpc_error_string = "{"created":"@1602259545.623281190","description":"Error received from peer ipv4:10.48.142.250:8470","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Socket closed","grpc_status":14}", maybe retrying the RPC
2020-10-09 16:06:09.642115: E    2449 tensorflow/compiler/xla/xla_client/xla_util.cc:76] >>> Dumping Computation 0
2020-10-09 16:06:09.642213: E    2449 tensorflow/compiler/xla/xla_client/xla_util.cc:76] HloModule SyncTensorsGraph.23679, input_output_alias={ {0}: (39, {}, may-alias), {1}: (37, {}, may-alias),
......

<!-- 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
Running without error!
<!-- A clear and concise description of what you would expect to happen. -->
Thanks for any help. 

Comments:
This doesn't seem to be an issue with `transformers` but with your TPU and its communication with your VM. You would probably have more help if you asked over at https://github.com/pytorch/xla

@LysandreJik Thanks for your comment. The problem rooted in the TPU software version. Setting it to `PyTorch-1.6` resolved the issue.