[WIP][examples/flax] add gradient accumulation
# What does this PR do?

Adds gradient accumulation in flax language modeling scripts.
Thanks a lot for adding this! That's super useful! It seems to require some bigger changes to core functionality to the script so I think we should be careful here. Also I'm starting to wonder whether the examples become to complicated to read with more and more functionality being added and whether we should maybe instead creating a new training script instead?

Also wouldn't it be better to use gradient accumulation functionality from `optax` such as https://optax.readthedocs.io/en/latest/api.html?highlight=ApplyEvery#optax.apply_every ? Think a lot of the code that is written here already exists in optax classes/functions no?

@sgugger - I'd love to have your feedback on the PR as well<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.<|||||>Unstale<|||||>Hi, I've just checked the flax examples on master branch and it seems that gradient accumulation is still missing, so I'm coming back to this PR :)

@patrickvonplaten mentioned to use `optax`, and I've found this (working?) implementation of gradient acc. for T5 MLM pre-training from @gsarti. This may could help here :hugs: 



