Problem with pretraining GPT-2 on TPU with Pytorch/XLA
Author: hakanbogan
Pull request: no

### **Environment info**
transformers version: 4.0.1
Platform: Ubuntu 18.04.4 LTS
Python version: 3.6.9
PyTorch version: 1.7.1
Torch-XLA version: 1.7
Tensorflow version: 2.3.1

### **Information**
Model I am intended to pretrain (Bert, XLNet ...): GPT2
The problem arises when using: 
- my own modified scripts: (give details below)

The tasks I am working on is:
- pretraining GPT-2 on Google TPU

### **To reproduce**
Steps to reproduce the behavior:

1. According to instructions [here](https://github.com/pytorch/xla/blob/master/README.md#-consume-prebuilt-compute-vm-images), creating instance of compute engine on Google Cloud with following spec:

- OS: Deep Learning on Linux
- Version: Debian GNU/Linux 9 Stretch + Pytorch/XLA

2. Running the code of:
`export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"`

3. Running following script.
```python
#!/usr/bin/env bash
pipenv run python3 xla_spawn.py --num_cores 8 \
    run_clm.py \
    --num_train_epochs 5 \
    --output_dir saved_model/ \
    --overwrite_output_dir \
    --logging_dir logs \
    --logging_steps 50 \
    --save_total_limit 2 \
    --save_steps 2000 \
    --model_type gpt2 \
    --config_name tokenized_data/ \
    --tokenizer_name tokenized_data/ \
    --block_size 1024 \
    --train_file dataset/dataset.txt \
    --per_device_train_batch_size=64 \
    --do_train 
```

Here is the exception occured:
```
Exception in device=TPU:3: tensorflow/compiler/xla/xla_client/mesh_service.cc:316 : Check failed: impl_->channel->WaitForConnected( std::chrono::system_clock::now() + std::chrono::seconds(connect_wait_seconds)) 
*** Begin stack trace ***
	tensorflow::CurrentStackTrace()
	xla::service::MeshClient::MeshClient(std::string const&)
	xla::service::MeshClient::Get()
	xla::ComputationClient::Create()
	
	
	xla::ComputationClient::Get()
	
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	
	PyObject_GetAttr
	_PyEval_EvalFrameDefault
	
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	PyObject_Call
	_PyEval_EvalFrameDefault
	
	
	PyObject_Call
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	
	_PyEval_EvalFrameDefault
	
	PyRun_StringFlags
	PyRun_SimpleStringFlags
	Py_Main
	main
	__libc_start_main
	_start
*** End stack trace ***
Failed to connect to client mesh master: workstation:41295
Traceback (most recent call last):
  File "/home/ws/.local/share/virtualenvs/gpt-2_pretrain_huggingface-aTVuIzXL/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 330, in _mp_start_fn
    _start_fn(index, pf_cfg, fn, args)
  File "/home/ws/.local/share/virtualenvs/gpt-2_pretrain_huggingface-aTVuIzXL/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 323, in _start_fn
    _setup_replication()
  File "/home/ws/.local/share/virtualenvs/gpt-2_pretrain_huggingface-aTVuIzXL/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 315, in _setup_replication
    device = xm.xla_device()
  File "/home/ws/.local/share/virtualenvs/gpt-2_pretrain_huggingface-aTVuIzXL/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 231, in xla_device
    devkind=devkind if devkind is not None else None)
  File "/home/ws/.local/share/virtualenvs/gpt-2_pretrain_huggingface-aTVuIzXL/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 136, in get_xla_supported_devices
    xla_devices = _DEVICES.value
  File "/home/ws/.local/share/virtualenvs/gpt-2_pretrain_huggingface-aTVuIzXL/lib/python3.6/site-packages/torch_xla/utils/utils.py", line 32, in value
    self._value = self._gen_fn()
  File "/home/ws/.local/share/virtualenvs/gpt-2_pretrain_huggingface-aTVuIzXL/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 18, in <lambda>
    _DEVICES = xu.LazyProperty(lambda: torch_xla._XLAC._xla_get_devices())
RuntimeError: tensorflow/compiler/xla/xla_client/mesh_service.cc:316 : Check failed: impl_->channel->WaitForConnected( std::chrono::system_clock::now() + std::chrono::seconds(connect_wait_seconds)) 
*** Begin stack trace ***
	tensorflow::CurrentStackTrace()
	xla::service::MeshClient::MeshClient(std::string const&)
	xla::service::MeshClient::Get()
	xla::ComputationClient::Create()
	
	
	xla::ComputationClient::Get()
	
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	
	PyObject_GetAttr
	_PyEval_EvalFrameDefault
	
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	PyObject_Call
	_PyEval_EvalFrameDefault
	
	
	PyObject_Call
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	_PyEval_EvalFrameDefault
	
	
	
	_PyEval_EvalFrameDefault
	
	PyRun_StringFlags
	PyRun_SimpleStringFlags
	Py_Main
	main
	__libc_start_main
	_start
*** End stack trace ***
Failed to connect to client mesh master: workstation:41295
Exception in device=TPU:2: tensorflow/compiler/xla/xla_client/mesh_service.cc:316 : Check failed: impl_->channel->WaitForConnected( std::chrono::system_clock::now() + std::chrono::seconds(connect_wait_seconds)) 
```

Comments:
I realized that I didn't follow [instructions](https://github.com/pytorch/xla/blob/master/README.md) properly.

@redrussianarmy  What`s the problem? I have the same issue but can not find a solution.