Slow evaluation using Trainer with TPUs in Colab
Author: finiteautomata
Pull request: no

## Environment info
- `transformers` version: 4.3.3
- Platform: Linux-4.19.112+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.10
- PyTorch version (GPU?): 1.9.0a0+7a178a8 (False)
- Tensorflow version (GPU?): 2.4.1 (False)
- Using GPU in script?: TPU
- Using distributed or parallel set-up in script?: NO

@sgugger @patrickvonplaten 

Model I am using (Bert, XLNet ...): BERT

I'm having very slow eval times using the `Trainer` API in conjunction with `XLA` in Google Colab. While the training epochs are running at a good speed, evaluating after each epoch it takes a very long time. I've tried restricting dataset size and tokenization max length with no success. 

I'm not sure how to check whether it's using `XLA` during evaluation. 

The task I am working on is NLI, using `multi-nli` from `datasets`

## To reproduce

Execute this notebook

https://colab.research.google.com/drive/1dVEfoxGvMAKd0GLnrUJSHZycGtyKt9mr?usp=sharing

## Expected behavior

Evaluation speed should be approximately the same as training.


Comments:
The notebook won't execute on TPU, you need to spawn a function on multiple processes for this (`xm.spawn(train_function)`). That function should contain all the training code including the `Trainer`, but `Trainer.train` by itself won't spawn multiple processes.

The recommended way to train on TPU is to follow the steps in the [examples](https://github.com/huggingface/transformers/tree/master/examples#running-on-tpus) to run the scripts.

Thanks for your answer @sgugger. Is there any plan to add an easier way to use TPUs in Colab?

I don't know of any easier way than launching the training function (in PyTorch). If you come across an easy example, please let me know and we will try to make the `Trainer` as easy to use.

Ok, sorry. I think I misunderstood. I thought that I should create a separate module for the training function because of the same reason that `multiprocessing` has issues with jupyter environments

I tried moving everything to a function and using `xmp.spawn(train_nli, args=())`, but I get this error which is not quite clear:

```python
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-4-d4081c64cb6f> in <module>()
      5 
----> 6 xmp.spawn(train_nli, args=())

2 frames
/usr/local/lib/python3.7/dist-packages/torch_xla/distributed/xla_multiprocessing.py in spawn(fn, args, nprocs, join, daemon, start_method)
    393         join=join,
    394         daemon=daemon,
--> 395         start_method=start_method)
    396 
    397 

/usr/local/lib/python3.7/dist-packages/torch/multiprocessing/spawn.py in start_processes(fn, args, nprocs, join, daemon, start_method)
    155 
    156     # Loop on join until it returns True or raises an exception.
--> 157     while not context.join():
    158         pass
    159 

/usr/local/lib/python3.7/dist-packages/torch/multiprocessing/spawn.py in join(self, timeout)
    110                 raise Exception(
    111                     "process %d terminated with exit code %d" %
--> 112                     (error_index, exitcode)
    113                 )
    114 

Exception: process 7 terminated with exit code 1
```

Any ideas? 

(Everything is on the same notebook as before)

Ok, I followed this notebook ([T5 on TPU](https://colab.research.google.com/github/patil-suraj/exploring-T5/blob/master/T5_on_TPU.ipynb)) and I managed to solve that error by using **`start_method="fork"`** on `xmp.spawn`. Thanks for your help @sgugger!

```python
def train_nli(index):
   # All the training code here
   ...
   
xmp.spawn(train_nli, args=(), start_method="spawn")
```

The notebook with the full code is [here](https://colab.research.google.com/drive/1dVEfoxGvMAKd0GLnrUJSHZycGtyKt9mr#scrollTo=k-e4NqfrtrJy)