Training on Tpu got stuck at 0%
I am trying to train a t5 model based on examples for Persian language. I cloned the example on Google Colab Pro, and chose TPU. I also add some code to the program to connect colab tpu via jax.


```
import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()
```

Then run:

```
!python run_t5_mlm_flax.py \
	--output_dir="/content/drive/MyDrive/Pouramini/" \
	--cache_dir="/content/drive/MyDrive/cache/" \
	--preprocessing_num_workers="20" \
  --model_type="t5" \
	--config_name="/content/drive/MyDrive/Pouramini/" \
	--tokenizer_name="/content/drive/MyDrive/Pouramini/" \
	--dataset_name="oscar" \
	--dataset_config_name="unshuffled_deduplicated_fa" \
	--max_seq_length="256" \
	--per_device_train_batch_size="16" \
	--per_device_eval_batch_size="16" \
	--adafactor \
	--learning_rate="0.005" \
	--weight_decay="0.001" \
	--warmup_steps="2000" \
	--logging_steps="500" \
	--save_steps="5000" \
	--eval_steps="50000" \
        --resume_from_checkpoint=True
```

Now over 1 hours, the progress got stuck at:

```
Epoch ... :   0% 0/3 [00:00<?, ?it/s]
Training...:   0% 0/161270 [00:00<?, ?it/s]

```
The number of steps indicates that the tokens were distributed on 8 Tpu core. The exact number of tokens is over 20M.
Anyway, I don't know what is doing, whether it's an stage that I must wait for it, or it hanged or something.

Just now it flushed hunderes of lines like below and stopped:
```
2021-10-10 20:07:11.931040: E external/org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/grpc_tpu_driver.cc:479] Resetting: 52612:2348383
2021-10-10 20:07:11.931048: E external/org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/grpc_tpu_driver.cc:479] Resetting: 52612:2357538
2021-10-10 20:07:11.931057: E external/org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/grpc_tpu_driver.cc:479] Resetting: 52612:2382164
2021-10-10 20:07:11.931065: E external/org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/grpc_tpu_driver.cc:479] Resetting: 52612:2391319
2021-10-10 20:07:11.931072: E external/org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/grpc_tpu_driver.cc:479] Resetting: 52612:2490748
2021-10-10 20:07:11.931080: E external/org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/grpc_tpu_driver.cc:479] Resetting: 52612:2418989
2021-10-10 20:07:11.931088: E external/org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/grpc_tpu_driver.cc:479] Resetting: 52612:2423186
2021-10-10 20:07:11.931095: E external/org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/grpc_tpu_driver.cc:479] Resetting: 52612:2471759
2021-10-10 20:07:11.931103: E external/org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/grpc_tpu_driver.cc:479] Resetting: 52612:2437978
2021-10-10 20:07:11.931112: E external/org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/grpc_tpu_driver.cc:479] Resetting: 52612:2456967
........
```

What does it mean?

Hi, this looks like a strange error.

> The number of steps indicates that the tokens were distributed on 8 Tpu core.

We specify `--per_device_train_batch_size=16` which means each TPU core will use 16 BS, so the effective total  BS is 16 * 8 = 256
and the number of steps are computed using this total BS.

It looks like the model is stuck at compilation, it could happen if the model cannot fit on the TPU core.
What is the size of the model, is it t5-base or t5-large or larger than that? 

Also, could you post the jax, jaxlib, and flax versions?<|||||>As I checked them using
```
!pip freeze | grep jax
!pip freeze | grep flax
```


it prints:
```
jax==0.2.21
jaxlib @ https://storage.googleapis.com/jax-releases/cuda111/jaxlib-0.1.71+cuda111-cp37-none-manylinux2010_x86_64.whl
flax==0.3.5
```

I also tried `--per_device_train_batch_size=8`, but I got the same error!<|||||>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.<|||||>I met the same error.<|||||>@patil-suraj @puraminy Same error here. I add these two lines in the script as well.
```
import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()
```

In Colab, Jax is still unable to connect TPU.
```
2022-05-13 02:37:09.265602: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:271] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected. 
[02:37:09] - INFO - absl -   Unable to initialize backend 'gpu': FAILED_PRECONDITION: No visible GPU devices.  
[02:37:09] - INFO - absl -   Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
```

And the training progress got stuck as well. I guess it is because the script uses CPU for training instead of TPU, so it is just too slow.. 
```
Epoch ... :   0% 0/3 [00:00<?, ?it/s]. 
Training...:   0% 0/6344 [00:00<?, ?it/s]
```
Do you have any solution? Many thanks.<|||||>Could you try to run these lines 

```python
import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()
```

and then print `jax.devies()` in colab cell to see if it can detect TPU ?
Adding those two lines in the script before any other JAX imports should resolve this issue.