Data Parallelism › The Data Parallelism Problem

Before we look at how data parallelism works, it helps to understand what problem it is solving and why naive approaches fail.

Suppose you have a model with parameters $\theta$ and you want to train on a dataset $\mathcal{D}$ with $N$ GPUs. The gradient on the full dataset is the average over all examples:

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

where:

• $\theta$ is the vector of all model parameters
• $\mathcal{D}$ is the full training dataset
• $\mathcal{L}_x(\theta)$ is the loss on a single example $x$
• $\nabla \mathcal{L}_x(\theta)$ is the per-example gradient, a vector with as many dimensions as there are parameters

If you split $\mathcal{D}$ into $N$ equal shards and compute the gradient on each shard separately, the average of those shard gradients equals the full-dataset gradient exactly. This is the mathematical foundation of data parallelism: gradient averaging over shards is equivalent to gradient computation over the union of those shards.

The challenge is how to perform that averaging efficiently across $N$ GPUs that are connected by a communication network.

The simplest parallelism strategy would be to split $\mathcal{D}$ into $N$ shards, run each on a separate GPU, collect all gradients on one designated "parameter server" GPU, average them there, and broadcast the updated parameters back to all GPUs. This approach, called parameter server training, was popular in early distributed machine learning research (Dean et al., 2012) and is straightforward to implement.

Parameter servers have a fatal flaw that becomes apparent at scale: the parameter server is a communication bottleneck. Consider the data flow. Each of the $N$ worker GPUs must send its full gradient vector to the parameter server. The parameter server must receive $N$ gradient vectors, sum them, and then broadcast the result to all $N$ workers. The parameter server's inbound bandwidth must handle $N$ gradient vectors simultaneously, and its outbound bandwidth must serve $N$ simultaneous downloads.

For a 7B parameter model with BF16 gradients (2 bytes per parameter), the gradient tensor is 14 GB. With 64 GPUs, the parameter server must receive $64 \times 14 = 896$ GB of gradient data and send back 14 GB to each of 64 workers (another 896 GB) every single training step. No single GPU has that kind of network bandwidth. Even on InfiniBand HDR with 200 Gb/s links, transferring 896 GB would take over 35 seconds per step, compared to a compute time of perhaps 250 ms. The ratio is roughly 140-to-1 in favor of the bottleneck.

Asynchronous parameter servers alleviate the synchronization requirement but introduce gradient staleness: workers apply gradients computed on stale parameters, which can harm convergence. The tradeoffs are poor.

Modern data parallelism solves the bottleneck problem with a different communication pattern called all-reduce, which distributes the aggregation work across all GPUs equally so that no single GPU is a bottleneck.

One subtle point before we proceed: for data parallelism to produce the correct result, every GPU must apply the same parameter update at the end of each step. If different GPUs apply different updates (because they each used only their local gradients), the model replicas would diverge. After a few steps, you would have $N$ different models rather than one model trained on $N$ times as much data.

This is why gradient synchronization is so important. Before any GPU applies an optimizer update, all GPUs must have agreed on a single averaged gradient. The moment any GPU deviates from this shared gradient, the replicas diverge and training becomes incoherent. DDP's design makes this guarantee by construction: no optimizer step can proceed until the all-reduce is complete on every rank.