xla_spawn.py crashes when training on TPU V3-32
## 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: 4.0.1
- Platform: Google Cloud debian-9-torch-xla-v20201215
- Python version: 3.7
- PyTorch version (GPU?): 1.7
- Tensorflow version (GPU?):
- Using GPU in script?: NO; using TPUS
- Using distributed or parallel set-up in script?:

### Who can help

@sgugger @patrickvonplaten


## Information

Model I am using (Bert, XLNet ...):

I am using an Albert base. 

The problem arises when using:
* [x] the official example scripts: (give details below)
Using the examples/xla_spawn.py together with run_mlm.py it crashes when we try to use it with v3-32. We're supposed to set num cores to either 1 or 8 but in our case we have 32 cores and it raises an error. We've also tried to let that variable to 1 or 8, but in both cases it raises errors: 

```
.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www
.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www
.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www
.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www
.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www
.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www
.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www
.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www
.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
Exception in device=TPU:0: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:1229 : Check failed: session.Run({tensorflow::Output(result, 0)}, &outputs) == ::tensorflow::Status::OK() (Internal: From /job:tpu_worker/replica:0/t
ask:0:
2 root error(s) found.
  (0) Internal: Invalid system configuration: 2x2 host topology with 0 missing hosts, but 1 hosts in total.
         [[{{node configure_distributed_tpu/_0}}]]
         [[ConfigureDistributedTPU_G3]]
  (1) Internal: Invalid system configuration: 2x2 host topology with 0 missing hosts, but 1 hosts in total.
         [[{{node configure_distributed_tpu/_0}}]]
0 successful operations.
0 derived errors ignored. vs. OK)
*** Begin stack trace ***
        tensorflow::CurrentStackTrace()
        xla::XrtComputationClient::InitializeAndFetchTopology(std::string const&, int, std::string const&, tensorflow::ConfigProto const&)
        xla::XrtComputationClient::InitializeDevices(std::unique_ptr<tensorflow::tpu::TopologyProto, std::default_delete<tensorflow::tpu::TopologyProto> >)
        xla::XrtComputationClient::XrtComputationClient(xla::XrtComputationClient::Options, std::unique_ptr<tensorflow::tpu::TopologyProto, std::default_delete<tensorflow::tpu::TopologyProto> >)
        xla::ComputationClient::Create()
        xla::ComputationClient::Get()
        _PyCFunction_FastCallDict
        _PyEval_EvalFrameDefault

        _PyEval_EvalFrameDefault
        PyEval_EvalCodeEx

        PyObject_Call

        _PyObject_GenericGetAttrWithDict
        _PyEval_EvalFrameDefault



        _PyEval_EvalFrameDefault


        _PyEval_EvalFrameDefault


        _PyEval_EvalFrameDefault


        _PyEval_EvalFrameDefault
        PyEval_EvalCodeEx

        PyObject_Call
        _PyEval_EvalFrameDefault
        PyEval_EvalCodeEx

        PyObject_Call
        _PyEval_EvalFrameDefault


        _PyEval_EvalFrameDefault


        _PyEval_EvalFrameDefault


        _PyEval_EvalFrameDefault



        _PyEval_EvalFrameDefault
        PyEval_EvalCodeEx
        PyEval_EvalCode

        PyRun_StringFlags
        PyRun_SimpleStringFlags
        Py_Main
        main
        __libc_start_main

*** End stack trace ***

Traceback (most recent call last):
  File "transformers/examples/xla_spawn.py", line 85, in <module>
    main()
  File "transformers/examples/xla_spawn.py", line 81, in main
    xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 395, in spawn
    start_method=start_method)
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 157, in start_processes
    while not context.join():
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 112, in join
    (error_index, exitcode)
Exception: process 0 terminated with exit code 17
```

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

## To reproduce

Steps to reproduce the behavior:

1. Initialize a TPU V3-32 and when running xla_spawn.py, set the number of cores to either 32, 8 or 1. In all three cases it raises an error. 


## Expected behavior

There should not be any problem in setting the number of cores to the number of TPU cores we actually have. It really does not make sense to be able to train only with either 1 core or 8 cores... 
Hmmm this may be because a TPU v3-32 regroups several TPU chips, as the error here seems to imply:
```
Internal: Invalid system configuration: 2x2 host topology with 0 missing hosts, but 1 hosts in total.
```

@sgugger can you confirm this is the source of the issue? Do you know the status of the Trainer/xla_spawn on TPU pods?<|||||>@LysandreJik If that's the source of the issue, what would be the procedure to solve it?<|||||>From the stack trace it doesn't look like it even gets to the training script, so I think there might be something wrong in your distributed TPU setup. Are you able to run another script (coming from official torch XLA for instance?) on this setup?<|||||>What do you mean by setup?? 

```{bash}
XRT_TPU_CONFIG="tpu_worker;0;10.157.150.13:8470"
```
This is the configuration parameter I set before calling the training script, which starts like this:

```{bash}
python transformers/examples/xla_spawn.py --num_cores 8 \
    transformers/examples/language-modeling/run_mlm.py \
```

I also tried setting the --num_cores to 1 (it only accepts 1 or 8) in the second code snippet. 
With v3-8 this setup works correctly, I don't know if you mean this by setup...<|||||>I mean the setup you are using. It can't be the same setup for one TPU (8cores) and a TPU pod (for the 32 cores). The second requires to launch different machines. That's why I was asking if you could run another example from someone else on your TPU v3-32.

Also, the launcher scrip `xla_spawn` only works for one TPU, not a TPU pod as fa as I know, so you will need to launch the script in a different way.<|||||>@sgugger does  ```xla_spawn ``` not support TPU pod? As many issues in this repo are related to TPU pod, so I have thought ```xla_spawn``` also support it. Do you know any examples of using TPU pod?<|||||>Sorry I didn't write it well. I meant the launcher script `xla_spawn` has only been tested for one TPU, not a TPU pod as far as I know. So you may need to launch the script in a different way.

I am not aware of anyone launching any of the example scripts on a TPU pod successfully, so I don't know if they work or not. <|||||>I see.
That information should be added to the document if it does not.

By the way, TFTrainer support TPU pod? I think it does, but I have not tested yet.<|||||>Same thing, it has not been tested. We don't have resources setup to test for more than a single TPU (so 8 cores).<|||||>I understand. 
Thank you for answering.
<|||||>Okay, so I'd need to change the setup for a TPU pod then... I don't understand why all this complication to go from 8 cores to 32 cores actually, I know that's on Google's side, but I don't think it makes sense to complicate things so much to be able to train on a 32 cores TPU. As I understood, not only the setup must be changed, but also the script to launch the xla, right? I mean, the xla_spawn.py from Transformers is thought for 1 TPU, and it may crash on multiple TPU nodes?<|||||>This issue has been automatically marked as stale and been closed because it has not had recent activity. Thank you for your contributions.

If you think this still needs to be addressed please comment on this thread.<|||||>I am having the same issue (xla_spawn works fine with 8 core TPU but fails with TPU pods v3-32 and above). Is there a way to utilize TPU pods with the transformers library? <|||||>Hi @sgugger, I am running xla_spawn, can I know is this correct way to run the hugginface examples on TPU pods like v3-32? Thanks
```
TPU_NAME=tpu-v3-32
python3 -m  torch_xla.distributed.xla_dist \
  --tpu=${TPU_NAME} --restart-tpuvm-pod-server  -- \
 python3 /transformers/examples/pytorch/xla_spawn.py --num_cores 8 /pytorch/text-classifications/run_glue.py \
  --model_name_or_path bert-base-cased \
  --dataset_name SetFit/mrpc  \
  --do_train \
  --do_eval \
  --max_seq_length 128 \
  --per_device_train_batch_size 1 \
  --per_device_eval_batch_size 32 \
  --learning_rate 2e-5 \
  --num_train_epochs 3 \
  --run_name mnli_v3-32_bs-64_lr-2e-5-bert \
  --output_dir /tmp/mrpc-bert/ \
  --overwrite_output_dir
``` 
Where I use `python3 xla_dist` to wrap the `python3 xla_spawn.py --num_cores 8`<|||||>The examples can't be launched directly on TPU pods. cc @muellerzr who has worked on them with accelerate and can share how to run an example on a TPU pod.<|||||>Hi @sgugger Thanks:). I do successfully run above command on TPU pods (V3-32 and V4-64), see the [wandb results](https://wandb.ai/jianguozhang/huggingface/reports/mrpc-for-text-classification--VmlldzozMzU2NTE1?accessToken=hce0jseir4d3x32cqbocha2936xes1r1hbeupgijjy6l2lhujtd0y2577xzcdn2c). But i am not sure it whether it is correct way to use the commands as the training loss is much higher than that on GPUs, and V4-64 shows lower running speed than v3-32. 
Hi @muellerzr, can you show an example that how to run huggingface torch_xla examples on TPU pods? Thanks:) 