Repository: huggingface/transformers

Issue:
Remove hard-coded uses of float32 to fix mixed precision use in Distilbert
In this commit [Remove hard-coded uses of float32 to fix mixed precision use](https://github.com/huggingface/transformers/commit/4fca874ea995f3d23ad7062b07b5ed7c4f87c0cd#diff-e3ab4f29f29fe1d243a6b55fafaab097), the mixed precision issue is fixed for [modeling_tf_bert.py](https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_tf_bert.py).

However, for [modeling_tf_distilbert.py](https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_tf_distilbert.py#L171), the line 171 is not fixed yet, and we get 

>    173         embeddings = inputs_embeds + position_embeddings  # (bs, max_seq_length, dim)
>    --> 174         embeddings = self.LayerNorm(embeddings)  # (bs, max_seq_length, dim)
>    InvalidArgumentError: cannot compute AddV2 as input #1(zero-based) was expected to be a bfloat16 tensor but is a float tensor

while using `mixed_bfloat16 mixed precision` with `TPU`.

A very quick fix is the same as the fix for `modeling_tf_bert.py`:

    position_embeddings = tf.cast(self.position_embeddings(position_ids), inputs_embeds.dtype)

@schmidek 


Discussion:
Indeed! Do you want to open a PR to fix this?
@LysandreJik 

I can do that. However @patrickvonplaten  has already self-assigned for this. How do you think, @patrickvonplaten?
Hey @chiapas, it would be great if you can open a PR for it :-) 
Hi @patrickvonplaten , OK, that would be my first contribution to transformers :)

Resolving patch:
diff --git a/src/transformers/modeling_tf_distilbert.py b/src/transformers/modeling_tf_distilbert.py
--- a/src/transformers/modeling_tf_distilbert.py
+++ b/src/transformers/modeling_tf_distilbert.py
@@ -76,7 +76,7 @@ def gelu(x):
         0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
         Also see https://arxiv.org/abs/1606.08415
     """
-    cdf = 0.5 * (1.0 + tf.math.erf(x / tf.math.sqrt(2.0)))
+    cdf = 0.5 * (1.0 + tf.math.erf(x / tf.cast(tf.math.sqrt(2.0), dtype=x.dtype)))
     return x * cdf
 
 
@@ -168,7 +168,9 @@ def _embedding(self, input_ids, position_ids, inputs_embeds, training=False):
 
         if inputs_embeds is None:
             inputs_embeds = tf.gather(self.word_embeddings, input_ids)
-        position_embeddings = self.position_embeddings(position_ids)  # (bs, max_seq_length, dim)
+        position_embeddings = tf.cast(
+            self.position_embeddings(position_ids), inputs_embeds.dtype
+        )  # (bs, max_seq_length, dim)
 
         embeddings = inputs_embeds + position_embeddings  # (bs, max_seq_length, dim)
         embeddings = self.LayerNorm(embeddings)  # (bs, max_seq_length, dim)
@@ -261,9 +263,12 @@ def unshape(x):
         scores = tf.matmul(q, k, transpose_b=True)  # (bs, n_heads, q_length, k_length)
         mask = tf.reshape(mask, mask_reshape)  # (bs, n_heads, qlen, klen)
         # scores.masked_fill_(mask, -float('inf'))            # (bs, n_heads, q_length, k_length)
-        scores = scores - 1e30 * (1.0 - mask)
 
-        weights = tf.nn.softmax(scores, axis=-1)  # (bs, n_heads, qlen, klen)
+        scores_dtype = scores.dtype
+        # calculate `scores` in `tf.float32` to avoid numeric overflow
+        scores = tf.cast(scores, dtype=tf.float32) - 1e30 * (1.0 - tf.cast(mask, dtype=tf.float32))
+
+        weights = tf.cast(tf.nn.softmax(scores, axis=-1), dtype=scores_dtype)  # (bs, n_heads, qlen, klen)
         weights = self.dropout(weights, training=training)  # (bs, n_heads, qlen, klen)
 
         # Mask heads if we want to


Test patch: