FLAX-T5 - TPU not found Colab
Author: GenVr
Pull request: no

Hello,
I'm using the code `run_t5_mlm_flax.py` on Google Colab in TPU mode. 

I have the following problem:
![image](https://user-images.githubusercontent.com/82599723/138088969-eb7fed61-67b0-4a57-aff3-cf42c3edb032.png)
And also:
`/usr/local/lib/python3.7/dist-packages/jax/__init__.py:27: UserWarning: cloud_tpu_init failed: ConnectionError(MaxRetryError("HTTPConnectionPool(host='metadata.google.internal', port=80): Max retries exceeded with url: /computeMetadata/v1/instance/attributes/agent-worker-number (Caused by NewConnectionError('<urllib3.connection.HTTPConnection object at 0x7f4ff0494790>: Failed to establish a new connection: [Errno 110] Connection timed out'))")) `
`This a JAX bug; please report an issue at https://github.com/google/jax/issues`
`_warn(f"cloud_tpu_init failed: {repr(exc)}\n This a JAX bug; please report "`

The TPU is not found, and the code switch in CPU mode. I'm using these libraries:
`pip install datasets`
`pip install transformers`
`pip install flax`
`pip install optax`

and also this configuration I read:
`import jax.tools.colab_tpu`
`jax.tools.colab_tpu.setup_tpu()`
`print(jax.local_devices())`
`export XRT_TPU_CONFIG="localservice;0;localhost:51011"`
`unset LD_PRELOAD`
`USE_TORCH=0`


How can I do to use this code on Colab or to use a FLAX-T5 with TPU on Colab?

Thank you!

Comments:
cc @patil-suraj 

Hi @IVannarI!

The `run_t5_mlm_flax.py` should run on colab TPU. Could you maybe share the colab so we can reproduce the issue?

Also, `jax.tools.colab_tpu.setup_tpu()` should be called before any other code. Could you run this in the notebook and see if it detects the TPU ?

Hi @patil-suraj, thank your for you support!

`jax.tools.colab_tpu.setup_tpu()` detects the TPU in colab.

This is the link of a colab, with some comments: https://colab.research.google.com/drive/1e8TrnPtLWU16dImozCuh-tt_OSxbdzDH?usp=sharing
I put a dummy dataset, that will give an error, but it's to see the TPU problem. 

Running, i have this:

![image](https://user-images.githubusercontent.com/82599723/139820288-6b83537c-b08b-405f-8e76-6453a79b9e3f.png)


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

Hi guys,
@GenVr @patil-suraj  did you find a solution?
Getting the same errors running run_t5_mlm_flax.py on Colab

![image](https://user-images.githubusercontent.com/70054668/159219882-ba29de58-cff2-4291-a0c4-235a288c6a0c.png)

jax.local_devices() finds the device but the script switches to CPU


@0syrys see [here](https://discuss.huggingface.co/t/how-to-use-tpu-for-model-training-using-example-script-run-mlm-py/15510), it's solved

@GenVr 
thanks for the help!
Unfortunately it's still not working
Seems to me the xla_spawn is only needed for the pytorch scrips?
Also the flax version run_t5_mlm_flax.py does not take a "--tpu_num_cores 8" argument or similar
So still a bit clueless

I created a public Colab with all the needed files in it:
https://drive.google.com/drive/folders/19TnTy_-h-MfccWWwKkGrll4wsN5X64CE?usp=sharing

For now the flax script ran just fine on the Colab GPU 
But when I scale up the model size TPU would be nice I guess...

Can also try to switch to generic pytorch or tensorflow scripts for TPU training
But the whole T5 Preprocessing Pipeline is just perfect for what I'm doing


@0syrys Hello. I'm meeting the same issues as you mentioned here. "run_t5_mlm_flax.py" script can't find Colab TPU. I found the link your provided has broken. Is it possible to explain more about how to solve this issue? Really appreciate. Thanks!

Found the issue. We need to call 
```python
import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()
```
in the script before importing anything JAX related. Calling `setup_tpu()` in the colab and then launching the script won't work because these are two different processes. So adding these two lines in the script before any JAX/Flax import should fix this issue.


@patil-suraj
 Thanks that's it!