Transformers need keras updates for tf-nightly + tf_keras (Keras 2.0) & keras 3.0

[jojivk]

System Info

With new changes for Keras (https://groups.google.com/g/keras-users/c/jgGcX730WGE), we are seeing issues with using Huggingface models (GPT-J, Stable Diffusion, & Vision Transformer). We saw that transformers restrict the versions for Tensorflow and Keras. https://github.com/huggingface/transformers/blob/main/setup.py#L129.

To use the models with tf-keras (keras 2.0), we could resolve issues in Transformers by changing import of keras in transformers to import tf_keras. -ie change
import keras.*
to
import tf_keras.*

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

1. Install tf-nightly
2. Install tf-keras-nightly
3. export TF_USE_LEGACY_KERAS=1
4. Run Huggingface Tensorflow GPT-J

Expected behavior

Fails to run

[amyeroberts]

Hi @jojivk, thanks for raising this issue!

We don't officially support nightly releases as they're not guaranteed to be stable. Our versions of tensorflow and keras are pinned because automatically installing the latest release tends to lead to many breaking changes. If there is a change in the import structure in a stable release, then it's something we could try to handle and of course help from the community on this is always welcome :)

cc @Rocketknight1

[jojivk]

Hi @amyeroberts Thanks for the quick response. Keras 3.0 will be default (default now in nightly) with TF 2.16 release. Also, to fall back to using keras 2.0 with TF 2.16, few fixes that I mentioned, is needed.
Best

[amyeroberts]

@jojivk OK, thanks for pointing out. There's a draft PR #27375 which you can track which is handling updating the supported TF versions

[Rocketknight1]

Hi @jojivk, yes, we're aware of the issues! We're going to stick to the version pin for now until 2.16 is more stable and we can test it properly.

Our long-term plan here is:

1. Start by pinning the version of TF so we don't break backward compatibility
2. Move and update imports so we can support TF>=2.16, although we still won't support Keras 3 at this point
3. Rewrite our core modelling code to support Keras 3 (TF only for now)
4. Finally, hopefully (!) add the ability to run our Keras 3 models in other frameworks, probably starting with JAX.

[jojivk]

@Rocketknight1 Thanks for the update

[Rocketknight1]

Quick update for this issue, we've filed our first big PR to prepare for Keras 3 / keras-nightly at #27794. There's likely more work to be done after this, but this resolves one major blocker.

[jojivk73]

@Rocketknight1 . Thanks for the update and fixes. I will update on any issues we see.

[jojivk]

@Rocketknight1 @amyeroberts. Thank you for all the fixes!
I see the following imports for keras which cause backward compatibility issues.

from keras import backend as K

loc1

loc2

We fixed it locally (as below) to use with Keras2.0 Models, due to reasons mentioned above in the issue report.
from tf_keras import backend as K

Other similar imports: which as of now we are not sure if it is an issue.

from tensorflow.keras.callbacks import Callback
./src/transformers/keras_callbacks.py
from tensorflow.keras.layers import Dense
./src/transformers/models/vision_text_dual_encoder/modeling_tf_vision_text_dual_encoder.py
from tensorflow.keras.activations import gelu
from tensorflow.keras.layers import Dense, Dropout, Embedding, Layer, LayerNormalization
./src/transformers/models/esm/modeling_tf_esm.py
from tensorflow.keras.preprocessing import image
./src/transformers/models/efficientnet/convert_efficientnet_to_pytorch.py
from tensorflow.keras.optimizers.legacy import Adam
from tensorflow.keras.optimizers import Adam
./src/transformers/optimization_tf.py

[Rocketknight1]

Hi @jojivk - yep, we noted those too! I think we're going to remove all instances of keras.backend in the TF files, and replace them with equivalent TF functions. We'll also probably replace imports of keras with from tensorflow import keras during the transition period - this should avoid issues where the Keras version differs significantly from the installed TF version.

[jojivk]

@Rocketknight1 . Thank you for the update!.

[fchollet]

We'll also probably replace imports of keras with from tensorflow import keras during the transition period

Prefer using import tf_keras as keras to make sure you're getting Keras 2.

[Rocketknight1]

Thanks @fchollet! We'll use that pattern instead.

[jojivk73]

@Rocketknight1
Thank you for the fixes. Have a few questions.
The fix (#28588) I understand supports transitioning to tf_keras. Does it also support Keras 3.0 for customers who want to move to Keras 3.0?
Will the information be added to your next release notes?
Also let me know if there is any additional setting(s) needed to with keep using tf_keras 2.0.?

[Rocketknight1]

You shouldn't need to change any settings @jojivk73. If you install tf_keras, everything should work the same as it did before.

The full transition to Keras 3 will be trickier. The reason for this is that all of our model code is written for TF/Keras, and mixes Keras layers with TF ops. To work in Keras 3, this code will need to replace all the TF ops with keras.ops. This will break compatibility with Keras 2, so we can't cleanly update the files in place.

Our two options are:

1. Wait some time (~1 year), then drop support for older versions of TensorFlow and begin to transition our TF codebase from Keras 2 + TF ops to Keras 3 + keras.ops code.
2. Add "Keras 3" as a fourth framework alongside TF, PyTorch and JAX/Flax, so we can start supporting Keras 3 before we deprecate TF + Keras 2.

Either way, we'll need a big community push to port all the models to keras.ops! Right now, we're being cautious, but we'll probably choose option 2 and accelerate a lot if we feel like a lot of users want Keras 3.

[jojivk73]

@Rocketknight1 Thanks for the quick and detailed reply.