Keras 3 and TF 16.1 support #4620
Labels
enhancement
New feature or request
Comments
|
An update on dependencies. Since 0.24, tensorflow-probability can be installed alongside Keras 3, but requires the now "legacy" tf_keras to be installed as well. As far as I understand, both libraries still rely completely on Keras 2, and I wonder if TFF actually cares about what these libraries use (e.g. uses their computation graph) or just rely on their black-boxed functionality. As for TFF itself, I see heavy dependence on Keras 2 API in implementation, which, as far as I understand, means considerable reimplementation for Keras 3... Sounds like a dependency hell if you ask me. |
|
Maybe we can simplify the dependency troubles. Lets brainstorm...
I think
Does this somewhat match your thinking? |
|
Yeah, I guess. Compression seems easy removable. Apart from a elias_gamma_encode.py that may be hidden behind a conditional import, there is only a weird dependency in worker_binary which I can't find being used even transitively. I'll try to brute-force transitive dependencies tomorrow to see if we can make it kinda work. I am afraid that Keras 3 claims a public API "compatibility", but I'd guess TFF uses slightly more than public API on the go. And taking all Keras 3 story is about the ubiquitous runtime, I'd assume the graph part being extensively updated. Nevertheless, there is only a single way to find out. |
|
Yeah I would recommend a few flags I imagine you are not using when trying to test... --build_tag_filters="${tag_filters}"
--test_size_filters=small,medium,large
--test_timeout_filters=short,moderate,long
--test_tag_filters="${tag_filters}"Where
I think you probably want something like this |
|
@michaelreneer Thanks for the filters! it works with them on main as a baseline. Small status update. It seems like usage of Compression and Privacy are quite limited to gamma Elias Gamma and Differential Privacy. After some consideration, I just removed them from dependencies now to see how the main body is affected by Keras. Well, TF 2.16.1 made a few major changes to Bazel dependencies, which I'm still trying to figure out. If interested, the latest for today is here: https://github.com/evgri243/tensorflow-federated/tree/evgri243/keras3-support Not much yet, stabilizing the build. It seems like there is a story around grpc in WORKSAPCE file, so that, just to ask: any chance it looks familiar? Otherwise, I'll continue exploring what has actually changed. `# bazel test //tensorflow_federated/...` (click to open)GRPS in TF hasn't been changed for ages, so that it is definitely not the source of the problem, updates to bazelrc are not revealing as well. So far I suspect the non-hermetic toolchain patch, but have to dig deeper. |
|
After some googling, it looks it is related to Cython update somewhere down the line; I'll double-check. |
|
Resolved the GRPC issue. There is a "hidden" regression in TensorFlow 2.16.1 due to Cython update. The latest release of GRPC cannot be built with Cython 3.0 due to the errors above. The fix on GRPC side is merged to master, but is yet to be released: grpc/grpc#33918 Selecting between unstable GRPC master and an older release of Cython, I picked the latter and rolled it back to the version as of TensorFlow 2.14.1. |
|
Okay, the update to TensorFlow 2.16.1 has built and passed an expectable number of tests. In attachment is an html test report. At first glance, apart from compression, we have one more dependency to handle: tensorflow-model-optimisation. And according to the issue tensorflow/model-optimization#1119 they have no plans on Keras 3 support apart from forcing Keras 2 legacy mode. The only algorithm affected is IBLT Heavy Hitters. For our own needs, we currently need TFF without any of the affected dependencies (even DiffP), so that I'll prefer stabilizing the core first. I'll hope to hack it around to support both Keras 2 and 3 as APIs are heavily compatible, and there is a trend on libraries forcing Keras 2 compatibility through a global What do you think? |
I've not seen this specific error before, but you are correct the grpc and protobuf deps are always hard to update because of the source dependency on TF. I've tried to update both of those deps recently independently from TF and hit different issue then the one you are hitting. |
Nice!
|
|
<removed, I need to think a little more> |
|
A recent update. I tried a few approaches to handle the problem, and so far the most reasonable one is to run with TF_USE_LEGACY_KERAS=1 for dependencies and adjust TFF's code to run for both Kerases using explicitly tf_keras and keras packages. While I managed to temporarily remove Amazing code quality, by the way! Meanwhile, due to the most logic going through tracing at some stage, it doesn't really matter for TFF, which Keras is being used until it is code-compatible with both Keras 2 and Keras 3. Accidentally, I managed to successfully run a few tests with Keras 3 sequential model and Keras 2 counter metrics... It worked, though this code path is definitely not to be officially supported. Nevertheless, the amount of work is quite dramatic and I'm slowly progressing through it. I'll keep you updated on the progress and hope to show the code soon. Though I am not the fun of TF_USE_LEGACY_KERAS, it seems to be the way for many libraries so far. In future, as all dependencies either update to Keras 3 or start explicitly using tf_keras, it will look much better. If I manage to make tests passing both Keras 2 and 3 variants through parametrisation, it should be quite easy to support both frameworks in the long run. While it came mostly to facelift with TFF, TensorFlow Privacy actually requires a much bigger update for Keras 3. It should either explicitly switch to tf_keras for now or be quite refactored as some parts of it are no longer available in Keras 3, most notably the TensorFlow Estimators (removed in TF 2.16) and tf.keras.legacy namespace (removed in Keras 3) used inside quite intensively and not only for internal implementation, but also for the public APIs. I will definitely need help with that to update the dependency. |
|
Regarding dependencies, #4669 should help. I'll try and ping some of the developers who are more actively working on TensorFlow Privacy. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
TensorFlow seems to be diverging towards Keras 3 after 2.16. Do you have any visibility when TFF may be migrated to Keras 3 as well? We would like to fine-tune a LLama-like model in TFF setting, and its implementation has been completed only in keras-nlp 0.9.3, which is pinned to TF 2.16.1 and Keras 3.
I know migration to Keras 3 is a big way forward even if strictly restricted to TF backends only. A few mentions:
Describe the solution you'd like
TFF supports Keras 3 out of the box.
Describe alternatives you've considered
No other easy alternative to use modern models with TFF to my knowledge.
Generally, supporting Keras 3 seems like a good way forward as it is slowly getting populated with popular models. I generally support the initiative and all the work to implement core models there.
Adding @michaelreneer for visibility. @michaelreneer I remember you mentioning goals to clean the dependency tree or at least make it optional. Can this work be of any help here to avoid being blocked by dependencies at least?
The text was updated successfully, but these errors were encountered: