THE PARALLELISM MESH ZOO

Summary from the TLDR AI newsletter

There are many parallelization strategies available for scaling training runs to work on more GPUs. This post discusses parallelization strategies in a more schematic way by focusing only on how they affect the device mesh. The device mesh is an abstraction used by PyTorch and JAX that takes GPUs and organizes them into an N-D tensor that expresses how the devices communicate with each other. Being able to explain why a device mesh is set up the way it is is a good check for whether you understand how the parallelization strategy works.