TFGPT2LMHeadModel fp16 support
Author: mymusise
Pull request: no

## Environment info

- `transformers` version:
- Platform: ubuntu 18.04
- Python version: python3.8
- Tensorflow version (GPU?): tf-nightly==2.5
- Using GPU in script?: Y
- Using distributed or parallel set-up in script?: N

### Who can help

 albert, bert, GPT2, XLM: @LysandreJik 
 Text Generation: @patrickvonplaten @TevenLeScao
 tensorflow: @jplu

## Information

Hi, there. If I want to use the mixed precision setting with keras apis when training `TFGPT2LMHeadModel`, [like this](https://www.tensorflow.org/guide/mixed_precision):
```python
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)
```
Then I will got this error:
```
  File "/home/mymusise/pro/fast-gpt2/env/lib/python3.8/site-packages/transformers/modeling_tf_gpt2.py", line 154, in call
    attn_outputs = self._attn(query, key, value, attention_mask, head_mask, output_attentions, training=training)
  File "/home/mymusise/pro/fast-gpt2/env/lib/python3.8/site-packages/transformers/modeling_tf_gpt2.py", line 101, in _attn
    w = w / tf.math.sqrt(dk)
  File "/home/mymusise/pro/fast-gpt2/env/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py", line 1181, in binary_op_wrapper
    raise e
  File "/home/mymusise/pro/fast-gpt2/env/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py", line 1165, in binary_op_wrapper
    return func(x, y, name=name)
  File "/home/mymusise/pro/fast-gpt2/env/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py", line 206, in wrapper
    return target(*args, **kwargs)
  File "/home/mymusise/pro/fast-gpt2/env/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py", line 1337, in truediv
    return _truediv_python3(x, y, name)
  File "/home/mymusise/pro/fast-gpt2/env/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py", line 1267, 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
```
Here's a [example](https://gist.github.com/mymusise/7192b7c252ff67ff84496cd8b27a91ff) to reappear this.

Please help me guys.

Comments:
Hello !

It is a known issue that will be fixed in a future release. Sorry.

Thanks, I'm looking forward to the future release.

Thank @jplu's work! :+1:  It works now with the mixed-precision policy when training.
But I think it still has some problem with `TextGenerationPipeline`, for example:

```python
from transformers import TextGenerationPipeline
from tensorflow.keras.mixed_precision import experimental as mixed_precision

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

text_generator = TextGenerationPipeline(model, tokenizer)
text_generator("Transformers: State-of-the-art Natural Language Processing for Pytorch and TensorFlow")
```

Then it raises an exception:
```
  File "./env/lib/python3.8/site-packages/transformers-4.0.0.dev0-py3.8.egg/transformers/generation_tf_utils.py", line 386, in generate
    output = self._generate_no_beam_search(
  File "./env/lib/python3.8/site-packages/transformers-4.0.0.dev0-py3.8.egg/transformers/generation_tf_utils.py", line 457, in _generate_no_beam_search
    next_token_logits = tf.math.multiply(next_token_logits, next_token_logits_penalties)
  File "./env/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py", line 206, in wrapper
    return target(*args, **kwargs)
  File "./env/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py", line 519, in multiply
    return gen_math_ops.mul(x, y, name)
  File "./env/lib/python3.8/site-packages/tensorflow/python/ops/gen_math_ops.py", line 6068, in mul
    _ops.raise_from_not_ok_status(e, name)
  File "./env/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 6867, in raise_from_not_ok_status
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: cannot compute Mul as input #1(zero-based) was expected to be a half tensor but is a float tensor [Op:Mul]
```

Sorry for this, the generation part is not yet compliant with mixed precision. It is in the pipeline to add this but we don't know when yet.

Okay, thank you!

Hello~ @jplu. 
Recently, I try to train my `TFGPT2LMHeadModel` model with mixed_precision again, it performs badly after many epochs, it seems didn't learn anything.
If I train without mixed_precision, it performs well after training with the same epochs.
I think maybe it will lose something when counting the loss with `logits` and `labels` in `fp16` here: https://github.com/mymusise/transformers/blob/master/src/transformers/models/gpt2/modeling_tf_gpt2.py#L742.

Here I make some change to `TFGPT2LMHeadModel` in https://github.com/huggingface/transformers/pull/10689
I'm not sure it's the right way to do it, correct me if it's wrong.

And I make a small test in colab, hope it can help to recur the problem.
https://colab.research.google.com/github/mymusise/gpt2-quickly/blob/main/examples/mixed_precision_test.ipynb

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.