Add cross entropy loss with stable custom gradient
Author: yhavinga
Pull request: yes

# What does this PR do?

This PR adds Flax code for a cross entropy loss calculation with an additional term to stabilize gradients for bfloat16 training.

The loss function is authored by the T5X Authors (https://github.com/google-research/t5x/blob/90d74fa703075d8b9808ae572602bc48759f8bcc/t5x/losses.py#L25)

Also add 'z_loss' as training argument to the T5 Flax pre-training script.

If z_loss > 0, then an auxiliary loss equal to z_loss*log(z)^2
will be added to the cross entropy loss (z = softmax normalization constant).

The two uses of z_loss are:
  1. To keep the logits from drifting too far from zero, which can cause
     unacceptable roundoff errors in bfloat16.
  2. To encourage the logits to be normalized log-probabilities.

While the z_loss function is only added to the t5 flax pretraining script, this loss function might be interesting for
other flax pre-training scripts. I did not test this.
Finally, there is a (currently unused) function `compute_weighted_cross_entropy` with z_loss and label smoothing,
which might be useful for other flax training scripts as well.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [*] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [*] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

No but I tested with the t5 flax norwegian script, bfloat16 and z_loss set to 1e-4, the setting used in t5 gin configs.

In my own tests I've seen the following consistently:
Without z_loss and bfloat16, the loss will either diverge or converge on a higher plateau than training with float32.
With z_loss and bfloat16, set to 1e-4, loss curves almost match curves with float32 training.

## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@patrickvonplaten, @patil-suraj

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.


Comments:
float32 compared with bfloat16 without and with z_loss
![image](https://user-images.githubusercontent.com/3098618/184906406-88786977-7a52-4f59-8319-c818c7c85050.png)


The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_18652). All of your documentation changes will be reflected on that endpoint.

Hey @yhavinga,

Thanks a lot for your PR. Could we maybe add this to the `examples/research_folder` instead of the official examples?
The reason is that we won't have time to maintain this example and we would have need to check those loss curves on more than just Norwegian. 

Would that be ok for you? 

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/main/CONTRIBUTING.md) are likely to be ignored.

Apologies for the long delay.
In the meantime I've noticed that term added to the loss (z_loss * jax.lax.square(log_z)) might in fact be (similar to) L2 regularization, and that this kind of regularization might in fact already be available through Optax Adafactors weight_decay_rate parameter. I currently do not have access to TRC so cannot test this, but thought it might be interesting to others training with bfloat16 and run_t5_mlm_flax.py if they might hit this page.