Mixed precision errors for GPT2 and other models

[mattdangerw]

Already have a pull request, just opening this to track! #958

[ADITYADAS1999]

Already have a pull request, just opening this to track! #958

hey @mattdangerw can you explain this a bit !

[mattdangerw]

Sure thing! Basically Keras supports a mixed precision mode that is a substantial boost in performance for GPU training. https://www.tensorflow.org/guide/mixed_precision

But supporting it requires making sure our tensorflow ops do not assume an float32 computation type (we will use 16 bit floating point numbers instead). This PR that landed just removed a few places we were making too many assumptions about dtype, and breaking the "mixed precision" workflow.