Fix TF Trainer for TPU
This PR try to fix the trainer for TensorFlow by updating the order of some steps:
1. Dataset preprocessing
2. Strategy creation
3. Model creation

Instead of

1. Strategy creation
2. Model creation
3. Dataset preprocessing

Fixes #6672 

I have added a model_init function in the PyTorch Trainer to support hp-search. Is it possible to use this instead of changing the `args`? This would make a very big difference between the PT Trainer and TF Trainer.<|||||>OK, I will check this. Thanks.<|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
<|||||>Hello did you have any success fixing this? Can I help? I'm on a tight art student collegiate budget and tpu speed would be awesome if not necessary. I've spent like 20-30 hours on fixing the tpu issue myself and no luck. Any help getting run_clm.py on a tpu so I can quickly iterate would be awesome. But generally I'd love to move mainly to tpu but I'm not sure its there yet. New to open source really want to learn as much as possible. Can I help?<|||||>@arccoxx This PR should be closed because we have identified two different issues:

1. The first one don't come from Transformers but from TensorFlow. To make it short, TPU don't handle `tf.data.Dataset.from_generator`, Google is currently working on it and we have to wait they release the fix once they have one.
2. Currently you cannot train a LM from scratch with any TF model. We are currently working on this, and it will be possible in our next release.

So for your project the best solution would be to use the PyTorch version that works on TPU and you can train from scratch any LM model.<|||||>> @arccoxx This PR should be closed because we have identified two different issues:
> 
> 1. The first one don't come from Transformers but from TensorFlow. To make it short, TPU don't handle `tf.data.Dataset.from_generator`, Google is currently working on it and we have to wait they release the fix once they have one.
> 2. Currently you cannot train a LM from scratch with any TF model. We are currently working on this, and it will be possible in our next release.
> 
> So for your project the best solution would be to use the PyTorch version that works on TPU and you can train from scratch any LM model.

I was not able to get any pytorch version to run on xla. Is there any reference notebook that could be linked? I tried finetuning in native pytorch, running (pytorch) tuner, run_language_modeling with multiple transformers library versions 2.1.0-2.9.1, and run_clm with 3.4.0 all with no luck. Ive also tried building a pytorch lightning module and no luck. As the speedup would be that helpful (provided it can handle gpt2 medium) it would be awesome to figure out reduce these compatibility issues. My hope is to then use the tpu in a more complicated model that will use this fine tuned model. Any help would be super appreciated. Thank you!