Fix mixed precision for BART

[abheesht17]

With

policy = keras.mixed_precision.Policy("mixed_float16")
keras.mixed_precision.set_global_policy(policy)

,

this error shows up:

TypeError: Exception encountered when calling layer "tf.linalg.matmul" (type TFOpLambda).

Input 'y' of 'BatchMatMulV2' Op has type float32 that does not match type float16 of argument 'x'.

Call arguments received by layer "tf.linalg.matmul" (type TFOpLambda):
  • a=tf.Tensor(shape=(None, None, 768), dtype=float16)
  • b=<tf.Variable 'token_embedding/embeddings:0' shape=(50265, 768) dtype=float32>
  • transpose_a=False
  • transpose_b=True
  • adjoint_a=False
  • adjoint_b=False
  • a_is_sparse=False
  • b_is_sparse=False
  • output_type=None
  • name=None

[mattdangerw]

I think we actually want this the other way. We should cast the variable to the compute dtype (the lower precision type when using mixed), before multiplying with x. So tf.cast(backbone.token_embedding.embeddings, x.dtype).

Can you check if that fixes things as well?

[abheesht17]

Hmmm, why does GPT-2 output tf.float32 instead of tf.float16, then?

[mattdangerw]

In what context does it? There might be some casting going on to the outputs.

My understanding is that we should be casting our variables to the compute_dtype. Take a look at
https://github.com/keras-team/keras/blob/v2.12.0/keras/engine/base_layer.py#L2216-L2235
and https://github.com/keras-team/keras/blob/v2.12.0/keras/mixed_precision/autocast_variable.py

But maybe I am missing something!

[abheesht17]

GPT2 backbone outputs tf.float32. BART backbone outputs tf.float16. Weird why that's happening.

https://colab.research.google.com/drive/18AbKIwbUAtJySgAYWXa0ggvRbp5TriEa?usp=sharing

[mattdangerw]

This actually may no longer be necessary with the port. We had to wrap things in a Layer to avoid some errors with keras-core, which might mean that variables are now autocasting automatically (I think it should).

Can you try again on the latest from master?

[mattdangerw]

We think this is no longer relevant with the ReverseEmbedding layer. Closing, can reopen if more issue.