You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Describe the bug
Training BERT using Keras NLP is significantly slower due to the keras.layers.Embedding not being XLA compatible by default on TensorFlow GPU. This is similar to an issue reported for Keras at keras-team/keras#19809.
To Reproduce
You can reproduce this issue by following the steps in this Colab Notebook: Link to Notebook
Expected behavior
I expect BERT training using Keras NLP on TensorFlow GPU with XLA to be optimized for performance, similar to native TensorFlow implementations.
Additional context
The lack of XLA compatibility affects the training speed and efficiency on GPU, crucial for model training scalability and practical application in production environments.
Would you like to help us fix it?
Yes, I am willing to contribute to resolving this issue by testing and suggesting implementations that ensure XLA compatibility.
The text was updated successfully, but these errors were encountered:
I think this could work, but I am not sure we would want to do it when XLA is off by default on the TF backend. So the first thing might be to look at enabling XLA with tf.distribute.
Is there a reason training on the Jax backend doesn't work for your use case? It is likely faster, and everything is XLA compatible as there is no other option on Jax.
Describe the bug
Training BERT using Keras NLP is significantly slower due to the
keras.layers.Embedding
not being XLA compatible by default on TensorFlow GPU. This is similar to an issue reported for Keras at keras-team/keras#19809.To Reproduce
You can reproduce this issue by following the steps in this Colab Notebook: Link to Notebook
Expected behavior
I expect BERT training using Keras NLP on TensorFlow GPU with XLA to be optimized for performance, similar to native TensorFlow implementations.
Additional context
The lack of XLA compatibility affects the training speed and efficiency on GPU, crucial for model training scalability and practical application in production environments.
See also:
Would you like to help us fix it?
Yes, I am willing to contribute to resolving this issue by testing and suggesting implementations that ensure XLA compatibility.
The text was updated successfully, but these errors were encountered: