Explain Different Approaches for Distributed Optimization

Introduction

Training modern deep learning models often requires more computational power than a single GPU can provide. A model like GPT-3 with 175 billion parameters would take years to train on one GPU. Distributed optimization solves this by spreading the training workload across multiple GPUs or machines, reducing training time from years to weeks or even days.

The core challenge in distributed training is coordination. When multiple GPUs each compute gradients on different data, those gradients must be combined correctly to update a single, consistent model. This gradient synchronization step determines both the correctness and efficiency of distributed training.

This article explains the main approaches to distributed optimization, starting with data parallelism (the most widely used method), then covering communication strategies like Ring AllReduce, and briefly introducing model parallelism techniques for extremely large models.

Data Parallelism

Data parallelism is the most common distributed training strategy. The idea is straightforward: replicate the entire model on each GPU, split the training data into batches, and have each GPU process a different batch simultaneously. After computing gradients locally, all GPUs synchronize their gradients by averaging them, then each GPU applies the same update to its model copy.

Consider training with 4 GPUs and a batch size of 128. Each GPU receives 32 samples, computes forward and backward passes independently, and produces gradients. These gradients are averaged across all GPUs, effectively computing the gradient over the full 128-sample batch. Each GPU then updates its model identically, keeping all copies synchronized.

The following code shows a simplified gradient synchronization step:

import torch
import torch.distributed as dist

def synchronize_gradients(model):
    """Average gradients across all processes."""
    world_size = dist.get_world_size()
    for param in model.parameters():
        if param.grad is not None:
            # Sum gradients across all GPUs, then average
            dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)
            param.grad.data /= world_size

import torch
import torch.distributed as dist

def synchronize_gradients(model):
    """Average gradients across all processes."""
    world_size = dist.get_world_size()
    for param in model.parameters():
        if param.grad is not None:
            # Sum gradients across all GPUs, then average
            dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)
            param.grad.data /= world_size

PyTorch’s DistributedDataParallel (DDP) automates this process and adds optimizations like overlapping gradient computation with communication. DDP is the recommended approach for multi-GPU training and typically provides near-linear scaling: 4 GPUs train roughly 4 times faster than 1 GPU.

Communication Strategies

The efficiency of data parallelism depends heavily on how gradients are communicated between GPUs. Two main architectures exist: Parameter Server and AllReduce.

In the Parameter Server architecture, a central server collects gradients from all workers, computes the average, and broadcasts updated parameters back. While simple to implement, the server becomes a bottleneck as the number of workers grows. Communication cost scales linearly with worker count.

The AllReduce approach eliminates this bottleneck. Instead of a central server, all GPUs communicate directly with each other to compute the gradient sum collaboratively. The Ring AllReduce algorithm achieves this with constant communication cost regardless of GPU count.

Ring AllReduce

In Ring AllReduce, GPUs are arranged in a logical ring. Each GPU has a left neighbor and a right neighbor, and only communicates with these two neighbors. The algorithm proceeds in two phases: scatter-reduce and allgather.

During scatter-reduce, each GPU divides its gradient tensor into N chunks (where N is the number of GPUs). In each step, every GPU sends one chunk to its right neighbor and receives a chunk from its left neighbor, accumulating partial sums. After N-1 steps, each GPU holds one complete chunk of the final reduced gradient.

During allgather, the GPUs circulate these completed chunks until every GPU has the full reduced gradient. This also takes N-1 steps.

The total data transferred per GPU is:

where M is the total gradient size and N is the number of GPUs. As N grows large, this approaches 2M, meaning each GPU transfers approximately twice its gradient size regardless of how many GPUs participate. This bandwidth-optimal property makes Ring AllReduce the foundation of modern distributed training frameworks.

NVIDIA’s NCCL (NVIDIA Collective Communications Library) provides highly optimized implementations of AllReduce that leverage high-speed GPU interconnects like NVLink and InfiniBand.

Model Parallelism

When a model is too large to fit on a single GPU, data parallelism is insufficient since it requires each GPU to hold a complete model copy. Model parallelism addresses this by splitting the model itself across GPUs.

Pipeline parallelism divides the model by layers. The first few layers run on GPU 0, the next layers on GPU 1, and so on. Data flows through the pipeline sequentially. The challenge is that naive pipelining leaves most GPUs idle while waiting for activations from earlier stages. Micro-batching mitigates this by splitting each batch into smaller pieces that can fill the pipeline, improving GPU utilization.

Tensor parallelism splits individual layers across GPUs. For example, a large matrix multiplication can be partitioned so each GPU computes part of the output. This approach is common in training massive transformer models where single attention or feed-forward layers may have billions of parameters.

ZeRO and FSDP

The Zero Redundancy Optimizer (ZeRO), developed by Microsoft and implemented in the DeepSpeed library, offers a middle ground between data and model parallelism. Standard data parallelism replicates the entire optimizer state, gradients, and parameters on each GPU, which is highly redundant. ZeRO partitions these components across GPUs while maintaining the simplicity of a data-parallel programming model.

ZeRO has three stages. Stage 1 partitions optimizer states (like Adam’s momentum and variance), Stage 2 adds gradient partitioning, and Stage 3 partitions the model parameters themselves. Each stage reduces memory usage at the cost of increased communication.

PyTorch’s Fully Sharded Data Parallel (FSDP) provides similar functionality as a native PyTorch feature. FSDP shards model parameters, gradients, and optimizer states across GPUs, gathering parameters just-in-time for computation and immediately releasing them afterward. For models with 500 million parameters or more, FSDP and ZeRO enable training that would otherwise be impossible due to memory constraints.

Summary

Distributed optimization enables training of modern deep learning models by spreading computation across multiple GPUs. Data parallelism, where each GPU holds a complete model copy and processes different data, is the most common approach and provides near-linear scaling. The Ring AllReduce algorithm makes gradient synchronization efficient with constant per-GPU communication cost of approximately 2M bytes regardless of GPU count.

For models too large for single-GPU memory, model parallelism techniques like pipeline and tensor parallelism split the model across devices. ZeRO and FSDP offer memory-efficient alternatives that partition optimizer states, gradients, and parameters while preserving the simplicity of data-parallel training.

In practice, PyTorch’s DistributedDataParallel handles most distributed training needs. For larger models, FSDP or DeepSpeed ZeRO provide the necessary memory efficiency without requiring complex model partitioning code.

Videos

• This playlist, “Distributed Data Parallel in PyTorch Tutorial Series” is the official PyTorch tutorial series that provides a clear walkthrough of DDP concepts and implementation, starting from single-GPU training and building up to multi-node distributed training. (Total runtime: 50 mins)

• The animation video, “Ring AllReduce” by Pakz Lessons visually explains how the Ring AllReduce algorithm achieves bandwidth-optimal gradient synchronization via a simple animation showing the scatter-reduce and allgather phases. (Runtime: 41 seconds)

References

• [1] Li, S., Zhao, Y., Varma, R., et al. (2020). PyTorch Distributed: Experiences on Accelerating Data Parallel Training. VLDB 2020. arXiv:2006.15704
• [2] Rajbhandari, S., Rasley, J., Ruwase, O., & He, Y. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC 2020. arXiv:1910.02054
• [3] Patarasuk, P., & Yuan, X. (2009). Bandwidth Optimal All-reduce Algorithms for Clusters of Workstations. Journal of Parallel and Distributed Computing, 69(2), 117-124.

Related Articles

• What is Backpropagation?
• Gradient Descent and Its Variants