[Flax] Add T5 pretraining script
Author: patrickvonplaten
Pull request: yes

# 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 adds T5 pretraining in Flax. Thanks to @craffel a lot of the preprocessing code was copied from  .
Tokenizer training code was largely copied from @SaulLu .

Also took this PR as a chance to integrate the new `push_to_hub` functionality that includes tensorboard logs to test out the new tensorboard functionality (cc @sgugger @LysandreJik @julien-c). 
The tensorboard logs aren't correctly displayed though :-/- an example can be seen [here](https://huggingface.co/patrickvonplaten/dummy-t5-test). 

Code is working, and the model seems to train. Will test it on a full training run over the weekend!

## Who can review?

@LysandreJik @sgugger - would be great if you could check the README.md and the `push_to_hub=True` logic / process to see if the workflow fits
@SaulLu - would be great if you could take a look at the tokenizer code, since it's 99% copied from yours :-) (it seems to work well)
@patil-suraj @sgugger - would be awesome if you could make a more general review to see if code is written according to examples and you are fine with having a rather model-specific training script in the general examples.

Comments:
I just wanted to ask one last question about the tokenizer for norwegian. :slightly_smiling_face: 

It seems to me that the tokenizer vocabularies of the "original" T5 models have `extras_ids` (100 for `T5-small`). It seems to me that in the proposed version of the tokenizer for norwegian, no extras_ids are introduced and I'm not sure where the [`extras_ids`](https://github.com/huggingface/transformers/blob/master/src/transformers/models/t5/tokenization_t5_fast.py#L112) argument would be redefined when initializing `T5TokenizerFast`.

> I just wanted to ask one last question about the tokenizer for norwegian. 
> 
> It seems to me that the tokenizer vocabularies of the "original" T5 models have `extras_ids` (100 for `T5-small`). It seems to me that in the proposed version of the tokenizer for norwegian, no extras_ids are introduced and I'm not sure where the [`extras_ids`](https://github.com/huggingface/transformers/blob/master/src/transformers/models/t5/tokenization_t5_fast.py#L112) argument would be redefined when initializing `T5TokenizerFast`.

Those additional tokens will be added when the tokenizer is loaded with `T5TokenizerFast` as follows:

```python
from transformers import T5TokenizerFast
tokenizer = T5TokenizerFast.from_pretrained("patrickvonplaten/t5-small-norwegian")
```

When you print out the tokenzier:

```python
print(tokenizer)
```

you should see that the extra ids have been added :-) It is done automatically [here](https://github.com/huggingface/transformers/blob/e27707488911a4bae5936a1bdad0cfdb2018cebd/src/transformers/models/t5/tokenization_t5_fast.py#L117)



@patrickvonplaten , run_t5_mlm_flax.py uses different lr schedule than paper. Any specific reason for that?

Hey @danshirron! 

Usually, it shouldn't make a big difference. Original T5 model was trained using an inverse square root scheduler. From my experiments, linear scheduler happens to be slightly more robust for faster convergence. 

Either way, there are various inverse square root scheduler implementations (e.g., [optax](https://github.com/formermagic/git-t5/blob/main/git_t5/core/schedulers.py#L124) or [pytorch](https://github.com/pytorch/fairseq/blob/master/fairseq/optim/lr_scheduler/inverse_square_root_schedule.py)).