As machine learning models grow in size and complexity, running computations on a single accelerator (like a GPU or TPU) becomes insufficient. Training large models or processing massive datasets often requires distributing the work across multiple devices. This chapter focuses on how JAX facilitates such distributed computing.

We will start with the fundamental concepts of parallelism relevant to machine learning workloads. You'll learn how JAX manages different compute devices and how to use its core primitive for multi-device execution: pmap (parallel map). We will cover the Single-Program Multiple-Data (SPMD) paradigm that pmap employs and demonstrate how to implement data parallelism, a common technique for accelerating training.

Furthermore, you'll explore essential collective communication operations (like psum, pmean) needed to aggregate information, such as gradients, across devices within a pmap'd function. We will also discuss the use of axis names for more explicit control over these collectives and touch upon advanced partitioning strategies and the concepts behind multi-host distribution. By the end of this chapter, you will understand how to use pmap to scale your JAX computations effectively across multiple accelerators.