Communication Optimization › Communication Overlap

Gradient compression reduces how much data must be transmitted. Communication overlap reduces how long that transmission blocks computation. The core idea is simple: rather than waiting until the entire backward pass completes before starting any communication, we begin transmitting gradient chunks as soon as they are computed.

In a naive synchronous data-parallel implementation, training proceeds in strict phases. The forward pass runs on all GPUs. The backward pass runs on all GPUs. Then all GPUs synchronize: they perform an all-reduce over all gradients. Only after the all-reduce completes does the optimizer step run, applying the now-synchronized gradients to the parameters. The all-reduce cannot begin until the backward pass finishes, and the optimizer step cannot begin until the all-reduce finishes.

This creates a sequential pipeline where communication sits between two computation phases, blocking both. The GPU is idle during communication, and the network is idle during computation. At scale, this is an enormous waste. A 256-GPU cluster spending 50% of time on communication is equivalent to operating a 128-GPU cluster for the same cost.

This pattern appears even in frameworks that use asynchronous training internally. The critical path through the training step passes through both the backward computation and the all-reduce. Even if you can dispatch the all-reduce on a separate CUDA stream, that stream must complete before the optimizer step reads the aggregated gradient. Any communication time that cannot be hidden behind computation adds directly to wall-clock training time.

The severity of this bottleneck grows with the number of GPUs. As you add more workers, each worker holds a smaller fraction of the total batch, which means each backward pass is shorter (fewer examples per worker). But the all-reduce time grows with the gradient size, which does not change. More GPUs means a higher ratio of communication time to compute time, making overlap increasingly important at scale.

The key insight is that backpropagation computes gradients in reverse order: the last layer's gradients are ready before the first layer's. We do not need to wait for the entire backward pass to finish before communicating the gradients from the final layers.

PyTorch's DistributedDataParallel (DDP) exploits this with a bucketing strategy. Parameters are grouped into buckets of a configurable size (default 25 MB). As soon as all gradients within a bucket are ready (computed by the backward pass), the all-reduce for that bucket is launched asynchronously on a separate CUDA stream. Meanwhile, the backward pass continues computing gradients for earlier layers, which goes into a different bucket.

By the time the backward pass completes, the all-reduces for the last several buckets have already finished or are nearly complete. The overlap between backward computation and gradient communication can eliminate most of the communication overhead, provided the computation time exceeds the communication time.

The timing relationship works as follows. Let $T_{\text{bwd}}$ denote total backward pass time, $T_{\text{comm}}$ denote total communication time, and $T_{\text{overlap}}$ denote the portion of communication that can be overlapped with backward computation. The effective training step time becomes approximately:

$$T_{\text{step}} \approx T_{\text{fwd}} + T_{\text{bwd}} + \max(0,\, T_{\text{comm}} - T_{\text{overlap}})$$

where:

• $T_{\text{fwd}}$: forward pass duration
• $T_{\text{bwd}}$: backward pass duration, typically 2-3x the forward pass
• $T_{\text{comm}}$: total time to communicate all gradients via all-reduce
• $T_{\text{overlap}}$: communication time that is hidden behind backward computation

When $T_{\text{comm}} \leq T_{\text{overlap}}$, communication is fully hidden and adds zero overhead. In practice, achieving full overlap requires careful bucket sizing and a sufficiently compute-heavy backward pass.

Bucket sizing involves a tradeoff. Smaller buckets mean more granular overlap (the first few layers' gradients are ready earlier relative to the remaining backward computation), but also more all-reduce operations, each with a fixed setup overhead (kernel launch latency, NCCL handshake). Larger buckets reduce this overhead but delay the start of each all-reduce until a larger chunk of the backward pass has completed, reducing the maximum achievable overlap. For most models, bucket sizes in the range of 25-500 MB work well, with larger models benefiting from larger buckets.

An often-overlooked consideration is that overlap also requires the GPU to run computation and communication kernels simultaneously. This is possible because modern GPUs have separate CUDA stream engines for compute and copy operations, and these can progress in parallel. However, if the compute kernels are already saturating all GPU execution units, there may not be spare capacity for the DMA engine to run communication in parallel with peak computation. In practice, backward passes often have irregular resource utilization, leaving bandwidth for overlap during the parts of backprop that are not matmul-heavy.

The buckets are also ordered carefully. DDP processes buckets in reverse order of how gradients are computed during backpropagation, which means the last layer's bucket is launched first for all-reduce. This ordering matters because the first all-reduce to finish is for the last layer's parameters, which are also the first parameters to be used in the next forward pass. Getting these gradients synchronized earliest maximizes the chance that they are ready before the next forward pass needs them.

In pipeline-parallel training, communication overlap takes a different form. Between pipeline stages, activations flow forward and gradients flow backward. If the send of layer $l$'s output to stage $l+1$ must complete before stage $l+1$ starts its computation, we again have a blocking sequential pipeline.

Techniques like 1F1B (one forward, one backward) scheduling interleave micro-batches so that while one stage is processing micro-batch $k$, the previous stage is already computing micro-batch $k+1$. This fills the pipeline bubbles with useful work. The communication between stages can be further overlapped using CUDA streams that run concurrently with computation kernels, provided the GPU has sufficient resources to multiplex both.

The 1F1B schedule reduces peak memory relative to GPipe's approach of accumulating all forward activations before starting backward. This memory reduction is significant because pipeline stages hold activations for multiple in-flight micro-batches simultaneously. With 1F1B, each stage holds at most two micro-batches of activations at once, enabling deeper pipelines with the same per-GPU memory budget. The communication overlap in 1F1B is a byproduct of the interleaved schedule: since stages are always processing something, the sends and receives between stages happen naturally in parallel with local computation.

For tensor-parallel training, where individual layers are split across GPUs within a single pipeline stage, the communication pattern is different again. Tensor parallel operations like matrix multiplications require all-reduce or all-gather within the tensor parallel group during both the forward and backward passes. These communications are inside the computation critical path, not outside it, so overlap is harder to achieve. Megatron-LM addresses this with sequence parallelism, which staggers the all-reduce operations across the forward and backward passes to reduce peak communication load.