[Demo] Slow down in TPU training
@avital @marcvanzee - I wanted to align `run_mlm_flax.py` more with `run_glue_flax.py` and noticed that by doing the change as shown in this PR, training on TPU slows down very significantly by ca. ~40%.

Currently, [`run_glue_flax.py`](https://github.com/huggingface/transformers/blob/master/examples/flax/text-classification/run_flax_glue.py) and [`run_mlm_flax.py`](https://github.com/huggingface/transformers/blob/master/examples/flax/language-modeling/run_flax_mlm.py) deal slightly differently with the PRNG key: `run_mlm_flax.py` splits the key inside the training step while `run_glue_flax.py` does so before the train step and shards it then before passing it to the train loop. It seems that `run_mlm_flax.py` is significantly faster on TPU.

Do you by any chance have good explanations for that? 
Basically prior to this change you were running a single jitted function in each train step, and because of asynchronous dispatch it didn't have to wait until the previous step was complete until dispatching the program for the next step. But if you split an RNG in between, then JAX blocks until the previous step was complete, then dispatches and executes the split command, and then dispatches the next training step.

In short, the guideline is that each step in a training loop should be a single jitted function. If done right this should lead to close to 100% device utilization.

This is a common gotcha -- people hit this regularly, and we should help catch the slow patterns early, such that you could detect this even with a local run with no accelerator/unit test. @jheek is working on a library that would allow you to annotate code such that you'd get that kind of error or warning for this, and other cases.

@jheek also said:

Yeah this is an example of my number 1 most common and most hurtful JAX performance gotcha that I want to catch automatically

In this case it stands out but there are more subtle variants where it's hard to spot in a review

This analysis is only true for TPU without async mode enabled btw. Because all other devices have a queue that is > 1
<|||||>Thanks a lot for this in-detail explanation @avital! 

Also pinging @sgugger @stas00 @mfuntowicz - might be interesting to read :-) <|||||>(I guess really this means that `run_glue_flax.py` could be made faster? /cc @marcvanzee )<|||||>> (I guess really this means that `run_glue_flax.py` could be made faster? /cc @marcvanzee )

Yeah, I'm currently testing it actually, see here: https://github.com/huggingface/transformers/pull/11778 .
Will report results tomorrow<|||||>Great that you discovered this! I actually didn't notice the bug, and since training was already fast enough I didn't look into it. Curious to see whether we will get even more speedup!<|||||>Reran, the experiments - got a small speed-up on TPU. Here the new numbers: https://github.com/huggingface/transformers/tree/master/examples/flax/text-classification#runtime-evaluation