Repository: huggingface/transformers

Issue:
TF BERT not FP16 compatible?
# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...): TFBertForQuestionAnswering

Language I am using the model on (English, Chinese ...): English

The problem arises when using:
* [x] my own modified scripts: 

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: SQUaD

## To reproduce
Simple example to reproduce error:
```
import tensorflow as tf
from transformers import TFBertForQuestionAnswering

# turn on mp (fp16 operations)
tf.keras.mixed_precision.experimental.set_policy('mixed_float16')

model = TFBertForQuestionAnswering.from_pretrained('bert-base-uncased')
```

The error occurs here: 
transformers/modeling_tf_bert.py", line 174, in _embedding
    embeddings = inputs_embeds + position_embeddings + token_type_embeddings

And this is the error:
tensorflow.python.framework.errors_impl.InvalidArgumentError: cannot compute AddV2 as input #1(zero-based) was expected to be a half tensor but is a float tensor [Op:AddV2] name: tf_bert_for_question_answering/bert/embeddings/add/

## Expected behavior
I want to use TF BERT with mixed precision (for faster inference on tensor core GPUs). I know that full fp16 is not working out-of-the-box, because the model weights need to be in fp16 as well. Mixed precision, however, should work because only operations are performed in fp16.

I get some dtype issue. Seems the mode is not fp16 compatible yet? Will this be fixed in the future?

## Environment info
- `transformers` version: 2.5.0
- Platform:  ubuntu 16.04
- Python version: 3.6.9
- PyTorch version (GPU?): 1.4.0 (GPU)
- Tensorflow version (GPU?): 2.1.0 (GPU)
- Using GPU in script?: sort of
- Using distributed or parallel set-up in script?: nope



Discussion:
I've aced same issue. Maybe it's hard coded the data type somewhere? Have you found solution?
Tried this on Colab TPU, same error.
Same here, would be convenient as hell :)
Having the same error also for `transformers` version 2.11.0. 
Here some code to easily reproduce the error:

```python 
#!/usr/bin/env python3
from transformers import TFBertModel, BertTokenizer
from tensorflow.keras.mixed_precision import experimental as mixed_precision

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

tok = BertTokenizer.from_pretrained("bert-base-uncased")
model = TFBertModel.from_pretrained("bert-base-uncased")
input_ids = tok("The dog is cute", return_tensors="tf").input_ids
model(input_ids)  # throws error on GPU
```
Encountering the same issue here:
```python3
import tensorflow as tf
from transformers.modeling_tf_distilbert import TFDistilBertModel

tf.keras.mixed_precision.experimental.set_policy('mixed_float16')
model = TFDistilBertModel.from_pretrained('distilbert-base-uncased')
Put this issue on my TF ToDo-List :-) 
+1
Hi @patrickvonplaten, is this problem fixed？
I got the same error recently with version 3.0.2
This is still an open problem...I didn't find the time yet to take a look! Will link this issue to the TF projects.

Resolving patch:
diff --git a/src/transformers/modeling_tf_bert.py b/src/transformers/modeling_tf_bert.py
--- a/src/transformers/modeling_tf_bert.py
+++ b/src/transformers/modeling_tf_bert.py
@@ -215,8 +215,8 @@ def _embedding(self, input_ids, position_ids, token_type_ids, inputs_embeds, tra
         if inputs_embeds is None:
             inputs_embeds = tf.gather(self.word_embeddings, input_ids)
 
-        position_embeddings = self.position_embeddings(position_ids)
-        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+        position_embeddings = tf.cast(self.position_embeddings(position_ids), inputs_embeds.dtype)
+        token_type_embeddings = tf.cast(self.token_type_embeddings(token_type_ids), inputs_embeds.dtype)
         embeddings = inputs_embeds + position_embeddings + token_type_embeddings
         embeddings = self.LayerNorm(embeddings)
         embeddings = self.dropout(embeddings, training=training)
@@ -281,7 +281,7 @@ def call(self, hidden_states, attention_mask, head_mask, output_attentions, trai
         attention_scores = tf.matmul(
             query_layer, key_layer, transpose_b=True
         )  # (batch size, num_heads, seq_len_q, seq_len_k)
-        dk = tf.cast(shape_list(key_layer)[-1], tf.float32)  # scale attention_scores
+        dk = tf.cast(shape_list(key_layer)[-1], attention_scores.dtype)  # scale attention_scores
         attention_scores = attention_scores / tf.math.sqrt(dk)
 
         if attention_mask is not None:
@@ -613,6 +613,8 @@ def call(
         if token_type_ids is None:
             token_type_ids = tf.fill(input_shape, 0)
 
+        embedding_output = self.embeddings(input_ids, position_ids, token_type_ids, inputs_embeds, training=training)
+
         # We create a 3D attention mask from a 2D tensor mask.
         # Sizes are [batch_size, 1, 1, to_seq_length]
         # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
@@ -626,7 +628,7 @@ def call(
         # Since we are adding it to the raw scores before the softmax, this is
         # effectively the same as removing these entirely.
 
-        extended_attention_mask = tf.cast(extended_attention_mask, tf.float32)
+        extended_attention_mask = tf.cast(extended_attention_mask, embedding_output.dtype)
         extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
 
         # Prepare head mask if needed
@@ -640,7 +642,6 @@ def call(
             head_mask = [None] * self.num_hidden_layers
             # head_mask = tf.constant([0] * self.num_hidden_layers)
 
-        embedding_output = self.embeddings(input_ids, position_ids, token_type_ids, inputs_embeds, training=training)
         encoder_outputs = self.encoder(
             embedding_output,
             extended_attention_mask,
diff --git a/src/transformers/modeling_tf_electra.py b/src/transformers/modeling_tf_electra.py
--- a/src/transformers/modeling_tf_electra.py
+++ b/src/transformers/modeling_tf_electra.py
@@ -134,8 +134,8 @@ def _embedding(self, input_ids, position_ids, token_type_ids, inputs_embeds, tra
 
         if inputs_embeds is None:
             inputs_embeds = tf.gather(self.word_embeddings, input_ids)
-        position_embeddings = self.position_embeddings(position_ids)
-        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+        position_embeddings = tf.cast(self.position_embeddings(position_ids), inputs_embeds.dtype)
+        token_type_embeddings = tf.cast(self.token_type_embeddings(token_type_ids), inputs_embeds.dtype)
 
         embeddings = inputs_embeds + position_embeddings + token_type_embeddings
         embeddings = self.LayerNorm(embeddings)
@@ -194,7 +194,7 @@ class TFElectraPreTrainedModel(TFBertPreTrainedModel):
     config_class = ElectraConfig
     base_model_prefix = "electra"
 
-    def get_extended_attention_mask(self, attention_mask, input_shape):
+    def get_extended_attention_mask(self, attention_mask, input_shape, dtype):
         if attention_mask is None:
             attention_mask = tf.fill(input_shape, 1)
 
@@ -211,7 +211,7 @@ def get_extended_attention_mask(self, attention_mask, input_shape):
         # Since we are adding it to the raw scores before the softmax, this is
         # effectively the same as removing these entirely.
 
-        extended_attention_mask = tf.cast(extended_attention_mask, tf.float32)
+        extended_attention_mask = tf.cast(extended_attention_mask, dtype)
         extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
 
         return extended_attention_mask
@@ -314,11 +314,11 @@ def call(
         if token_type_ids is None:
             token_type_ids = tf.fill(input_shape, 0)
 
-        extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape)
-        head_mask = self.get_head_mask(head_mask)
-
         hidden_states = self.embeddings(input_ids, position_ids, token_type_ids, inputs_embeds, training=training)
 
+        extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape, hidden_states.dtype)
+        head_mask = self.get_head_mask(head_mask)
+
         if hasattr(self, "embeddings_project"):
             hidden_states = self.embeddings_project(hidden_states, training=training)
 


Test patch: