Data Parallelism › Distributed Data Parallel: The Algorithm

Distributed Data Parallel (DDP), introduced by Li et al. (2020) and implemented in PyTorch's torch.nn.parallel.DistributedDataParallel, is the standard approach for data-parallel training today. DDP was a significant improvement over PyTorch's earlier DataParallel module, which used a parameter-server-like approach on a single node and suffered from exactly the bottleneck problems described above. DDP's key properties are:

• Each GPU (called a "rank") holds a complete copy of the model
• Each rank processes a different shard of the mini-batch
• Gradients are synchronized via all-reduce before the optimizer step
• All ranks apply identical parameter updates, so all copies stay identical

Let us walk through one training step in detail, because each phase has interesting properties worth understanding.

Each rank $r$ receives a different mini-batch $\mathcal{B}_r$ of size $b$. The global effective batch size is $N \times b$, where $N$ is the number of ranks. Each rank computes its local loss:

$$\mathcal{L}_r = \frac{1}{|\mathcal{B}_r|} \sum_{x \in \mathcal{B}_r} \mathcal{L}_x(\theta)$$

where:

• $\mathcal{L}_r$ is the local loss on rank $r$, a scalar
• $\mathcal{B}_r$ is the mini-batch assigned to rank $r$, a set of examples
• $\theta$ is the current parameter vector, identical on all ranks at the start of the step

This forward pass is entirely local. No inter-GPU communication happens during the forward pass. Each rank can run its forward computation at full speed without waiting for any other rank. This is one of the reasons data parallelism is so efficient: the most compute-intensive phase of training (the forward pass through a large transformer) is completely parallel.

The forward passes on different ranks are not strictly synchronized either. If rank 2 finishes its forward pass slightly before rank 1, it can immediately start backpropagation without waiting. Synchronization only becomes necessary when gradients are ready and need to be averaged.

During backpropagation, each rank computes local gradients $\nabla_r \mathcal{L}_r$ with respect to every parameter. Backpropagation visits layers in reverse order, starting from the output layer and working back toward the input embeddings. As each layer's backward pass completes, its parameter gradients become available.

DDP does not wait for the entire backward pass to finish before communicating. Instead, it overlaps gradient communication with backward computation using a technique called gradient bucketing. Understanding why this matters requires thinking about the timeline of a training step.

In a naive implementation, the full sequence of events would be: run the entire forward pass, run the entire backward pass (computing all gradients), communicate all gradients via all-reduce, then run the optimizer step. The communication phase sits between backpropagation and the optimizer update and cannot begin until all gradients are ready. During the communication phase, the GPU's compute units sit idle.

With gradient bucketing, this idle time is eliminated. Here is how it works. DDP groups parameters into buckets, with each bucket containing roughly 25 MB of gradients by default. The buckets are ordered in reverse layer order, so the last layers (whose gradients become available first during backward) form the first buckets. As soon as all gradients in a bucket are ready, DDP immediately launches an all-reduce operation on that bucket, while the backward pass continues computing gradients for earlier layers in parallel.

By the time backpropagation reaches the earliest layers (embeddings, first transformer blocks), several all-reduce operations have already completed for the later layers. When the backward pass finally finishes, the last all-reduce is either already done or nearly done. The total step time is reduced from compute + communication to approximately max(compute, communication) plus a small serial residual.

The choice of 25 MB per bucket involves a tradeoff. Smaller buckets reduce the latency before the first all-reduce starts (because buckets fill up faster), but they increase the total number of collective communication operations, each with fixed launch overhead. Larger buckets amortize that overhead but delay the start of communication and require more memory in the communication buffers. For most language model training, 25 MB is a reasonable default, but you can tune it with the bucket_cap_mb parameter.

The all-reduce operation takes a tensor from every rank, computes an element-wise sum (or average), and writes the result back to every rank simultaneously. After all-reduce on gradients:

$$\nabla_{\text{sync}} = \frac{1}{N} \sum_{r=0}^{N-1} \nabla_r \mathcal{L}_r$$

where:

• $\nabla_{\text{sync}}$ is the synchronized gradient, identical on all ranks after the operation completes
• $N$ is the number of ranks (GPUs)
• $\nabla_r \mathcal{L}_r$ is the local gradient computed by rank $r$, a vector of the same shape as $\theta$

This synchronized gradient is mathematically equivalent to the gradient you would have computed if you had processed all $N$ mini-batches on a single device. The averaging ensures that the gradient scale does not grow with the number of GPUs, so the same learning rate works regardless of $N$ (assuming the global batch size is proportionally increased). This is a critical design choice: without the averaging, gradients would grow proportionally to $N$, requiring you to scale down the learning rate by $N$ every time you added more GPUs, which would break any learning rate schedule you had tuned.

The all-reduce operation is a collective operation, meaning it requires all ranks to participate simultaneously. Every rank must call all-reduce at the same time (with the same tensor shape), and every rank blocks until the operation completes. This synchronization barrier is what guarantees that all ranks end up with identical gradients.

After gradient synchronization, every rank has identical $\nabla_{\text{sync}}$. Each rank independently runs its optimizer (Adam, AdamW, SGD, or any other) using this synchronized gradient:

$$\theta \leftarrow \theta - \alpha \cdot \text{Optimizer}(\nabla_{\text{sync}})$$

where $\alpha$ is the learning rate. Because every rank starts the step with identical $\theta$ and identical $\nabla_{\text{sync}}$, every rank computes an identical update and every rank ends the step with identical $\theta$. No parameter synchronization is needed after the optimizer step. The model copies remain perfectly in sync as long as they were in sync at the start, which is guaranteed by the all-reduce.

This design is elegant. You never need to communicate parameters. You only need to communicate gradients (once per step), and the rest of the training loop is identical to single-GPU training. The optimizer, gradient clipping, learning rate schedule, and all other components of the training pipeline require no modification.