TFBert activation layer will be casted into float32 under mixed precision policy
Author: jlei2
Pull request: no

### To reproduce

```
import tensorflow as tf
from transformers import TFBertForQuestionAnswering
tf.keras.mixed_precision.experimental.set_policy('mixed_float16')
model = TFBertForQuestionAnswering.from_pretrained('bert-base-uncased')
```
User will receive warning saying

```
WARNING:tensorflow:Layer activation is casting an input tensor from dtype float16 to the layer's dtype of float32, which is new behavior in TensorFlow 2.  The layer has dtype float32 because its dtype defaults to floatx.

If you intended to run this layer in float32, you can safely ignore this warning. If in doubt, this warning is likely only an issue if you are porting a TensorFlow 1.X model to TensorFlow 2.

To change all layers to have dtype float16 by default, call `tf.keras.backend.set_floatx('float16')`. To change just this layer, pass dtype='float16' to the layer constructor. If you are the author of this layer, you can disable autocasting by passing autocast=False to the base Layer constructor.
```
This log means that under mixed-precision policy, the gelu activation layer in TFBert will be executed with float32 dtype. So the input x([link](https://github.com/huggingface/transformers/blob/d6c08b07a087e83915b4b3156bbf464cebc7b9b5/src/transformers/modeling_tf_bert.py#L91)) will be casted into float32 before sent into activation layer and casted back into float16 after finishing computation.

### Causes
This issue is mainly due to two reasons:

-  The dict of [ACT2FN](https://github.com/huggingface/transformers/blob/d6c08b07a087e83915b4b3156bbf464cebc7b9b5/src/transformers/modeling_tf_bert.py#L121) is defined outside of [TFBertIntermediate class](https://github.com/huggingface/transformers/blob/d6c08b07a087e83915b4b3156bbf464cebc7b9b5/src/transformers/modeling_tf_bert.py#L347). The mixed-precision policy will not be broadcasted into the activation layer defined in this way. This can be verified by adding `print(self.intermediate_act_fn._compute_dtype)` below this [line](https://github.com/huggingface/transformers/blob/d6c08b07a087e83915b4b3156bbf464cebc7b9b5/src/transformers/modeling_tf_bert.py#L360). float32 will be given.

- tf.math.sqrt(2.0) used in [gelu](https://github.com/huggingface/transformers/blob/d6c08b07a087e83915b4b3156bbf464cebc7b9b5/src/transformers/modeling_tf_bert.py#L98) will always return a float32 tensor even if it is under mixed-precision policy. So there will be dtype incompatibility error after we move the ACT2FN into TFBertIntermediate class. This can be further solved by specifying dtype of tf.math.sqrt(2.0) to be consistent with `tf.keras.mixed_precision.experimental.global_policy().compute_dtype
`.

### Pros & Cons 
Pros: activation layer will have higher precision as it is computed with float32 dtype
Cons: the two additional Cast operation will increase latency.

### PR
The PR that can help getting rid of warning log and make the activation layer be executed in float16 under mixed-precision policy: https://github.com/jlei2/transformers/pull/3 This PR is not a concrete one because there are some other models importing this ACT2FN dict so it can only solve the issue existed in TFBert related model.

### Benchmark Results
As Huggingface official benchmark tests don't support mixed-precision, the comparison was made using Tensorflow Profiling Tool(https://www.tensorflow.org/tfx/serving/tensorboard).

1 GPU-V100
Model: Huggingface FP16 Bert-Base with batch_size=128 and seq_len = 128.

<img width="350" alt="image (5)" src="https://user-images.githubusercontent.com/70337521/92769405-cae10f80-f34d-11ea-87c0-df4ec9ecdae2.png">
<img width="350" alt="image (5)" src="https://user-images.githubusercontent.com/70337521/92769851-39be6880-f34e-11ea-9926-2e8014c17e4e.png">

After applying the PR, the two Cast ops disappear and the time spent on activation layer is reduced from 3.745 ms to 1.75ms. 

This issue is not like a bug or error and the pros & cons are listed above. Will you consider make changes corresponding to this?  @patrickvonplaten  Thanks!

Comments:
Hello @jlei2!

This issue should be fixed in https://github.com/huggingface/transformers/pull/7022 

Gotcha! Thanks!

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


Hi @jplu , you mentioned that this issue should be fixed in #7022. But I noticed that that PR has been closed. Are you still planning to look at this? Thanks!

Did you try with the master version?

Yes. I got same warning indicating that the gelu activation layer will be executed in fp32.

Thanks! I will rework this.

Hey @jlei2 ! Can you let me know if your issue is fixed in this PR please https://github.com/huggingface/transformers/pull/9163

Hi @jplu,

Thanks for creating the PR! Unfortunately the PR doesn't fix the issue. It resolves the 2nd causes I mentioned in this PR description but not the 1st one. It means after applying your PR, the activation layer will always has a compute dtype of fp32, no matter we are in mixed_fp16 mode or pure fp16 mode. As I said for the 1st cause, I think it is because the activation layer is defined outside of the BertIntermediate layer and the computation policy is not broadcasted into it.

For example, under your branch, with the code below:
```
import tensorflow as tf
from transformers import TFBertForQuestionAnswering
tf.keras.mixed_precision.experimental.set_policy('float16')
model = TFBertForQuestionAnswering.from_pretrained('bert-base-uncased')
```

We still got
```
WARNING:tensorflow:Layer activation is casting an input tensor from dtype float16 to the layer's dtype of float32, which is new behavior in TensorFlow 2.  The layer has dtype float32 because its dtype defaults to f
loatx.

If you intended to run this layer in float32, you can safely ignore this warning. If in doubt, this warning is likely only an issue if you are porting a TensorFlow 1.X model to TensorFlow 2.

To change all layers to have dtype float16 by default, call `tf.keras.backend.set_floatx('float16')`. To change just this layer, pass dtype='float16' to the layer constructor. If you are the author of this layer, you can disable autocasting by passing autocast=False to the base Layer constructor.
```
indicating that the activation layer is casted into float32.

The behavior is different from Google's official Bert codebase where the activation layer will be executed as float16. And the two cast ops converting float16 to float32 and then converting it back can slow down the forward computation but provide higher precision.

Did you pull the Last version? I don't get the same behavior than you.

I think so. I am on your branch of fix-activations. Does it require any specific version of TF? I am using tf==2.3.0. Do you mean you don't get the warning info?

I'm using 2.4 and no, I don't have the warning. Try to git pull once again to check.

Just tried with 2.3 as well and I still don't get the warning.

With TF2.4, I don't get the warning because in your PR after tf2.4 we start to use tf.keras.activations.gelu which is introduced in tf2.4. But I still think you should be able to get the warning with tf 2.3, because in that case we are using tf.keras.layers.Activation to create the activation layer. This is weird. Could you double-check you are using tf 2.3 and the execution flow does go through here(https://github.com/huggingface/transformers/blob/4ba0b792767c2dfff61afbf4fd36b8bc521518a1/src/transformers/activations_tf.py#L21)?

I confirm, including with gelu_new

Ok, I know why, I was setting the mixed precision before to import the transformers lib, now if I switch the two I get the warning with 2.3. I will remove the usage of the Activation layer and use only the function then.

I just did the update, now it should works!

Oh good catch! I didn't know the order would make a difference. It now works well on my side!

Awesome! Happy to know that the issue is fixed on your side as well!! The PR should be merged this week :)

Great! I really appreciate your effort on this issue!