Add tf_keras imports to prepare for Keras 3

[Rocketknight1]

Keras 3 will break backward compatibility for our TF code, and is becoming the default Keras in TF 2.16. This PR uses the tf_keras package to maintain backward compatibility - it imports tf_keras if available, and if not then it attempts to import keras, but raises an issue if the version is >= 3.

Our future plan is to ensure that TF code remains backward compatible, but to support Keras 3 in all its framework-independent glory with new Keras classes (e.g. TFBertModel -> KerasBertModel). The PR for this is at #26224, but it's on hold until we handle the urgent issue of preserving backward compatibility. It was also blocked by the need for a couple of other PRs, but those are mostly in now. Because the full Keras 3 PR will require TF models to be rewritten with 100% Keras ops instead of TF ones, we'll likely need to do a community push once the core modelling code is ready to port everything.

cc @fchollet

Fixes #27377
Fixes #28296

TODO:

• Replace keras or tf.keras with tf_keras falling back to keras in core modelling code
• Replace keras or tf.keras with tf_keras falling back to keras in all model files
• Confirm versions - should we reject Keras >= 3, or Keras >= 2.16?

[Rocketknight1]

Quick note to myself: The failing test is likely caused by replacing K.set_value() with tf.assign() somewhere, since tf.assign() lacks an implicit dtype cast. The test fails whether or not other tests are run beforehand, so it's not a case of the mixed precision policy accidentally carrying over between tests (which happened to us before)

[Rocketknight1]

Quick note for @amyeroberts: Although there are a lot of changes in this PR, the core idea is simple, we import tf_keras as keras if possible, and if not we fall back to import keras and raise an error if the version is >= 3. To avoid copying the boilerplate all over the modeling files, I just do it once in modeling_tf_utils and then from modeling_tf_utils import keras in all files, but that's probably bad - let me know if you want me to just copy the import try/except block around instead!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[amyeroberts]

Thanks for all these updates! It is indeed a hefty PR 😅

A few comments / question but overall looks good!

[amyeroberts]

Is .assign backwards compatible?

[Rocketknight1]

Yes - .assign() is a TF Tensor method that's been there since 2.0.

[amyeroberts]

Ah, sorry, what I mean is this guaranteed to be equivalent? (I'm assuming yes, but I'm not very familiar with keras.backend)

[Rocketknight1]

The only difference is that .assign() doesn't cast different dtypes together automatically. The dtypes should always be float32 here, but I'll add a cast for safety (also this code is ancient and looks like TF1, so I'm not sure if it runs at all on modern TF anyway)

[amyeroberts]

I'm not sure I understand the rational for when the version compatible import is just handled e.g. in src/transformers/keras_callbacks.py and when the version is checked e.g. here

[Rocketknight1]

I've rewritten things to always import from modeling_tf_utils. The only exceptions, where I copy-pasted it instead, are here in activations_tf, because that would create a circular import, and in the example files, which are designed for users to read and modify so I want to make it clear exactly which keras they're getting.

[amyeroberts]

Perfect :)

[amyeroberts]

I'm slightly concerned that raising this exception at the top-level of a module is going to cause lots of issues. Users, who might not even be using tensorflow with transformers, but will have it in their environment will start having exception raised if they do e.g. from transformers import * (or something less stupid which ends up loading this module).

[Rocketknight1]

Yeah, I was worried about this too! I think that specific problem shouldn't happen, because if TF isn't available then all the TF objects will be dummies and TF-specific files shouldn't be executed.

I'm still a bit unsure about the import from modeling_tf_utils, though - it feels dubious to me, but it was the only way I could think of to avoid pasting boilerplate everywhere. Maybe for peace of mind it'd be better to just do that, though, which would resolve this issue (and most of the other comments here)

[amyeroberts]

We can just install the different possibilities: no TF, old keras, new keras and check to see if the exceptions / warnings are egregious. I suspect the issue will not be users that don't have tensorflow, but for people that have it installed in their environment but aren't necessarily using it with transformers.

As TF is always a pain to find the set of compatible libraries, I can imagine a lot of people complaining.

Could we have the if parse logic put into a function and then add that as validation when people instantiate objects or use functions?

Or perhaps, create something like require_keras_lt_2 we can use to decorate objects?

[Rocketknight1]

I did some testing with this! I set up an environment that should trigger the exception (TF 2.15 + Keras 3, no tf_keras). Initializing TF objects caused the exception to be thrown. However, I was able to initialize torch models and run them fine. I think our lazy loading protects us here - the exception should only appear when the user explicitly initializes TF objects.

[Rocketknight1]

(Also in general I think the library maintainers segregated TF stuff pretty hard in the past to make sure that it didn't spam the console or allocate GPU memory)

[amyeroberts]

@Rocketknight1 Have you run the slow model tests on this branch?

[Rocketknight1]

I haven't, but I'll try it with a few models today!

[Rocketknight1]

Quick update @amyeroberts - I tested a couple of models with TF 2.14 + Keras 2.14, as well as TF 2.15 + Keras 3. Behaviour was as expected - the TF 2.15 run threw an error and asked me to pip install the tf-keras backwards compatibility package. After that was installed, the slow tests all passed.

[Rocketknight1]

@amyeroberts all comments should be addressed now! The last issue is the pytest stuff that should be resolved by rebasing once the PR for it is merged.

[amyeroberts]

Thanks for all the work handling this tricky TF compatibility case and testing in different environments!

[echan5]

In case this isn't already on your radar, just wanted to let you know that the changes in this PR now cause an error for an import like from transformers import AdamWeightDecay. The error is:

RuntimeError: Failed to import transformers.optimization_tf because of the following error (look up to see its traceback):
module 'keras.optimizers.schedules' has no attribute 'LearningRateSchedule'

I'm guessing due to the change in these lines here: https://github.com/huggingface/transformers/blob/main/src/transformers/optimization_tf.py#L32
https://github.com/huggingface/transformers/blob/main/src/transformers/optimization_tf.py#L215
(so more generally, the error comes up with from transformers import optimization_tf too)

It looks like changing keras.optimizers.schedules.LearningRateSchedule to keras.optimizers.schedules.learning_rate_schedule.LearningRateSchedule might fix it?

[ArthurZucker]

good catch! Would you like to open a PR for a fix @echan5 ?

[Rocketknight1]

Thanks for the warning @echan5! This is quite urgent, I'm on it, and I'm surprised our tests didn't catch it!

[echan5]

Thank you @Rocketknight1! (sorry, I didn't get a chance to get to it earlier)