-
Notifications
You must be signed in to change notification settings - Fork 225
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Distributed training not working (batch size calculation) #1630
Comments
Here is Colab with reproducible example: https://colab.research.google.com/drive/1aCJVUNfro68fek-o0i_7Iojl6Qtix6NK?usp=sharing |
I can reproduce the error using just keras, so maybe I should open issue there? Or maybe it should be fixed in tensorflow? But the documentation for https://colab.research.google.com/drive/1IxVNDcNoIK4SiX2wuDQKfqR_6Or9P40I#scrollTo=0Hf6qJOxXsqI |
Hi @SuryanarayanaY, the related ticket in keras was closed with the recommendation that this be fixed in keras-nlp. Per @hertschuh: "One should simply apply batch_size after the map and not in _convert_inputs_to_dataset". I can't quite figure out the best way for this to work with keras-nlp API. In particular, it seems like there are several combinations of (1) distribution strategy, (2) input types (e.g., tf.data.Dataset, NumPy arrays), and (3) batching (e.g., pre-batched dataset, explicit batch_size). Currently, in _convert_inputs_to_dataset it will raise an error if you attempt to pass a tf.data.Dataset with explicit batch_size argument. It also looks like there is error handling to prevent you from passing unbatched inputs, but the string matching on the error message may be oudated and not functioning. |
Sorry for the delay, I'm still working on this. It turned out to be more complex to fix than I expected. |
I experimented with a few things, but I could not find a fix in keras-nlp that would work in all cases. However, I do have an easy workaround: If you don't do If you're concerned about not using the last few examples, you can |
@hertschuh do you have a working example? This doesn't seem to work for me. Edit: It seems like the workaround works outside of keras-nlp. Maybe there is something specific to keras-nlp that still needs to be resolved? https://colab.research.google.com/drive/1IxVNDcNoIK4SiX2wuDQKfqR_6Or9P40I#scrollTo=0Hf6qJOxXsqI |
@SuryanarayanaY can we reopen this issue please? |
Describe the bug
This is an issue I am having with keras-nlp, but I am not sure if it can be solved here or should be reported under keras or tensorflow.
Currently, the batch size is not calculated correctly when performing multi-worker distributed training with JAX backend:
To Reproduce
Run (multi-worker?) distributed training with JAX backend.
The issue seems to stem from
keras-nlp/keras_nlp/src/utils/pipeline_model.py
Line 181 in 778ccd7
Here is minimal example where
tensorflow.python.data.experimental.ops.distribute.compute_batch_size()
returns -1 after mapping:Expected behavior
A batched tf.data.Dataset() object is recognized as being batched.
The text was updated successfully, but these errors were encountered: