TPU not initialized when running official `run_mlm_flax.py` example.
Author: patrickvonplaten
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: 4.9.0.dev0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29
- Python version: 3.8.5
- PyTorch version (GPU?): 1.9.0+cu102 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (tpu)
- Jax version: 0.2.16
- JaxLib version: 0.1.68
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No


### Who can help

@avital @marcvanzee 

## Information

I am setting up a new TPU VM according to the [Cloud TPU VM JAX quickstart](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm) and the following the installation steps as described here: https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects#how-to-install-relevant-libraries to install `flax`, `jax` `transformers`, and `datasets`.

Then, when running a simple example using the [`run_mlm_flax.py`](https://github.com/huggingface/transformers/blob/master/examples/flax/language-modeling/run_mlm_flax.py) script, I'm encounting an error/ warning:

```
INFO:absl:Starting the local TPU driver.
INFO:absl:Unable to initialize backend 'tpu_driver': Not found: Unable to find driver in registry given worker: local://
INFO:absl:Unable to initialize backend 'gpu': Not found: Could not find registered platform with name: "cuda". Available platform names are: TPU Interpreter Host
```

=> I am now unsure whether the code actually runs on TPU or instead on CPU.

## To reproduce

The problem can be easily reproduced by: 

1. sshing into a TPU, *e.g.* `patrick-test` (Flax, JAX, & Transformers should already be installed)

If one goes into `patrick-test` the libraries are already installed - on an "newly" created TPU VM, one can follow [these](https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects#how-to-install-relevant-libraries) steps to install the relevant libraries.

2. Going to home folder

```
cd ~/
```

3. creating a new dir:

```
mkdir test && cd test
```

4. cloning a dummy repo into it

```
GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/patrickvonplaten/norwegian-roberta-als
```

4. Linking the `run_mlm_flax.py` script

```
ln -s $(realpath ~/transformers/examples/flax/language-modeling/run_mlm_flax.py) ./
```

5. Running the following command (which should show the above warning/error again):

```
./run_mlm_flax.py \
    --output_dir="norwegian-roberta-als" \
    --model_type="roberta" \
    --config_name="norwegian-roberta-als" \
    --tokenizer_name="norwegian-roberta-als" \
    --dataset_name="oscar" \
    --dataset_config_name="unshuffled_deduplicated_als" \
    --max_seq_length="128" \
    --per_device_train_batch_size="8" \
    --per_device_eval_batch_size="8" \
    --learning_rate="3e-4" \
    --overwrite_output_dir \
    --num_train_epochs="3"
```

=>

You should see a console print that says:

```
[10:15:48] - INFO - absl - Starting the local TPU driver.
[10:15:48] - INFO - absl - Unable to initialize backend 'tpu_driver': Not found: Unable to find driver in registry given worker: local://
[10:15:48] - INFO - absl - Unable to initialize backend 'gpu': Not found: Could not find registered platform with name: "cuda". Available platform names are: TPU Host Interpreter
```

## Expected behavior

I think this warning / error should not be displayed and the TPU should be correctly configured.



Comments:
Given that when running this command I'm getting ~30 training iterations per second, I'm assuming though that it's just a warning and not an error.

I am getting the same errors here. Also noting that the parameters from TrainingArguments are written at the start of the run_mlm_script, I see a few weird settings, like:
```
no_cuda=False,
tpu_num_cores=None,
``` 
I am not getting a stable report on iterations per second, so it is hard to see if this is a going well. Progress is still at 0%. Occationally, I am also getting error messages  like this written to the screen roughly each minute:

```
tcmalloc: large alloc 25737314304 bytes == 0x770512000 @  0x7ff6c92be680 0x7ff6c92df824 0x7ff6c92dfb8a 0x7ff49fbb6417 0x7ff49a9c43d0 0x7ff49a9d1ef4 0x7ff49a9d4e77 0x7ff49a9261dd 0x7ff49a6a0563 0x7ff49a68e460 0x5f5b29 0x5f66f6 0x50ad17 0x570296 0x56951a 0x5f60b3 0x5f6b6b 0x664e8d 0x5f556e 0x56ca9e 0x56951a 0x5f60b3 0x5f54e7 0x56ca9e 0x5f5ed6 0x56b3fe 0x5f5ed6 0x56b3fe 0x56951a 0x5f60b3 0x5f54e7
```
It might also be mentioned that I am initially getting "Missing XLA configuration". I had to manually set the environment variable: `export XRT_TPU_CONFIG="localservice;0;localhost:51011".` for the script to run. I am not sure if this really does the right thing. Maybe also the tpu driver needs to be specified?



...and ... I see this warning, tat does not look good:
`[10:28:02] - WARNING - absl - No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)`

If of interest, this is my notes on installing this. It is mainly based on Patricks tutorial:
https://github.com/NBAiLab/notram/blob/master/guides/flax.md

The simplest way to see if you're running on TPU is to call `jax.devices()`. E.g. you may see:

```
[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),
 TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),
 TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),
 TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),
 TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),
 TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),
 TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),
 TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]
 ```

@avital. "import jax;jax.devices()" gives me exactly the same response.

I am also able to make simple calculations on the TPU. The problem seem only to be related to the run_mlm_flax-script.

@avital. I also found this log file. It seems to think that the TPU is busy.
```
E0630 13:24:28.491443   20778 kernel_dma_mapper.cc:88] Error setting number simples with FAILED_PRECONDITION: ioctl failed [type.googleapis.com/util.ErrorSpacePayload='util::PosixErrorSpace::Device or resource busy']
E0630 13:24:28.491648   20778 tensor_node.cc:436] [0000:00:04.0 PE0 C0 MC-1 TN0] Failed to set number of simple DMA addresses: FAILED_PRECONDITION: ioctl failed [type.googleapis.com/util.ErrorSpacePayload='util::PosixErrorSpace::Device or resource busy']
E0630 13:24:28.491660   20778 driver.cc:806] [0000:00:04.0 PE0 C0 MC-1] tensor node 0 open failed: FAILED_PRECONDITION: ioctl failed [type.googleapis.com/util.ErrorSpacePayload='util::PosixErrorSpace::Device or resource busy']
E0630 13:24:28.491678   20778 driver.cc:194] [0000:00:04.0 PE0 C0 MC-1] Device has failed. Status:FAILED_PRECONDITION: ioctl failed [type.googleapis.com/util.ErrorSpacePayload='util::PosixErrorSpace::Device or resource busy']
E0630 13:24:28.492048   20777 kernel_dma_mapper.cc:88] Error setting number simples with FAILED_PRECONDITION: ioctl failed [type.googleapis.com/util.ErrorSpacePayload='util::PosixErrorSpace::Device or resource busy']
E0630 13:24:28.492093   20777 tensor_node.cc:436] [0000:00:05.0 PE0 C1 MC-1 TN0] Failed to set number of simple DMA addresses: FAILED_PRECONDITION: ioctl failed [type.googleapis.com/util.ErrorSpacePayload='util::PosixErrorSpace::Device or resource busy']
E0630 13:24:28.492100   20777 driver.cc:806] [0000:00:05.0 PE0 C1 MC-1] tensor node 0 open failed: FAILED_PRECONDITION: ioctl failed [type.googleapis.com/util.ErrorSpacePayload='util::PosixErrorSpace::Device or resource busy']
E0630 13:24:28.492110   20777 driver.cc:194] [0000:00:05.0 PE0 C1 MC-1] Device has failed. Status:FAILED_PRECONDITION: ioctl failed [type.googleapis.com/util.ErrorSpacePayload='util::PosixErrorSpace::Device or resource busy']
E0630 13:24:28.492996   20778 driver.cc:165] [0000:00:04.0 PE0 C0 MC-1] Transitioned to State::FAILED, dumping core
E0630 13:24:28.493215   20777 driver.cc:165] [0000:00:05.0 PE0 C1 MC-1] Transitioned to State::FAILED, dumping core
E0630 13:24:28.494112   20786 kernel_dma_mapper.cc:88] Error setting number simples with FAILED_PRECONDITION: ioctl failed [type.googleapis.com/util.ErrorSpacePayload='util::PosixErrorSpace::Device or resource busy']
E0630 13:24:28.494225   20786 tensor_node.cc:436] [0000:00:07.0 PE0 C3 MC-1 TN0] Failed to set number of simple DMA addresses: FAILED_PRECONDITION: ioctl failed [type.googleapis.com/util.ErrorSpacePayload='util::PosixErrorSpace::Device or resource busy']
```


I am now able to get this to run, and think I understand why this is happening.

The initial error I am seeing is: "RuntimeError: tensorflow/compiler/xla/xla_client/computation_client.cc:273 : Missing XLA configuration"

I have been getting around this error by setting "export XRT_TPU_CONFIG="localservice;0;localhost:51011". This is probably the right way of doing it on torch system, but it also leads to torch stealing the TPU from Jax, and letting the run_mlm_flax.py train on CPU.

However, it seems like there is a function in transformers/training_args.py called "is_torch_tpu_available". If this returns TRUE, it also asks for the XRT_TPU_CONFIG. I am really not sure why it is returning TRUE on my system but it might be because the VM I am using have other preinstalled software. 

Lots of ways of fixing this of course. You guys probably know the best way. 

I've seen the same `INFO` message from `absl`:

![image](https://user-images.githubusercontent.com/20651387/124166237-dd763200-daa2-11eb-910d-fb86cf16cff5.png)

But training seems to work (batch size 128 as from the mlm flax documentation)

Closing this issue now as it is expected

I am having the exact same issue described here. I see that @peregilk found a work around, but he/she hasn't shared what it was.

Could you describe how you overcame this issue? @peregilk 

@erensezener I think a lot has changed in the code here since this was written. I am linking to my internal notes above. I have repeated that one several times, and know it gets a working system up and running.

Just a wild guess: Have you tried setting ```export USE_TORCH=False```

> Just a wild guess: Have you tried setting `export USE_TORCH=False`

This solves the issue indeed! Thank you, you saved me many more hours of debugging :)