Data Parallelism › Gradient Synchronization in Depth

The gradient synchronization step is where most of the complexity in DDP lives. Let us examine several aspects that matter in practice.

DDP synchronizes gradients during the backward pass, not after it. This is achieved through autograd hooks that DDP registers on each parameter's .grad_fn. When the autograd engine finishes computing the gradient for a parameter, the hook fires immediately. DDP accumulates gradients into the appropriate bucket, and when the bucket is full (or when the last parameter in the bucket gets its gradient), DDP launches an asynchronous NCCL all-reduce operation on that bucket.

The hooks fire in the order that autograd computes gradients, which is roughly reverse layer order for a standard sequential model. This means the gradients for output layers (which backpropagate first) are placed into buckets and communicated first, while the gradients for early layers (embeddings, first transformer blocks) are computed and communicated last. By the time the backward pass finishes, many all-reduce operations have already completed.

This design has two important benefits. First, compute and communication run simultaneously: while the CPU continues the backward pass for earlier layers, the GPU's communication unit handles gradient transfers for later layers. Second, the critical path is shortened: the wall-clock time for a training step approaches max(compute_time, communication_time) rather than their sum, as long as the communication can be fully overlapped.

The practical degree of overlap depends on the model architecture. Deep models with many layers provide more opportunity for overlap than shallow models. Attention-heavy layers with long compute times provide more hiding of communication than lightweight MLP layers. For large transformer models (where each layer has hundreds of millions of parameters), the overlap is substantial and gradient bucketing provides large speedups.

Sometimes you want to accumulate gradients over multiple mini-batches before synchronizing. This is called gradient accumulation, and it is used when the target global batch size is larger than what fits in GPU memory in a single batch. The idea is to run $k$ forward-backward passes and only call optimizer.step() after $k$ passes, effectively computing a gradient over $k$ times as many samples per optimizer step.

With DDP, calling backward() triggers an all-reduce after every pass, even when you do not want to synchronize yet. This wastes communication bandwidth: you are averaging gradients $k$ times per optimizer step but only need to average them once. DDP provides the no_sync() context manager to suppress gradient synchronization:

Gradient accumulation is an important practical technique because many language model training recipes target specific global batch sizes (for example, 4 million tokens per step) that far exceed what fits in GPU memory in a single forward pass. By accumulating over $k$ steps with no_sync(), you achieve the target batch size while only communicating once. The communication overhead is the same as without accumulation, but the effective batch size is $k$ times larger.

One thing to be careful about: the loss / accumulation_steps scaling in the code above. Because gradients accumulate additively across the $k$ backward passes, you need to scale each loss by $1/k$ so that the accumulated gradient equals the gradient of the mean loss over the full accumulated batch. Without this scaling, the accumulated gradient would be $k$ times the mean gradient, which would require reducing the learning rate by $k$ to compensate.

DDP averages gradients by default, dividing the sum by $N$ (the number of ranks). This normalization ensures that the gradient scale is the same regardless of how many GPUs you use, which makes learning rate tuning independent of the degree of parallelism.

There is an important subtlety here. The per-rank loss is already averaged over the local batch:

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

where $b = |\mathcal{B}_r|$ is the local batch size. The gradient of this loss is:

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

After all-reduce averaging over $N$ ranks, the synchronized gradient is:

$$\begin{aligned} \nabla_{\text{sync}} &= \frac{1}{N} \sum_{r=0}^{N-1} \nabla \mathcal{L}_r \\ &= \frac{1}{N} \sum_{r=0}^{N-1} \frac{1}{b} \sum_{x \in \mathcal{B}_r} \nabla \mathcal{L}_x(\theta) \\ &= \frac{1}{Nb} \sum_{r=0}^{N-1} \sum_{x \in \mathcal{B}_r} \nabla \mathcal{L}_x(\theta) \end{aligned}$$

where:

• $\nabla_{\text{sync}}$ is the final synchronized gradient used for the optimizer step
• $b$ is the per-rank mini-batch size
• $N$ is the number of ranks (GPUs)

This expression is exactly the gradient of the mean loss over the global batch of size $Nb$. The synchronized gradient is the true full-batch gradient for the combined mini-batch across all ranks, divided by the global batch size. This is the correct gradient for the loss function you are optimizing, which is the mean loss per token (or per example) over the global batch.