Repository: huggingface/transformers

Issue:
tf.function and half precision fails with Roberta models
## Environment info

- `transformers` version: 4.5.1
- Platform: Linux-5.4.0-71-generic-x86_64-with-glibc2.29
- Python version: 3.8.5
- PyTorch version (GPU?): 1.7.1 (False)
- Tensorflow version (GPU?): 2.4.1 (False)
- Using GPU in script?: No (but also fails with GPU)
- Using distributed or parallel set-up in script?: No

### Who can help
As far as I can tell, this worked before #9788, so maybe @jplu can help. Also this is a TF issue so @Rocketknight1 .

## Information

Model I am using (Bert, XLNet ...): TFRoberta, this also happens with TFXLMRoberta

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

```py3
import tensorflow as tf
from transformers.models.roberta import RobertaTokenizerFast, TFRobertaModel


@tf.function
def get_embeddings(
    tokenizer: RobertaTokenizerFast, model: TFRobertaModel, text: str
) -> tf.Tensor:
    return model(**tokenizer(text, return_tensors="tf")).last_hidden_state


if __name__ == "__main__":
    tf.keras.mixed_precision.set_global_policy("float16")
    name = "roberta-base"
    tokenizer = RobertaTokenizerFast.from_pretrained(name)
    model = TFRobertaModel.from_pretrained(name)
    embeddings = get_embeddings(
        tokenizer=tokenizer,
        model=model,
        text="tf.function and mixed precision",
    )
    print(embeddings)
```

Traceback:
```
  File "roberta_bug.py", line 17, in <module>
    embeddings = get_embeddings(
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 828, in __call__
    result = self._call(*args, **kwds)
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 871, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 725, in _initialize
    self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 2969, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3361, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3196, in _create_graph_function
    func_graph_module.func_graph_from_py_func(
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 990, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 634, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 977, in wrapper
    raise e.ag_error_metadata.to_exception(e)
TypeError: in user code:

    roberta_bug.py:9 get_embeddings  *
        return model(**tokenizer(text, return_tensors="tf")).last_hidden_state
    /home/arthur/reinfer/env/lib/python3.8/site-packages/transformers/models/roberta/modeling_tf_roberta.py:744 call  *
        outputs = self.roberta(
    /home/arthur/reinfer/env/lib/python3.8/site-packages/transformers/models/roberta/modeling_tf_roberta.py:544 call  *
        extended_attention_mask = tf.multiply(tf.subtract(1.0, extended_attention_mask), -10000.0)
    /home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper  **
        return target(*args, **kwargs)
    /home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py:561 subtract
        return gen_math_ops.sub(x, y, name)
    /home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/ops/gen_math_ops.py:10316 sub
        _, _, _op, _outputs = _op_def_library._apply_op_helper(
    /home/arthur/reinfer/env/lib/python3.8/site-packages/tensorflow/python/framework/op_def_library.py:555 _apply_op_helper
        raise TypeError(

    TypeError: Input 'y' of 'Sub' Op has type float16 that does not match type float32 of argument 'x'.
```

## Expected behavior

The model should calculate embeddings correctly. This is due to `tf.subtract(1.0, extended_attention_mask)` checking that `1.0` and `extended_attention_mask` have the same type, but in `float16` mode they do not. Reverting to `1.0 - extended_attention_mask` fixes the issue.



Discussion:
Thanks for sharing this issue!

The issue here indeed comes from 1.0 that is not from the same dtype than `extended_attention_mask`. A better fix would be to align this line such as the other models by replacing it with:
```
extended_attention_mask = tf.cast(extended_attention_mask, dtype=embedding_output.dtype)
one_cst = tf.constant(1.0, dtype=embedding_output.dtype)
ten_thousand_cst = tf.constant(-10000.0, dtype=embedding_output.dtype)
extended_attention_mask = tf.multiply(tf.subtract(one_cst, extended_attention_mask), ten_thousand_cst)
```
Extracted from BERT. I will do a fix.

Resolving patch:
diff --git a/src/transformers/models/roberta/modeling_tf_roberta.py b/src/transformers/models/roberta/modeling_tf_roberta.py
--- a/src/transformers/models/roberta/modeling_tf_roberta.py
+++ b/src/transformers/models/roberta/modeling_tf_roberta.py
@@ -541,7 +541,9 @@ 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, dtype=embedding_output.dtype)
-        extended_attention_mask = tf.multiply(tf.subtract(1.0, extended_attention_mask), -10000.0)
+        one_cst = tf.constant(1.0, dtype=embedding_output.dtype)
+        ten_thousand_cst = tf.constant(-10000.0, dtype=embedding_output.dtype)
+        extended_attention_mask = tf.multiply(tf.subtract(one_cst, extended_attention_mask), ten_thousand_cst)
 
         # Prepare head mask if needed
         # 1.0 in head_mask indicate we keep the head


Test patch: