Fix grad loss computation
Author: andyehrenberg
Pull request: yes

In the flax summarization fine-tuning example, the loss is only computed where the decoder_attention_mask is 1. Different batches on different devices will have different decoder_attention_masks, and `jax.lax.pmean(loss, axis_name="batch")` doesn't take this into account, so it won't be equivalent to if the loss was computed on all batches put together on a single device. To fix this, this PR computes the number of tokens for which the loss was computed on each device, and multiply the per-device losses and gradients by these weights, then `lax.psum` the losses, gradients and weights before then dividing by the psummed weights.


- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] 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?


## Who can review?

@patil-suraj @sgugger 

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

Hey @andyehrenberg! This is a great spot! Indeed the loss isn't computed correctly. I've formalised an argument for this mathematically here (not rendering very well... easier viewed in markdown or LaTex):

<details>

<summary> Mathematical Proof </summary>

Technically speaking, in the `train_step`, the `pmap` won't compute a 'true' mean over devices. Here, what we're doing is computing a normalised loss on each device, and then averaging these losses over devices. This isn't strictly equal to summing the losses over all devices, and then dividing by the number of samples. 

Let $K$ denote the number of devices. Denote the loss on the $i$-th device as $L_i$ (`loss.sum()`) and the number of samples $N_i$ (`label_mask.sum()`). In the `loss_fn`, we compute the normalised loss on each device (`loss.sum() / label_mask.sum()`):

$$\bar{L}_i = \frac{L_i}{N_i}$$

and then average over devices with the `pmap`:

$$\mathcal{L} = \frac{1}{K} \sum_{i=1}^{K} \frac{L_i}{N_i}$$

Whereas, for a 'true' loss, we should first add up all the losses over devices:

$$L_{tot} = \sum_{i=1}^{K} L_i $$

and then divide by the total number of labels:

$$\mathcal{L}' = \frac{L_{tot}}{N} =  \frac{1}{N}\sum_{i=1}^{K} L_i $$

where $N$ is the total number of labels:

$$ N = \sum_{i=1}^{K} N_i $$

If we compare the two and ignore the constant $K$ in the `pmap` average:

$$\mathcal{L} =  \sum_{i=1}^{K} \frac{L_i}{N_i}$$

$$ \mathcal{L}' =  \frac{1}{N}\sum_{i=1}^{K} L_i $$

we see that the losses are in-fact different. The first expression is what you get if you average the losses on each device, then average these terms over devices with a `pmap`. The second expression is a 'true' loss, what you get by summing the losses on each device, summing these losses over devices, and then dividing by the total number of terms in your batch (= sum of the `label_mask` per device, summing these terms over devices).

</details>

A PR to address this was merged here: https://github.com/huggingface/transformers/pull/19504

Here, we compute the 'true' number of loss terms by summing the num labels on each device, and then normalising our loss by the sum of the labels over devices.

You should be able to rebase onto main to get these changes and compute a 'true' loss, both for summarisation and all the other Flax training examples 🤗

Closing this PR as https://github.com/huggingface/transformers/pull/19504 (merged) addresses this issue. 

All the best with your Flax experiments!