TFGPT2LMHeadModel fp16 support

[mymusise]

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:

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 to reappear this.

Please help me guys.

[jplu]

Hello !

It is a known issue that will be fixed in a future release. Sorry.

[mymusise]

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

[mymusise]

Thank @jplu's work! 👍 It works now with the mixed-precision policy when training.
But I think it still has some problem with TextGenerationPipeline, for example:

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]

[jplu]

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.

[mymusise]

Okay, thank you!

[mymusise]

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.

[mymusise]

Here I make some change to TFGPT2LMHeadModel in #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

[github-actions]

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 are likely to be ignored.