[Flax T5] Speed up t5 training
# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

This PR makes sure that no jax functionality is used in the preprocessing to make sure the TPU is not unnecessarly blocked. This small change leads to a **5x** factor speed-up in training T5.

🚨🚨 **Note**: It is extremely important to verify that no DeviceArrays are created during the data preprocessing to make sure that the training step can run asynchronously on TPU while the preprocessing runs on CPU. A good rule is to make sure that in the training step look, only the function `p_train_step` uses JAX/Flax code and all other functions run on CPU. Other relevant links: https://jax.readthedocs.io/en/latest/async_dispatch.html#async-dispatch 🚨🚨


The PR is tested here: https://huggingface.co/patrickvonplaten/t5-base-norwegian/tensorboard (check train loss graph which shows that time is reduced to < 5h now)