

TITLE
TPU pod training with BERT 

COMMENTS
9

REACTIONS
+1: 0
-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.3.1
- Platform: TPU
- Python version: 
- PyTorch version (GPU?): 1.6 TPU v3-32
- Tensorflow version (GPU?):
- Using GPU in script?:
- Using distributed or parallel set-up in script?: Yes

### 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
 -->

## Information

Model I am using (Bert):

The problem arises when using:
* [x] the official example scripts: (give details below)
* [] my own modified scripts: (give details below)

https://huggingface.co/blog/how-to-train

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


After starting training on TPU Pod using Pytorch XLA, I get following long message

```## ❓ Questions and Help
2020-10-14 14:14:33 10.164.0.60 [2] Traceback (most recent call last):
2020-10-14 14:14:33 10.164.0.60 [2]   File "/home/developer/XainBERT/train.py", line 111, in <module>
2020-10-14 14:14:33 10.164.0.60 [2]     prediction_loss_only=True,
2020-10-14 14:14:33 10.164.0.60 [2]   File "/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/transformers/trainer.py", line 248, in __init__
2020-10-14 14:14:33 10.164.0.60 [2]     self.model = model.to(args.device) if model is not None else None
2020-10-14 14:14:33 10.164.0.60 [2]   File "/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/transformers/file_utils.py", line 944, in wrapper
2020-10-14 14:14:33 10.164.0.60 [2]     return func(*args, **kwargs)
2020-10-14 14:14:33 10.164.0.60 [2]   File "/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/transformers/training_args.py", line 408, in device
2020-10-14 14:14:33 10.164.0.60 [2]     return self._setup_devices[0]
2020-10-14 14:14:33 10.164.0.60 [2]   File "/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/transformers/file_utils.py", line 934, in __get__
2020-10-14 14:14:33 10.164.0.60 [2]     cached = self.fget(obj)
2020-10-14 14:14:33 10.164.0.60 [2]   File "/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/transformers/file_utils.py", line 944, in wrapper
2020-10-14 14:14:33 10.164.0.60 [2]     return func(*args, **kwargs)
2020-10-14 14:14:33 10.164.0.60 [2]   File "/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/transformers/training_args.py", line 379, in _setup_devices
2020-10-14 14:14:33 10.164.0.60 [2]     device = xm.xla_device()
2020-10-14 14:14:33 10.164.0.60 [2]   File "/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 167, in xla_device
2020-10-14 14:14:33 10.164.0.60 [2]     devkind=[devkind] if devkind is not None else None)
2020-10-14 14:14:33 10.164.0.60 [2]   File "/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 72, in get_xla_supported_devices
2020-10-14 14:14:33 10.164.0.60 [2]     xla_devices = _DEVICES.value
2020-10-14 14:14:33 10.164.0.60 [2]   File "/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/torch_xla/utils/utils.py", line 30, in value
2020-10-14 14:14:33 10.164.0.60 [2]     self._value = self._gen_fn()
2020-10-14 14:14:33 10.164.0.60 [2]   File "/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 18, in <lambda>
2020-10-14 14:14:33 10.164.0.60 [2]     _DEVICES = xu.LazyProperty(lambda: torch_xla._XLAC._xla_get_devices())
2020-10-14 14:14:33 10.164.0.60 [2] RuntimeError: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:244 : Check failed: default_device_target != options_.global_device_map.end() 
2020-10-14 14:14:33 10.164.0.60 [2] *** Begin stack trace ***
2020-10-14 14:14:33 10.164.0.60 [2]     tensorflow::CurrentStackTrace()
2020-10-14 14:14:33 10.164.0.60 [2]     xla::XrtComputationClient::XrtComputationClient(xla::XrtComputationClient::Options, std::unique_ptr<tensorflow::tpu::TopologyProto, std::default_delete<tensorflow::tpu::TopologyProto> >)
2020-10-14 14:14:33 10.164.0.60 [2]     xla::ComputationClient::Create()
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     xla::ComputationClient::Get()
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     _PyCFunction_FastCallDict
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     _PyEval_EvalFrameDefault
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     _PyEval_EvalFrameDefault
2020-10-14 14:14:33 10.164.0.60 [2]     PyEval_EvalCodeEx
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     PyObject_Call
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     _PyObject_GenericGetAttrWithDict
2020-10-14 14:14:33 10.164.0.60 [2]     _PyEval_EvalFrameDefault
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     _PyEval_EvalFrameDefault
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     _PyEval_EvalFrameDefault
2020-10-14 14:14:33 10.164.0.60 [2]     PyEval_EvalCodeEx
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     PyObject_Call
2020-10-14 14:14:33 10.164.0.60 [2]     _PyEval_EvalFrameDefault
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     _PyEval_EvalFrameDefault
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     _PyFunction_FastCallDict
2020-10-14 14:14:33 10.164.0.60 [2]     _PyObject_FastCallDict
2020-10-14 14:14:33 10.164.0.60 [2]     PyObject_CallFunctionObjArgs
2020-10-14 14:14:33 10.164.0.60 [2]     _PyObject_GenericGetAttrWithDict
2020-10-14 14:14:33 10.164.0.60 [2]     _PyEval_EvalFrameDefault
2020-10-14 14:14:33 10.164.0.60 [2]     PyEval_EvalCodeEx
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     PyObject_Call
2020-10-14 14:14:33 10.164.0.60 [2]     _PyEval_EvalFrameDefault
2020-10-14 14:14:33 10.164.0.60 [2]     PyEval_EvalCodeEx
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     PyObject_Call
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     _PyObject_GenericGetAttrWithDict
2020-10-14 14:14:33 10.164.0.60 [2]     _PyEval_EvalFrameDefault
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     _PyFunction_FastCallDict
2020-10-14 14:14:33 10.164.0.60 [2]     _PyObject_FastCallDict
2020-10-14 14:14:33 10.164.0.60 [2]     _PyObject_Call_Prepend
2020-10-14 14:14:33 10.164.0.60 [2]     PyObject_Call
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     _PyObject_FastCallDict
2020-10-14 14:14:33 10.164.0.60 [2]     _PyObject_FastCallKeywords
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     _PyEval_EvalFrameDefault
2020-10-14 14:14:33 10.164.0.60 [2]     PyEval_EvalCodeEx
2020-10-14 14:14:33 10.164.0.60 [2]     PyEval_EvalCode
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2]     PyRun_FileExFlags
2020-10-14 14:14:33 10.164.0.60 [2]     PyRun_SimpleFileExFlags
2020-10-14 14:14:33 10.164.0.60 [2]     Py_Main
2020-10-14 14:14:33 10.164.0.60 [2]     main
2020-10-14 14:14:33 10.164.0.60 [2]     __libc_start_main
2020-10-14 14:14:33 10.164.0.60 [2] 
2020-10-14 14:14:33 10.164.0.60 [2] *** End stack trace ***
020-10-14 14:27:11 10.164.0.61 [0] /anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/transformers/trainer.py:267: FutureWarning: Passing `prediction_loss_only` as a keyword argument is deprecated and won't 
be possible in a future version. Use `args.prediction_loss_only` instead.
2020-10-14 14:27:11 10.164.0.61 [0]   FutureWarning,
2020-10-14 14:27:11 10.164.0.61 [0] Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future version. Using `--per_device_train_batch_size` is preferred.
2020-10-14 14:27:11 10.164.0.61 [0] Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future version. Using `--per_device_train_batch_size` is preferred.
2020-10-14 14:27:11 10.164.0.61 [0] Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future version. Using `--per_device_train_batch_size` is preferred.
Epoch:   0%|          | 0/1 [00:00<?, ?it/s]
2020-10-14 14:28:09 10.164.0.61 [0] 
2020-10-14 14:28:20 10.164.0.61 [0] 
2020-10-14 14:28:32 10.164.0.61 [0] 
2020-10-14 14:28:33 10.164.0.61 [0] 
2020-10-14 14:28:38 10.164.0.61 [0] 
2020-10-14 14:28:38 10.164.0.61 [0] 
2020-10-14 14:28:38 10.164.0.61 [0] 
2020-10-14 14:28:40 10.164.0.61 [0] 
2020-10-14 14:28:42 10.164.0.61 [0] 
2020-10-14 14:28:54 10.164.0.61 [0] 
2020-10-14 14:28:54 10.164.0.61 [0] 
2020-10-14 14:28:56 10.164.0.61 [0] 
2020-10-14 14:28:57 10.164.0.61 [0] 
2020-10-14 14:28:58 10.164.0.61 [0] 
2020-10-14 14:28:58 10.164.0.61 [0] 
2020-10-14 14:28:59 10.164.0.61 [0] 
2020-10-14 14:29:01 10.164.0.61 [0] 
2020-10-14 14:29:01 10.164.0.61 [0] 
2020-10-14 14:29:01 10.164.0.61 [0] 
2020-10-14 14:29:02 10.164.0.61 [0] 
2020-10-14 14:29:02 10.164.0.61 [0] 
```

I think the problem lies here 

```2020-10-14 14:14:33 10.164.0.60 [2] RuntimeError: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:244 : Check failed: default_device_target != options_.global_device_map.end() ```

But I do not understand what is wrong. 

I basically use the script from here https://huggingface.co/blog/how-to-train with some modifications, own data but Bert vocabulary. I have used the same script for training on TPUv3-8 as well which works fine.The transformer library version is 3.3.1 and I am using BERT tokeniser and BERT config. The distributed training on TPU pod worked with ImageNet example mentioned in https://cloud.google.com/tpu/docs/tutorials/pytorch-pod#configure-gcloud. 


## To reproduce
I can't really provide the whole script. But the training command is as follow:

python -m torch_xla.distributed.xla_dist  --tpu=$TPU_NAME --conda-env=torch-xla-1.6 -- python /home/developer/XainBERT/xla_spawn.py --num_cores 8 /home/developer/XainBERT/train.py 

Steps to reproduce the behavior:

1.
2.
3.

<!-- 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

<!-- A clear and concise description of what you would expect to happen. -->
