Tensorflow Mixed Precision Training
Author: zuyezheng
Pull request: no

## Environment info

- `transformers` version: 4.5.1
- Platform: Linux-5.4.0-74-generic-x86_64-with-glibc2.27
- Python version: 3.8.8
- PyTorch version (GPU?): 1.8.1+cu111 (True)
- Tensorflow version (GPU?): 2.6.0-dev20210604 (True)
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help

@LysandreJik @Rocketknight1 @sgugger

## Information

Model I am using (Bert, XLNet ...): Bert

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Use TF directly with `model.fit` or `TFTrainer` with policy `mixed_float16` for mixed precision training. 
2. Due to this tensorflow [cast](https://github.com/tensorflow/tensorflow/issues/50964) issue in SparseCategoricalCrossentropy loss used in many of the huggingface TF models, incorrect label encodings could result in `nan` or errors in loss. 
3. Errors can start with token (or class) indexes at 2k+ and `nan` loss with labels closer to the max. 

## Expected behavior

Correct loss and no `nan`. 

Changing `compute_loss` to use `CategoricalCrossentropy` vs sparse and manually one hot encoding solves this:

```
def compute_loss(labels, logits):
    loss_fn = tf.keras.losses.CategoricalCrossentropy(
        from_logits=True, reduction=tf.keras.losses.Reduction.NONE
    )
    # make sure only labels that are not equal to -100 affect the loss
    active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100)
    reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
    labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)
    
    **one_hot_labels = tf.one_hot(labels, tf.shape(logits)[-1], dtype=logits.dtype)**

    return loss_fn(one_hot_labels, reduced_logits)
```

Changing the last output layer to be float32 also solves this:

```
class TFBertMLMHead(tf.keras.layers.Layer):
    def __init__(self, config: BertConfig, input_embeddings: tf.keras.layers.Layer, **kwargs):
        super().__init__(**kwargs)

        self.predictions = TFBertLMPredictionHead(config, input_embeddings, name="predictions")
        **self.finalCast = tf.keras.layers.Activation('linear', dtype='float32')**

    def call(self, sequence_output: tf.Tensor) -> tf.Tensor:
        prediction_scores = self.predictions(hidden_states=sequence_output)
        **prediction_scores = self.finalCast(prediction_scores)**

        return prediction_scores
```

But given the recommendation that output be accumulated in float32 to be numerically stable, perhaps `transform_act_fn` and everything after needs to be `float32`?

Comments:
One-hot encoding the labels for a language model will get you OOM super fast since the vocab size is often large, so that's not an option. I think casting the prediction before the loss back to float32 is probably the safest option?

Ah, good point, I assumed since logits were of the same dimensionality, it wouldn't be too bad, but digging deeper in TF's sparse implementation, definitely more optimal. Interestingly, TF's internal [implementation](https://github.com/tensorflow/tensorflow/blob/57da85f8870bc8dee1b77225b3e30ea3f314d304/tensorflow/python/ops/nn_ops.py#L4185) even notes a requirement for labels to be of "dtype `int32` or `int64`" so I think it's their cast that needs to be fixed since it's still going from `int -> float32 -> int64` currently.

I settled with this loss function that does a cast of the logits in the meantime which also has a benefit (I think) of performing the final softmax in float32 vs float16.
```
@tf.function
def compute_loss(labels, logits):
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
        # make sure only labels that are not equal to -100 affect the loss
        active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100)
        reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
        labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)

        return loss_fn(labels, tf.cast(reduced_logits, tf.float32))
```

However, still curious if using float32 explicitly in an earlier layer such as for the activation function in `TFBertPredictionHeadTransform` might still be better?


This is a fascinating bug in Keras. It's a known issue that softmaxes can be unstable in float16 or bfloat16, but I didn't realize that this issue could also smear the labels around too. Tagging #12332 as well, which is a relevant PR. (And maybe this might finally explain my confusion with what was going on in that case!)

I think you're right that computing the logits in float32 across our models might still be an improvement for numerical stability reasons even if the label cast bug is fixed, though, and so it would be worth making that change even if the upstream Keras bug gets fixed. @sgugger wdyt?

In PyTorch, we always compute the softmax in FP32 as it's better for numerical stability. So yes, if possible, we should the same on the TF side.

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.