search

keras-team / keras

Deep Learning for humans
http://keras.io/
Apache License 2.0
61.68k stars 19.42k forks source link

Think about Keras3 preference, should I develop my distributed training CTR model based on Jax or TensorFlow? #19866

Open MoFHeka opened 3 months ago

MoFHeka commented 3 months ago

I am a developer of tensorflow recommenders-addons and I now need to develop an all-to-all embedding layer for multi-GPU distributed training of recommendation models. The old tensorflow distributed strategy clearly did not meet this need. So the question is, should I develop on TF DTensor or Jax? Because it seems that Keras support for TF DTensor is not friendly. But Jax lacks the ability to online inference services and the functional components used by various recommendation algorithms. Also recommenders-addons has a lot of custom operators.

fchollet commented 3 months ago

If you need a SPMD API, I 100% recommend JAX. It's more mature and better optimized than DTensor. You can use it via the easy keras.distribution API with the JAX backend. https://keras.io/guides/distribution/