How to do distributed training?

[ayaka14732]

I have been searching hard to find a tutorial about doing distributed training in JAX (e.g. with 100 v2-8 Cloud TPUs). It seems that Ray can achieve this goal (mesh-transformer-jax, swarm-jax), but I don't quite understand how to make it work.

[functionstackx]

check out https://github.com/sholtodouglas/scalingExperiments for data and tensor parallelism. pipeline parallelism through k8s/ray is coming soon-ish according to the repo owner

[mwitiderrick]

Check out Distributed training with JAX & Flax: https://www.machinelearningnuggets.com/distributed-training-with-jax-and-flax/