Source-linked AI summary
PyTorch Distributed: Experiences on Accelerating Data Parallel Training
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, Soumith Chintala
TL;DR
The paper addresses how to scale deep-learning training with data parallelism while managing the dependencies between computation and communication. It designs and evaluates PyTorch’s distributed data parallel module using gradient bucketing, overlap, and synchronization skipping. Properly configured, the module achieves near-linear scalability on 256 GPUs.
Problem
Scaling larger models and datasets requires distributed training, but dependencies between computation and communication make efficient data-parallel training non-trivial.
Method
The paper designs and evaluates PyTorch’s distributed data parallel module with gradient bucketing, computation–communication overlap, and skipped gradient synchronization.
Results
Near-linear scalability was achieved on 256 GPUs when PyTorch distributed data parallel was configured appropriately.
Takeaways & Limitations
Communication is the dominant training-latency contributor, and tuning bucket sizes and synchronization frequency is important for distributed-training performance.
Takeaways & Limitations
Variable autograd graphs can require consistent reduction ordering and may cause backward hangs when gradients are absent from iterations.
Abstract
from arXiv · showhide
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.
1. INTRODUCTION
The paper presents PyTorch v1.5’s distributed data parallel module as a minimally intrusive way to scale training while preserving local-training equivalence. Its design addresses computation–communication dependencies through gradient bucketing, overlap, and selective synchronization, achieving strong measured scalability.
- Distributed data parallelism replicates models across resources, processes separate data portions, and synchronizes gradients or updated parameters during training.
- Mathematical equivalence requires distributed training to produce the same model as local training without model replication.
- Communication and computation have subtle dependencies, making high-performance distributed training a design challenge.
- PyTorch exposes distributed data parallel as an nn.Module with the same forward API as the user model, minimizing application-code changes while enabling internal interception.
- More than 2X speedup could result from properly configured bucket sizes, while appropriate synchronization skipping reduces amortized communication overhead without noticeably degrading convergence speed.
- The paper reports a widely adopted industrial solution, real-world caveats, and performance-tuning experience, with more than 60% of Facebook production GPU hours using the package during the cited period.
2. BACKGROUND
The background frames DistributedDataParallel as PyTorch’s multi-process, multi-device data-parallel option and contrasts gradient synchronization with parameter averaging. Gradient synchronization preserves mathematical equivalence and permits communication–computation optimization opportunities that parameter averaging lacks.
- PyTorch provides DistributedDataParallel for multi-process data parallel training across GPUs and machines, alongside DataParallel and RPC for other distributed settings.
- Parameter averaging can differ substantially from local training and sometimes harm accuracy because optimizer states may diverge across replicas.
- Parameter averaging separates backward computation from parameter communication, leaving one resource type idle and foregoing optimization opportunities.
- PyTorch therefore implements distributed training by synchronizing gradients rather than parameters, while still allowing applications to construct parameter averaging explicitly.
- AllReduce computes gradient summation across processes and returns the same result tensor to every participant through collective synchronized communication.
3. SYSTEM DESIGN
DDP provides a minimally intrusive module that preserves mathematical equivalence by synchronizing replicated models’ gradients, while exposing hooks for communication optimizations. Its design addresses small-tensor inefficiency, computation–communication overlap, ordering mismatches, skipped gradients, and gradient accumulation.
- API design: Mathematical equivalence requires all replicas to start identically and receive the same gradients after every backward pass.Independent local optimizers can therefore update corresponding replicas consistently.
- Gradient reduction: AllReduce computes gradient summation across processes and returns the same result tensor to every participant.DDP relies on collective communication libraries including NCCL, Gloo, and MPI.
- API design: DDP wraps a local model as an nn.Module, synchronizes gradients during backward, and lets applications reuse the forward API with minimal changes.Replicas start from the same model state, while autograd hooks trigger gradient reduction during backward execution.
- Performance design: Separating gradient computation from synchronization creates a hard boundary that leaves either computation or communication idle.Overlapping them is therefore a central performance opportunity, especially because collective communication performs poorly on small tensors.
- Gradient reduction: Gradient bucketing combines multiple ready gradients before asynchronous AllReduce, improving throughput and latency compared with reducing each tensor separately.The design avoids one giant reduction so communication can begin before backward computation finishes.
- Gradient reduction: DDP must use the same bucketing order across processes because different gradient-ready orders can mismatch AllReduce contents and cause incorrect results or crashes.PyTorch v1.5.0 uses the reverse order of model.parameters() as an approximation to a consistent bucketing order.
- Gradient reduction: Unused gradients can prevent buckets from becoming ready, causing the backward pass to hang when different iterations execute different model sub-graphs.DDP’s construction-time gradient-to-bucket mapping creates this dependency.
- Gradient accumulation: Reducing synchronization frequency supports local iterations or microbatch accumulation, but requires a no-sync interface because DDP cannot infer the optimizer’s intended accumulation boundary.The approach conflicts with marking unused parameters ready at the end of every forward pass.
4. IMPLEMENTATION
PyTorch DDP combines a Python-facing module with a C++ gradient-reduction core, coordinating gradient buckets through autograd hooks and AllReduce. Its behavior depends on configuration choices such as bucket size, device affinity, buffer handling, and unused-parameter detection.
- Implementation architecture: DDP exposes a user-facing Python API while delegating the core gradient-reduction algorithm to C++ through Pybind11.The Python layer composes non-performance-critical components, while C++ handles the performance-critical reduction logic.
- Core gradient reduction: The reducer builds parameter-to-bucket mappings, installs autograd hooks, launches bucket AllReduce operations, and detects globally unused parameters.These four components constitute the main gradient-reduction implementation.
- Core gradient reduction: DDP accounts for device affinity when creating buckets, broadcasts model buffers from rank 0 before forward passes, and can gather global unused-parameter information with an additional AllReduce.Buckets are placed on parameter devices, while unused-parameter detection uses a bitmap to collect process-wide information.
- Core gradient reduction: Autograd post-hooks track gradient readiness and mark a bucket ready for AllReduce when all gradients assigned to it are available.Because gradient readiness order is not guaranteed, DDP uses per-bucket counts to determine readiness.
- Core gradient reduction: Bucket size trades lower amortized communication overhead against longer reduction lead time; the default bucket size is 25MB.Applications are expected to measure this trade-off and tune the bucket-size setting for their workloads.
5. EVALUATION
The evaluation examines DDP latency, bucket-size choices, scalability, skipped synchronization, and Round-Robin process groups across models, backends, and GPU counts. Results show that communication and configuration strongly influence performance, with selected techniques reducing distributed-training overhead.
- Latency Breakdown: Communication dominates DDP training latency, especially as model size increases, making the backward pass the main optimization target.In the backward pass, communication takes more than half of total delay; NCCL is considerably faster than Gloo.
- Bucket Size: Intermediate bucket sizes provide the best per-iteration latency: 10–25MB with NCCL on ResNet50, while 5MB is fastest with Gloo.The best bucket size depends on the communication backend and application.
- Bucket Size: Above 5MB bucket sizes avoid noticeable speed regression when scaling from 16 to 32 GPUs, whereas 0MB buckets become substantially slower.Asynchronous execution and parallelism can hide some additional AllReduce delay at larger bucket sizes.
- Scalability: At 256 GPUs, NCCL reaches a scaling factor of 128 relative to local training, while Gloo incurs about 3X ResNet50 and 6X BERT per-iteration slowdowns.The larger slowdown for BERT with Gloo indicates that network capacity is the bottleneck in that experiment.
- Skip Gradient Synchronization: Skipping gradient synchronization reduces amortized latency, but its convergence impact depends on configuration and can worsen final loss.No synchronization caused negligible convergence-speed degradation in one MNIST setting but hurt final loss with a larger batch size and learning rate.
- Round-Robin Process Group: Round-Robin process groups deliver their largest reported gain for BERT with NCCL, where three groups achieve 33% speedup over one group on 16 GPUs.ResNet50 with NCCL shows negligible differences across group counts, while ResNet50 with Gloo benefits consistently from rr3 over rr1.
6. DISCUSSION
The discussion finds that DDP performance depends strongly on configuration and deployment conditions rather than one universal setting. It highlights backend choice, bucket sizing, resource placement, and prospective communication reductions as key considerations.
- No single DDP configuration works for every use case because optimal settings depend on model size, model structure, and network bandwidth.Developers can narrow the search using summarized intuitions, but deployment-specific empirical measurements remain necessary.
- NCCL is considerably faster than Gloo in most use cases, making it the preferred backend when available.
- Both excessively small and large gradient buckets hurt communication performance, so the optimal size lies between them and depends on the backend.Optimal bucket sizes likely increase sub-linearly with model size.
- Scaling across machine boundaries can significantly slow NCCL when inter-machine bandwidth is much lower than intra-machine bandwidth.Keeping the DDP group within one machine is recommended in this case; no-sync mode is an option for larger-scale training if convergence remains acceptable.
- Randomly skipped layers do not reduce communicated data under fixed parameter-to-bucket mappings because AllReduce communicates at bucket granularity.Bucket-level skipping requires extra coordination across DDP processes, such as shared random seeds or an authority process broadcasting the plan.
- Adaptive gradient compression could reduce communication volume by using only the precision necessary for gradients, with prior work communicating one bit per gradient at a small accuracy cost.
7. RELATED WORK
Related work classifies distributed training across update synchronization, iteration overlap, and parallelism dimensions, then compares DDP with alternative communication and parallelization approaches.
- Distributed training is categorized as synchronous or asynchronous, cross-iteration or intra-iteration, and data-parallel or model-parallel.
- Data parallelism distributes input data across replicated models, whereas model parallelism divides a model across devices or machines when it cannot fit on one.
- Prior work explores alternative communication algorithms, including tree-based AllReduce, heterogeneity-aware interconnection structures, and AllReduce decomposition.
- PyTorch DDP was implemented with bucketing, computation-communication overlap, and skipped synchronizations before similar techniques appeared in TensorFlow and GradientFlow.GradientFlow additionally requires consensus on which gradients to synchronize.
- Consensus overhead can outweigh gradient-synchronization speedups for small models or networks with large round-trip delays.
- PipeDream combines stage-wise data parallelism with pipeline and model parallelism, while Mesh-TensorFlow and ZeRO combine data and model parallelism differently.PipeDream trades some accuracy for speed by using the latest gradients from its pipeline stages.
8. CONCLUSION
The paper explains and evaluates PyTorch DDP, identifying its main acceleration techniques and practical configuration constraints. Properly configured NCCL DDP achieves near-linear scalability on 256 GPUs, while backward computation remains the most expensive step.
- PyTorch DDP accelerates training by aggregating gradients into buckets, overlapping communication with computation, and skipping synchronizations.
- Near-linear scalability on 256 GPUs is achievable with the NCCL backend when DDP is configured properly.
- The backward pass is DDP’s most expensive training step, requiring optimization efforts from framework developers and empirical knob configuration by application developers.
- The paper reports real-world gradient-synchronization caveats and shares lessons from applications to support broader adoption and future improvements.