[Paper Review] PyTorch Distributed: Experiences on Accelerating Data Parallel Training

Li Shen, Yanli Zhao+9 more|arXiv (Cornell University)|Jun 28, 2020

Software System Performance and Reliability26 references116 citations

TL;DR

The paper presents the design, implementation, and evaluation of PyTorch DistributedDataParallel (DDP) to accelerate data-parallel training, including gradient bucketing, overlap of computation and communication, and skipping gradient synchronization, achieving near-linear scalability on large GPU counts.

ABSTRACT

This paper presents the design, implementation, and evaluation of the PyTorch distributed data parallel module. PyTorch is a widely-adopted scientific computing package used in deep learning research and applications. Recent advances in deep learning argue for the value of large datasets and large models, which necessitates the ability to scale out model training to more computational resources. Data parallelism has emerged as a popular solution for distributed training thanks to its straightforward principle and broad applicability. In general, the technique of distributed data parallelism replicates the model on every computational resource to generate gradients independently and then communicates those gradients at each iteration to keep model replicas consistent. Despite the conceptual simplicity of the technique, the subtle dependencies between computation and communication make it non-trivial to optimize the distributed training efficiency. As of v1.5, PyTorch natively provides several techniques to accelerate distributed data parallel, including bucketing gradients, overlapping computation with communication, and skipping gradient synchronization. Evaluations show that, when configured appropriately, the PyTorch distributed data parallel module attains near-linear scalability using 256 GPUs.

Motivation & Objective

• Demonstrate the design and implementation of PyTorch's distributed data-parallel module (DDP).
• Show how to achieve mathematical equivalence between distributed and local training by synchronizing gradients.
• Identify performance bottlenecks and present optimization techniques to maximize training throughput.
• Provide real-world insights and measurements from internal and external deployments.
• Highlight practical caveats and directions for future improvements in industrial-scale distributed training.

Proposed method

• Present DDP as an nn.Module that wraps a local model to ensure non-intrusive integration.
• Describe gradient reduction techniques including autograd hooks and AllReduce-based gradient averaging.
• Introduce gradient bucketing to improve AllReduce efficiency by grouping small gradients into larger buckets.
• Explain overlapping computation with communication to hide latency in gradient reductions.
• Discuss the no_sync context manager to enable gradient accumulation across multiple iterations.
• Detail collective backends (NCCL, Gloo, MPI) and a ProcessGroup abstraction to route communications.

Experimental results

Research questions

• RQ1 How can PyTorch's DDP guarantee mathematical equivalence to local training while remaining non-intrusive to user code?
• RQ2 What optimizations (bucketing, overlap, skip_sync) best improve distributed data-parallel training performance?
• RQ3 How do different communication backends (NCCL, Gloo, MPI) impact scalability and throughput?
• RQ4 What are the practical caveats and failure modes when deploying DDP at large scale?
• RQ5 What runtime configurations (bucket size, process groups, unused-parameter handling) affect convergence and speed?

Key findings

• DDP can achieve near-linear scalability on up to 256 GPUs when configured appropriately.
• Gradient bucketing and overlapping computation with communication substantially improve performance, especially for models with many small parameters.
• Skipping gradient synchronization (no_sync) reduces amortized communication overhead with minimal impact on convergence speed.
• Communication is the dominant latency component and bucket size substantially influences efficiency; improper bucket sizing can negate benefits.
• NCCL and Gloo backends show different performance characteristics; bucket sizing and process-group configuration are critical for optimal throughput.
• Experiments corroborate substantial adoption and impact of DDP in production workloads.