Fix mixed precision issue in TF DistilBert
Author: ydshieh
Pull request: yes

Fix mixed precision issue in TF DistilBert by removing hard-coded uses of float32.

<!-- This line specifies which issue to close after the pull request is merged. -->
Fixes #6858


Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/6915?src=pr&el=h1) Report
> Merging [#6915](https://codecov.io/gh/huggingface/transformers/pull/6915?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/4ebb52afdb4dc4bcd599e7cb503763e5d4afc962?el=desc) will **increase** coverage by `2.01%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/6915/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/6915?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #6915      +/-   ##
==========================================
+ Coverage   77.81%   79.83%   +2.01%     
==========================================
  Files         157      157              
  Lines       28853    28853              
==========================================
+ Hits        22452    23034     +582     
+ Misses       6401     5819     -582     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/6915?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/modeling\_tf\_distilbert.py](https://codecov.io/gh/huggingface/transformers/pull/6915/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9kaXN0aWxiZXJ0LnB5) | `98.82% <100.00%> (ø)` | |
| [src/transformers/modeling\_tf\_electra.py](https://codecov.io/gh/huggingface/transformers/pull/6915/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9lbGVjdHJhLnB5) | `25.13% <0.00%> (-73.83%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_openai.py](https://codecov.io/gh/huggingface/transformers/pull/6915/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9vcGVuYWkucHk=) | `22.58% <0.00%> (-72.26%)` | :arrow_down: |
| [src/transformers/tokenization\_marian.py](https://codecov.io/gh/huggingface/transformers/pull/6915/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fbWFyaWFuLnB5) | `32.20% <0.00%> (-66.95%)` | :arrow_down: |
| [src/transformers/tokenization\_pegasus.py](https://codecov.io/gh/huggingface/transformers/pull/6915/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fcGVnYXN1cy5weQ==) | `46.03% <0.00%> (-49.21%)` | :arrow_down: |
| [src/transformers/modeling\_marian.py](https://codecov.io/gh/huggingface/transformers/pull/6915/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19tYXJpYW4ucHk=) | `60.00% <0.00%> (-30.00%)` | :arrow_down: |
| [src/transformers/modeling\_t5.py](https://codecov.io/gh/huggingface/transformers/pull/6915/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190NS5weQ==) | `76.70% <0.00%> (-6.07%)` | :arrow_down: |
| [src/transformers/activations.py](https://codecov.io/gh/huggingface/transformers/pull/6915/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9hY3RpdmF0aW9ucy5weQ==) | `85.00% <0.00%> (-5.00%)` | :arrow_down: |
| [src/transformers/tokenization\_auto.py](https://codecov.io/gh/huggingface/transformers/pull/6915/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fYXV0by5weQ==) | `95.55% <0.00%> (-2.23%)` | :arrow_down: |
| [src/transformers/modeling\_openai.py](https://codecov.io/gh/huggingface/transformers/pull/6915/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19vcGVuYWkucHk=) | `80.96% <0.00%> (-1.30%)` | :arrow_down: |
| ... and [23 more](https://codecov.io/gh/huggingface/transformers/pull/6915/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/6915?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/6915?src=pr&el=footer). Last update [4ebb52a...481baa3](https://codecov.io/gh/huggingface/transformers/pull/6915?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


I run a test with this change on my ubuntu 18.04 with a 2080Ti GPU, tensorflow-gpu 2.2.0:
```
from tensorflow.keras.layers import Input, Embedding, Bidirectional, GRU, Dense
from tensorflow.keras.models import Model
from transformers import TFDistilBertModel
from tensorflow.keras.mixed_precision import experimental as mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

bert = TFDistilBertModel.from_pretrained('distilbert-base-uncased')
inputs = Input(shape=(None,), dtype='int32')
bert_out = bert(inputs)[0]
output = Dense(9, activation='softmax', dtype='float32')(bert_out)
model = Model(inputs, output)
model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
model.summary()
x = [[5, 2, 3] * 3] * 100
y = [[1, 2, 3] * 3] * 100
model.fit(x=x, y=y, epochs=20, batch_size=16)
```
And get error info:
```
Traceback (most recent call last):
  File "test.py", line 8, in <module>
    bert = TFDistilBertModel.from_pretrained('distilbert-base-uncased')
  File "/home/xingya/transformers/src/transformers/modeling_tf_utils.py", line 602, in from_pretrained
    model(model.dummy_inputs, training=False)  # build the network with dummy inputs
  File "/home/xingya/.conda/envs/transformers/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py", line 968, in __call__
    outputs = self.call(cast_inputs, *args, **kwargs)
  File "/home/xingya/transformers/src/transformers/modeling_tf_distilbert.py", line 615, in call
    outputs = self.distilbert(inputs, **kwargs)
  File "/home/xingya/.conda/envs/transformers/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py", line 968, in __call__
    outputs = self.call(cast_inputs, *args, **kwargs)
  File "/home/xingya/transformers/src/transformers/modeling_tf_distilbert.py", line 508, in call
    tfmr_output = self.transformer(
  File "/home/xingya/.conda/envs/transformers/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py", line 968, in __call__
    outputs = self.call(cast_inputs, *args, **kwargs)
  File "/home/xingya/transformers/src/transformers/modeling_tf_distilbert.py", line 401, in call
    layer_outputs = layer_module(hidden_state, attn_mask, head_mask[i], output_attentions, training=training)
  File "/home/xingya/.conda/envs/transformers/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py", line 968, in __call__
    outputs = self.call(cast_inputs, *args, **kwargs)
  File "/home/xingya/transformers/src/transformers/modeling_tf_distilbert.py", line 355, in call
    ffn_output = self.ffn(sa_output, training=training)  # (bs, seq_length, dim)
  File "/home/xingya/.conda/envs/transformers/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py", line 968, in __call__
    outputs = self.call(cast_inputs, *args, **kwargs)
  File "/home/xingya/transformers/src/transformers/modeling_tf_distilbert.py", line 304, in call
    x = self.activation(x)
  File "/home/xingya/.conda/envs/transformers/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py", line 968, in __call__
    outputs = self.call(cast_inputs, *args, **kwargs)
  File "/home/xingya/.conda/envs/transformers/lib/python3.8/site-packages/tensorflow/python/keras/layers/core.py", line 420, in call
    return self.activation(inputs)
  File "/home/xingya/transformers/src/transformers/modeling_tf_distilbert.py", line 79, in gelu
    cdf = 0.5 * (1.0 + tf.math.erf(x / tf.math.sqrt(2.0)))
  File "/home/xingya/.conda/envs/transformers/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py", line 984, in binary_op_wrapper
    return func(x, y, name=name)
  File "/home/xingya/.conda/envs/transformers/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py", line 1081, in _truediv_python3
    raise TypeError("x and y must have the same dtype, got %r != %r" %
TypeError: x and y must have the same dtype, got tf.float16 != tf.float32
```
I made a modification to L299:
`self.activation = (
            tf.keras.layers.Activation(gelu, dtype='float32') if config.activation == "gelu" else tf.keras.activations.relu
        )`
And then the model began to train, however the loss don't decrease and the accuracy is always 0:
```
7/7 [==============================] - 0s 28ms/step - loss: 2.1972 - accuracy: 0.0000e+00
Epoch 2/20
7/7 [==============================] - 0s 29ms/step - loss: 2.1972 - accuracy: 0.0000e+00
Epoch 3/20
7/7 [==============================] - 0s 30ms/step - loss: 2.1972 - accuracy: 0.0000e+00
Epoch 4/20
7/7 [==============================] - 0s 31ms/step - loss: 2.1972 - accuracy: 0.0000e+00
```

I have trid this code in float32 precision, and it works. 
```
Epoch 1/20
7/7 [==============================] - 0s 31ms/step - loss: 2.5418 - accuracy: 0.2800
Epoch 2/20
7/7 [==============================] - 0s 33ms/step - loss: 1.2452 - accuracy: 0.3356
Epoch 3/20
7/7 [==============================] - 0s 31ms/step - loss: 1.1438 - accuracy: 0.3267
Epoch 4/20
7/7 [==============================] - 0s 33ms/step - loss: 1.1219 - accuracy: 0.3400
```

@xuxingya , the accuracy not improved during training is due to a line 

    > scores = scores - 1e30 * (1.0 - mask)

while `1e30` with `half precision` will cause `nan` values. I am still trying to figure out a way to deal with it.

@xuxingya Would you mind to run the test on your side again, please? I tested it with your example, and it is fine now.

@chiapas Yes, I run the test and now it's fine.