Cast logits from bf16 to fp32 at the end of TF_T5
Author: szutenberg
Pull request: yes

# What does this PR do?

This change enables tf.keras.mixed_precision with bf16

I found that T5 model does not follow the [official TF guidelines regarding mixed precision](https://www.tensorflow.org/guide/mixed_precision). Therefore it's impossible to use 
`tf.keras.mixed_precision.set_global_policy('mixed_bfloat16')` which is the recommended way of training on bfloat16.

I took a notebook [snapthat/TF-T5-text-to-text](https://github.com/snapthat/TF-T5-text-to-text/blob/master/snapthatT5/notebooks/TF-T5-%20Training.ipynb) and added `tf.keras.mixed_precision.set_global_policy('mixed_bfloat16')`.

Experiments were done on tensorflow-cpu == 2.4.1, datasets == 1.8.0, transformers == 4.6.1.

The first issue is with the loss curve:
![loss_t5](https://user-images.githubusercontent.com/37601244/123172656-b81f6d80-d47d-11eb-886a-f3974321f9d8.PNG)

And the second is with inference (also included in the notebook):
```
  File "bf16_experiment.py", line 136, in <module>
    max_length=decoder_max_len, top_p=0.95, top_k=50, repetition_penalty=2)
  File "/home/mszutenberg/venv24/lib/python3.6/site-packages/transformers/generation_tf_utils.py", line 417, in generate
    use_cache=use_cache,
  File "/home/mszutenberg/venv24/lib/python3.6/site-packages/transformers/generation_tf_utils.py", line 472, in _generate_no_beam_search
    next_token_logits = tf.math.multiply(next_token_logits, next_token_logits_penalties)
  File "/home/mszutenberg/venv24/lib/python3.6/site-packages/tensorflow/python/util/dispatch.py", line 201, in wrapper
    return target(*args, **kwargs)
  File "/home/mszutenberg/venv24/lib/python3.6/site-packages/tensorflow/python/ops/math_ops.py", line 518, in multiply
    return gen_math_ops.mul(x, y, name)
  File "/home/mszutenberg/venv24/lib/python3.6/site-packages/tensorflow/python/ops/gen_math_ops.py", line 6068, in mul
    _ops.raise_from_not_ok_status(e, name)
  File "/home/mszutenberg/venv24/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 6862, in raise_from_not_ok_status
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: cannot compute Mul as input #1(zero-based) was expected to be a bfloat16 tensor but is a float tensor [Op:Mul]
```

This fix solves both issues. I see that the problem may also occur in other models. I can check and prepare fixes if this PR is approved.

## Who can review?

@LysandreJik @patrickvonplaten

Comments:
cc @Rocketknight1 

Hey! Firstly, good catch with this issue - this seems like a good PR. Two questions before we merge it, though:

- You tested with tensorflow-cpu - I'm guessing this means you were running on TPU, correct? (It's not like most CPUs support `bfloat16`, after all)
- The code only checks for the dtype `bfloat16` and not `float16`. I'm guessing the same issue might occur on GPUs with float16 dtype, so we should probably cast to float32 in either case. If you don't have access to a GPU or GPU instance, would you like to send me your code or a notebook so I can test it?

Hi @Rocketknight1 ,

I used Kaggle Code to run my script on TPU and:
- performance was the same on fp32 and bf16 (_TPU uses bf16 under the hood_)
- accuracy issue did not occur (I suspect that TPU modifies SparseSoftmaxCrossEntropyWithLogits precision: bf16->fp32)
- inference issue was reproduced with mixed_bfloat16

I used `os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'` in order to run bfloat16 on CPU. The loss curve in my PR comes from such execution.

I guess that adding cast to `float32` is required for `float16` too but I was getting `loss = nan` while attempting to run my script with `mixed_float16`. Maybe something is still broken or I do the conversion to fp16 incorrectly.

You can find my script in https://gist.github.com/szutenberg/80f30b980c15e384200d86ae242a1067

Output on TPU:
```
 1/10 [==>...........................] - ETA: 14:11 - accuracy: 0.2245 - loss: 12.9219
step 1: 94563.1 ms
 2/10 [=====>........................] - ETA: 0s - accuracy: 0.3733 - loss: 7.3379    
step 2: 104.5 ms
 3/10 [========>.....................] - ETA: 0s - accuracy: 0.4637 - loss: 5.2188
step 3: 84.6 ms
 4/10 [===========>..................] - ETA: 0s - accuracy: 0.5263 - loss: 4.1115
step 4: 86.1 ms
 5/10 [==============>...............] - ETA: 0s - accuracy: 0.5731 - loss: 3.4062
step 5: 84.6 ms
 6/10 [=================>............] - ETA: 0s - accuracy: 0.6100 - loss: 2.9387
step 6: 85.4 ms
 7/10 [====================>.........] - ETA: 0s - accuracy: 0.6399 - loss: 2.5959
step 7: 85.9 ms
 8/10 [=======================>......] - ETA: 0s - accuracy: 0.6644 - loss: 2.3573
step 8: 85.7 ms
 9/10 [==========================>...] - ETA: 0s - accuracy: 0.6849 - loss: 2.1528
step 9: 85.7 ms
10/10 [==============================] - 95s 88ms/step - accuracy: 0.7167 - loss: 1.9842
...
InvalidArgumentError: cannot compute Mul as input #1(zero-based) was expected to be a bfloat16 tensor but is a float tensor [Op:Sub]
```

Output on XLA_CPU:
```
 1/10 [==>...........................] - ETA: 10:50 - accuracy: 0.2245 - loss: 15.4375
step 1: 72235.4 ms
 2/10 [=====>........................] - ETA: 5:28 - accuracy: 0.3718 - loss: 9.4609
step 2: 41093.1 ms
 3/10 [========>.....................] - ETA: 4:43 - accuracy: 0.4576 - loss: 7.0469
step 3: 40026.9 ms
 4/10 [===========>..................] - ETA: 4:02 - accuracy: 0.5160 - loss: 5.7227
step 4: 39995.5 ms
 5/10 [==============>...............] - ETA: 3:20 - accuracy: 0.5595 - loss: 4.8281
step 5: 38984.9 ms
 6/10 [=================>............] - ETA: 2:39 - accuracy: 0.5933 - loss: 4.3073
step 6: 38997.0 ms
 7/10 [====================>.........] - ETA: 1:59 - accuracy: 0.6205 - loss: 3.9721
step 7: 39454.5 ms
 8/10 [=======================>......] - ETA: 1:19 - accuracy: 0.6420 - loss: 3.7764
step 8: 38780.1 ms
 9/10 [==========================>...] - ETA: 39s - accuracy: 0.6598 - loss: 3.5469
step 9: 39538.0 ms
10/10 [==============================] - 428s 40s/step - accuracy: 0.6877 - loss: 3.3266
...
tensorflow.python.framework.errors_impl.InvalidArgumentError: cannot compute Mul as input #1(zero-based) was expected to be a bfloat16 tensor but is a float tensor [Op:Mul]
```

You can see that loss behaves differently. After applying my patch everything is ok:
```
 1/10 [==>...........................] - ETA: 10:00 - accuracy: 0.2245 - loss: 12.3656
step 1: 66688.5 ms
 2/10 [=====>........................] - ETA: 5:13 - accuracy: 0.3773 - loss: 7.0029  
step 2: 39146.4 ms
 3/10 [========>.....................] - ETA: 4:32 - accuracy: 0.4685 - loss: 4.9963
step 3: 38585.8 ms
 4/10 [===========>..................] - ETA: 3:51 - accuracy: 0.5314 - loss: 3.9416
step 4: 38006.6 ms
 5/10 [==============>...............] - ETA: 3:12 - accuracy: 0.5789 - loss: 3.2643
step 5: 37900.7 ms
 6/10 [=================>............] - ETA: 2:33 - accuracy: 0.6164 - loss: 2.8159
step 6: 38225.8 ms
 7/10 [====================>.........] - ETA: 1:55 - accuracy: 0.6465 - loss: 2.4948
step 7: 38605.7 ms
 8/10 [=======================>......] - ETA: 1:16 - accuracy: 0.6711 - loss: 2.2820
step 8: 38495.9 ms
 9/10 [==========================>...] - ETA: 38s - accuracy: 0.6914 - loss: 2.0957 
step 9: 38342.7 ms
10/10 [==============================] - 413s 38s/step - accuracy: 0.7229 - loss: 1.9338
...
We went on a trip to Europe. We had our breakfast at 7 am in the morning at the nearby coffee shop. Wore a dark blue over coat for our first visit to Louvre Museum to experience history and art.
At what time did we had breakfast?
Answer:  <pad> 7 am in the morning

```

Thanks for the detailed log! I have a 30-series GPU here, so I'll try mixed_float16 and mixed_bfloat16 with your script when I get a chance and see if I get the same issues.

Hi @Rocketknight1,

Any updates?

I managed to get rid of nans on mixed_float16 by adding fp32 casts in:
* TFT5MainLayer for calculating extended_attention_mask:
```
if extended_attention_mask.dtype == tf.float16:
    extended_attention_mask = tf.cast(extended_attention_mask, tf.float32)

extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
```
* TFT5Attention before softmax:
```
if scores.dtype == tf.float16:
    scores = tf.cast(scores, tf.float32)

scores += position_bias
weights = tf.nn.softmax(scores, axis=-1)
```

but it's still not training (accuracy is 1.0 and loss around 10 - they are the same in each step, it seems that forward part is broken).

What do you think about merging my fix for bf16 and fixing fp16 later, by a separate PR?

I tried testing this. 'mixed_bfloat16' doesn't actually get run on the GPU on Keras, even though I believe 30-series GPUs support it. Very few CPUs support bfloat16 arithmetic, so I presume that no bfloat16 operations are used on CPU either, and that 'mixed_bfloat16' only actually runs bfloat16 operations on TPUs. As such, I'm confused about what's going on here - I suspect the differences you see with the 'mixed_bfloat16' policy on CPU are caused by some other side effect rather than true bfloat16 computation.

I'd like to resolve this before approving the bfloat16 PR - if it turns out that TPUs already handle this issue and no other hardware actually runs bfloat16, then this PR isn't necessary, although your effort is appreciated anyway!

Also the float16 fix might still be very useful - if you get it working, and you notice a difference on GPU with and without the cast to float32, please let us know!

@Rocketknight1, bfloat16 is not supported by GPU in TF. Ampere supports bfloat16 but the support for this datatype wasn't added to TensorFlow. For example [this article](https://developer.nvidia.com/blog/accelerating-ai-training-with-tf32-tensor-cores/) says only about adding TensorFloat32 and using bf16 directly from CUDA, not TF. 

I'm working on a custom accelerator which does support bf16 and I have exactly the same issue as on CPU (broken loss and broken inference). Inference is broken on a graph level (the model does not follow the rule _make sure the (model) output is float32_).

Did you run my reproducer on CPU? It should work on tensorflow-cpu==2.4.1. You can decrease batch_size to 1 and still be able to see the difference. In order to run on TPU (you can use free Kaggle code) set "use_tpu=True" in the script.

Hi, I'm sorry for the delay here! I'm probably explaining myself badly, though - what I want to know before I merge anything here is what exactly your reproducer is doing on CPU. The reason I'm asking is that bfloat16 is not supported except by a few rare CPUs, and I don't think bfloat16 on CPU is supported by TF/Keras at all. So I don't really understand what's happening when you run that code on CPU - I realize something changes, but I don't know what or why!

Hi @Rocketknight1 

Sorry for the delay caused by the summer season ;) 

My reproducer allows using bfloat16 by enabling XLA_CPU device which registers kernels for bfloat16 too. In the current TF version, bfloat16 kernels are not being registered for the CPU device. This is just to show that something is wrong also with the training accuracy. In my opinion, proof that something is wrong with inference is enough to accept this change.

Other models and templates should be reviewed as well. What do you think?

Hi @szutenberg - after the conversation in #12898 I think I'm happy to accept this PR. Could you change it to check if the dtype is either `float16` or `bfloat16`? Alternatively, you could just run the casts without an `if` statement - it will do nothing if the dtype is already `float32`.

Hi @Rocketknight1 , thanks! This PR is ready for merge.

Done! Thank you for this, and thanks for your patience with the review process too!