Add Flax BART pretraining script
Author: duongna21
Pull request: yes

# What does this PR do?

Fixes #6743 #18030 #4151 #5096. Adds Flax script for BART pretraining.

Inspired by @patil-suraj's [suggestion](https://github.com/huggingface/transformers/issues/6743#issuecomment-1046717336), I modified the @morganmcg1's [DataCollatorForDenoisingTasks](https://github.com/morganmcg1/rotobart/blob/main/data_collator.py#L223) to create a BART denoising pretraining script in Flax.

Implementation details from the paper:
- [x] Text infilling
- [x] Sentence permutation
- [ ] Large training batch sizes (will add gradient accumulation for this and other Flax language modeling scripts in the next PR)

[Training statistics](https://tensorboard.dev/experiment/Maw62QlaSXWS0MOf2V2lbg/) when pre-train bart-base in Norwegian on a single TPUv3-8 pod: 

<img src="https://imgur.com/N3PnWtn.png" width="300" height="220" />
<img src="https://i.imgur.com/1SKROC2.png" width="300" height="220" />
<img src="https://i.imgur.com/KDACVb4.png" width="300" height="220" />

## Who can review?

cc potential reviewers: @patrickvonplaten, @patil-suraj, @sgugger, @LysandreJik 

Comments:
_The documentation is not available anymore as the PR was closed or merged._

This looks very nice to me! Just a small question, why is the file called `run_bart_dlm_flax.py`, *i.e.* why the `dlm` and not just `lm` ?

@patil-suraj could you take a look here as well ?

> This looks very nice to me! Just a small question, why is the file called `run_bart_dlm_flax.py`, _i.e._ why the `dlm` and not just `lm` ?

@patrickvonplaten It stands for `denoising language modeling`, which is consistent with `mlm` and `clm`. What do you think?

Your reviews are extremely helpful @sanchit-gandhi! All of them have been resolved. Thank you! 

@patil-suraj Could you have a look at this PR? I'd love to hear your feedback.

Hey @duongna21! Great job on this PR, and thank you for addressing the comments!

For the time being, let's hold off on gradient accumulation. If you require it for your personal experiments, it can be achieved quite easily using the Optax wrapper [MultiSteps](https://optax.readthedocs.io/en/latest/api.html#multi-step-update). However, in my experience, this wrapper is pretty memory inefficient and does not yield particularly good performance; it applies a dummy update of zeros for $K-1$ train steps (redundant), and then the accumulated gradient update step on the $K$-th train step. Instead, writing a custom gradient accumulation training loop is more efficient (c.f. [seq2seq-speech](https://github.com/sanchit-gandhi/seq2seq-speech/blob/95857f52b4d8dc5c1ab48835899d6f66e86ba9b1/run_flax_speech_recognition_seq2seq.py#L1226)), but this involves quite a lot of additional code and is significantly more involved, so I'm not particularly in favour of using it for these streamlined examples scripts!

Otherwise, all the implementation TODOs are complete, the code review approved, and the training results on track, so happy to go ahead and merge!

@sanchit-gandhi Yeah, I agree that gradient accumulation shouldn't be added until there is a more elegant way to implement it, so feel free to merge this PR. Thank you for the helpful advice!

@duongna21 I wonder why the [`permute_sentences`](https://github.com/huggingface/transformers/blob/main/examples/flax/language-modeling/run_bart_dlm_flax.py#L318) only contains the pad token rather the full stop token?