TF2 TPU slow?
## Environment info
- `transformers` version: 3.0.2
- Platform: ubuntu 18.04
- Python version: 3.6
- Tensorflow version (GPU?): TF 2.3.0

### Who can help
@jplu 

## Information

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

The problem arises when using:
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [x] task: MLM

## To reproduce
I am using the TFTrainer from `transformers` for MLM pretraining. However, it seems that even for TPU training each batch is fed separately to the TPU, while it's usually more common to feed a bunch of batches to the TPU for efficiency (see https://github.com/tensorflow/models/blob/master/official/nlp/bert/model_training_utils.py#L226).

I am not sure that's the only problem, but MLM pretraining BERT is around 3x slower on TPU with the TFTrainer compared to the official implementation (https://github.com/google-research/bert).

For better TPU usage, we probably need sth like here:
https://github.com/tensorflow/models/blob/master/official/nlp/bert/model_training_utils.py#L345-L361
Hello!

Training on TPU is slow indeed because the training loop is not optimized for TPU. For two reasons:
1) we don't want to have a version that is specifically optimized for each device which will create to much confusion for maintenance.
2) having a training loop optimized for TPU will limit the possibility of logging, the logging will be at every epoch instead of every X steps. Which is a not wanted behavior.

Nevertheless if you have a solution that respect these two points, I will be happy to review your PR :)<|||||>You are right, this will introduce some behavior changes because logging/saving is not possible while the TPU is processing several batches.

I actually played around with this a bit and introduced a variable `steps_per_loop` that I set to `200` when using TPU. I then used an iterator for the dataset to only do batch feeding/optimization during this period. However, this only improved the training speed marginally, so I don't think it's worth a PR.

What gives me a bigger speedup (around 80%) is actually using 'keras compile/fit', where we can set `experimental_steps_per_execution` to different values depending on the device, for the training loop. I would be curious though if we could substitute part of the trainer or even the whole training by a "keras style" implementation. However, this would change the current architecture quite significantly and I am not sure it will retain the generic properties of the TFTrainer class.

Maybe it becomes clearer if I just show you how I implemented an option to use keras for the training loop in the TFTrainer (pretty hacky of course): https://gist.github.com/volker42maru/20641970599c27dc9503161f52aa67c9#file-tf_trainer_train-py-L85-L113


I will close this for now, because it's probably more of a long term prospect.