Distributed JAX Training with Custom Data Loading in Multi-Host and Multi-Process Environments

[Lyuliyao]

I want to use JAX in multi-host and multi-process environments. I have found the tutorial at "https://jax.readthedocs.io/en/latest/multi_process.html". However, it doesn't seem to explain how to load data on different hosts. Suppose I have 100 datasets and 10 processes with 10 devices, labeled as "data_001" to "data_100". I want the first process to load "data_001", "data_011", and so on. Then, I would like to train the model based on the dataset in each process and average the gradients at the end.

[wen020]

Do you currently know the best practices for implementing multi-host data parallelism with JAX?

[Lyuliyao]

I still don't have some idea about how to do this.

[uvikas]

@Lyuliyao Did you end up figuring this out? Any good examples to follow?