As mentioned in the chapter introduction, the scale of modern machine learning often outstrips the capabilities of a single compute device. Whether dealing with enormous datasets or models containing billions of parameters, distributing the computational load becomes a necessity. This section introduces the fundamental concepts of parallelism employed to tackle these large-scale problems, providing the conceptual groundwork for understanding JAX's distributed computing features, particularly pmap.

Parallelism, in essence, involves dividing a computational task into smaller, independent or semi-independent parts that can be executed simultaneously on multiple processing units. In the context of machine learning training and inference, several distinct strategies have emerged, primarily driven by whether the bottleneck is the amount of data or the size of the model itself.

Data Parallelism

This is arguably the most common parallelization strategy for accelerating model training. The core idea is simple: if you have a large dataset, split it into smaller chunks and process each chunk concurrently.

Model Replication: The same model architecture and its current parameters are replicated across multiple compute devices (e.g., GPUs or TPU cores).
Data Sharding: The global batch of training data is divided into mini-batches, with each device receiving a unique mini-batch.
Parallel Forward/Backward Pass: Each device independently performs the forward and backward pass computations on its local mini-batch using its copy of the model, calculating local gradients.
Gradient Aggregation: The gradients calculated on each device are aggregated across all devices. Common aggregation methods include averaging (pmean) or summing (psum). This aggregated gradient represents the gradient computed over the entire global batch.
Parameter Update: The model parameters on all devices are updated using the aggregated gradient. This ensures that all model replicas remain synchronized.

This approach effectively increases the total batch size that can be processed per step, often leading to faster convergence and better utilization of multiple accelerators. It works best when the model fits comfortably into the memory of a single device, but the dataset is very large. JAX's pmap function is primarily designed to implement this Single-Program Multiple-Data (SPMD) style of data parallelism.

Flow of data parallelism. Data is sharded, processed in parallel on model replicas, gradients are aggregated, and parameters are updated synchronously.

Model Parallelism

When a model becomes so large that its parameters, activations, or intermediate states cannot fit into the memory of a single accelerator, data parallelism alone is insufficient. Model parallelism addresses this by partitioning the model itself across multiple devices.

Model Partitioning: Different parts of the model (e.g., layers or specific operations within layers) are assigned to different devices.
Data Flow: Input data flows through the model parts sequentially. The output of a computation on one device might become the input for the next part of the model residing on another device.
Inter-Device Communication: Significant communication is often required between devices to pass activations forward during the forward pass and gradients backward during the backward pass.

This strategy allows training extremely large models but introduces complexities in partitioning the model effectively and managing the communication overhead between devices. Common approaches include tensor parallelism (splitting individual weight matrices) and pipeline parallelism.

Flow of model parallelism. A single data sample flows through model parts distributed across multiple devices.

Pipeline Parallelism

Pipeline parallelism is a more sophisticated form of model parallelism designed to improve device utilization. Instead of processing a single batch sequentially through model parts spread across devices (leaving some devices idle while others work), pipeline parallelism divides the batch into smaller micro-batches.

Staging: The model is divided into sequential stages, each assigned to a device or group of devices.
Micro-batching: The input data batch is split into smaller micro-batches.
Pipelined Execution: Devices process micro-batches in a staggered fashion. As Device 1 finishes processing micro-batch 1 for Stage 1 and starts processing micro-batch 2, Device 2 can begin processing micro-batch 1 for Stage 2. This creates a pipeline effect, keeping more devices busy simultaneously.

This approach helps mitigate the "bubble" of idle time inherent in naive model parallelism but requires careful management of dependencies, scheduling, and state (like activations needed for backward passes).

Illustration of pipeline parallelism across three devices/stages over time steps (T1-T5). Micro-batches (MB1, MB2, MB3) enter the pipeline sequentially, allowing devices to operate concurrently on different micro-batches.

Focus on Data Parallelism with pmap

While all these parallelism strategies are relevant in modern machine learning, this chapter will primarily focus on data parallelism. JAX provides powerful tools for this through its pmap transformation, which maps a function across multiple devices, automatically handling data distribution and providing mechanisms for collective communication (like gradient aggregation). Understanding data parallelism and pmap is fundamental for scaling most standard training workloads in JAX. We will delve into device management, the SPMD execution model, collective operations, and practical implementation details in the following sections. Conceptual understanding of model and pipeline parallelism provides valuable context, especially when considering extremely large models discussed later in Chapter 6.