Data pipeline solution for jax? #19436
Unanswered
louis-shao
asked this question in
General
Replies: 1 comment
-
Do you currently know the best practices for implementing multi-host data parallelism with JAX? |
Beta Was this translation helpful? Give feedback.
0 replies
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
Hi, if we use multi-GPU or multi-host training with jax, what's the expected data loading pipeline? When we use tensorflow we use the tf.data module with the tf distributed strategies, so when switched to jax we are using the same pipeline with dlpack to convert the result for jax to consume. Not sure if this is the optimal solution.
Beta Was this translation helpful? Give feedback.
All reactions