Source-linked AI summary
Stochastic Gradient Push for Distributed Deep Learning
Mahmoud Assran, Nicolas Loizou, Nicolas Ballas, Michael Rabbat
TL;DR
Distributed deep-learning training relies on synchronization that can become costly and straggler-sensitive, while gossip offers lower-coupling but approximate averaging. This paper combines PUSHSUM with stochastic gradient updates in SGP, proves convergence for smooth non-convex objectives, and evaluates it on image classification and translation. SGP variants can train better models in less time in communication-bound settings.
Problem
Exact distributed averaging can be communication-constrained and sensitive to stragglers, motivating decentralized methods that use approximate averaging.
Method
SGP combines PUSHSUM gossip with stochastic gradient descent and supports directed, sparse, and time-varying communication topologies.
Results
The paper proves convergence to stationary points of smooth non-convex objectives at SGD-matching sub-linear rates and finds better models in less time on ImageNet and WMT16 EN-DE workloads.
Takeaways & Limitations
SGP and Overlap SGP can accelerate distributed DNN training and improve fixed-runtime accuracy over ALLREDUCE SGD in communication-constrained settings.
Takeaways & Limitations
The accuracy of D-PSGD and SGP degrades for larger topologies, which the authors hypothesize is due to greater divergence from approximate averaging.
Abstract
from arXiv · showhide
Distributed data-parallel algorithms aim to accelerate the training of deep neural networks by parallelizing the computation of large mini-batch gradient updates across multiple nodes. Approaches that synchronize nodes using exact distributed averaging (e.g., via AllReduce) are sensitive to stragglers and communication delays. The PushSum gossip algorithm is robust to these issues, but only performs approximate distributed averaging. This paper studies Stochastic Gradient Push (SGP), which combines PushSum with stochastic gradient updates. We prove that SGP converges to a stationary point of smooth, non-convex objectives at the same sub-linear rate as SGD, and that all nodes achieve consensus. We empirically validate the performance of SGP on image classification (ResNet-50, ImageNet) and machine translation (Transformer, WMT'16 En-De) workloads. Our code will be made publicly available.
1. Introduction
Distributed training commonly uses exact gradient averaging, but communication bottlenecks make such synchronization vulnerable in constrained networks. The paper introduces SGP and related variants to reduce overhead while retaining convergence and improving fixed-runtime training outcomes.
- Motivation: Large mini-batch distributed SGD averages local gradients across workers, but ALLREDUCE is not robust when network bandwidth is a significant bottleneck.Parameter servers can also introduce a bottleneck and central point of failure.
- Motivation: PUSHSUM replaces exact averaging with approximate, less-coupled message passing, motivating decentralized inexact SGD for lower communication overhead.Approximate averaging introduces a tradeoff that the paper addresses through SGP.
- Contributions: SGP blends parallel SGD with PUSHSUM and supports directed, sparse, and time-varying communication topologies.This broadens the communication settings beyond symmetric gossip methods.
- Contributions: The paper analyzes SGP and Overlap SGP, proving convergence to stationary points of smooth non-convex objectives at an O(1/√nK) rate for an appropriate step-size.Here n is the number of nodes and K is the number of iterations.
- Empirical evaluation: ImageNet and WMT’16 En-De experiments show that SGP variants can reduce communication overhead and mitigate stragglers while improving fixed-runtime model accuracy over ALLREDUCE SGD.The evaluation covers ResNet-50 and Transformer workloads.
2. Preliminaries
The preliminaries formulate distributed training as consensus optimization and contrast exact ALLREDUCE averaging with approximate PUSHSUM gossip. PUSHSUM uses column-stochastic mixing and a de-biasing weight to recover averages under asymmetric communication.
- Problem formulation: The network seeks DNN parameters minimizing the average loss across local data distributions while requiring all nodes to reach consensus.Each node can evaluate stochastic gradients locally but must communicate to access information about other objectives.
- Distributed averaging: Parallel SGD computes local mini-batch gradients and uses ALLREDUCE to obtain the exact stochastic gradient of the global objective.For an overall mini-batch of size nb, each node computes gradients from b samples.
- Approximate distributed averaging: PUSHSUM instead applies gossip iterations to approximate the average of vectors distributed across nodes.The vectors are arranged as rows of a matrix and repeatedly mixed using communication-dependent matrices.
- PUSHSUM: Column-stochastic mixing matrices can represent sparse communication, and under mild connectivity conditions their products converge to a rank-one limit.The resulting limit generally weights nodes according to an ergodic vector π rather than uniformly.
- PUSHSUM: Unlike symmetric mixing, PUSHSUM permits asymmetric communication by maintaining an auxiliary scalar weight and using a de-biased ratio to recover averages.Symmetric matrices can be doubly stochastic but require implementation care to avoid deadlocks.
- Connection to optimization: PUSHSUM approximately synchronizes parameters and can be adapted to decentralized versions of momentum SGD and Adam.The method is presented as an alternative to exact distributed averaging for data-parallel training.
3. Stochastic Gradient Push
SGP interleaves local stochastic gradient updates with PUSHSUM communication, while Overlap SGP hides communication by allowing bounded-delay computation. The approach supports flexible sparse and asymmetric network topologies and reduces to parallel SGD under uniform mixing.
- SGP algorithm: SGP interleaves one local stochastic gradient descent update with one PUSHSUM iteration at every node.Each node maintains model parameters, a PUSHSUM weight, and de-biased parameters.
- SGP algorithm: Gradients are evaluated at de-biased parameters, then used to update local model parameters before communication exchanges PUSHSUM numerators and weights.Communication occurs through messages containing the parameter numerator and scalar weight.
- Communication topology: SGP can use sparse, asymmetric, or time-varying communication topologies, with sparse mixing matrices reducing communication overhead.Nodes need only their outgoing mixing weights rather than the entire mixing matrix.
- Relationship to parallel SGD: Uniform mixing makes SGP mathematically equivalent to parallel SGD using ALLREDUCE.This identifies parallel SGD as a special case of the broader SGP framework.
- Overlapping communication and computation: Overlap SGP makes communication non-blocking and permits τ gradient updates during communication, producing bounded-delay messages when k − k′ ≤ τ.The paper guarantees convergence with bounded delays and reports that τ=1 or 2 is usually sufficient to hide communication overhead.
4. Theoretical guarantees.
The paper extends SGP convergence analysis to smooth, non-convex objectives under bounded delays and establishes convergence of both network averages and individual node variables.
- Assumptions: SGP is analyzed for smooth, non-convex objectives with bounded stochastic-gradient variance, bounded delays, and time-varying mixing connectivity assumptions.The connectivity condition requires sufficiently frequent strongly connected graph unions with bounded diameter.
- Scope of the guarantees: The convergence analysis explicitly extends prior PushSum-based results from strongly convex objectives to smooth non-convex objectives with bounded message delays.The proof and precise constants are provided in Appendix E.
- Convergence of the average: Theorem 1 proves that SGP converges in the stated stationary-point sense under the paper’s assumptions and an appropriately chosen step-size.The convergence bound depends on network diameter, delay bounds, and the sequence of mixing matrices.
- Convergence of the average: The average of the nodes’ parameters converges, but this result alone does not establish convergence of each node’s parameters.The paper separates network-average convergence from node-wise consensus.
- Consensus: Theorem 2 shows that de-biased variables at each node converge to the node-wise average and therefore to a stationary point.For fixed n and large K, the 1/(nK) term dominates the other factors in the bound.
5. Related Work
Related work includes communication-reduction methods, symmetric and asynchronous gossip algorithms, and decentralized SGD; the paper positions SGP as a PushSum-based alternative with broader guarantees and strong empirical comparisons.
- Communication-efficient training: Gradient quantization and multiple local SGD steps are complementary to SGP’s use of approximate rather than exact distributed averaging.The paper describes these approaches as injecting additional approximation noise, creating a communication-versus-accuracy tradeoff.
- Theoretical positioning: The paper claims the first convergence analysis for a PushSum-based method in the smooth non-convex case with bounded message delays.Earlier SGP analysis assumed strongly convex local objectives.
- Gossip protocols: Unlike symmetric gossip methods, SGP permits directed message passing without requiring a response before a node proceeds.The paper associates symmetric message passing with blocking and slower communication in constrained settings.
- Relation to D-PSGD: SGP generalizes D-PSGD under static, undirected, connected topologies with symmetric mixing weights, where PushSum weights remain equal to 1.Under those conditions, the two algorithms are mathematically equivalent.
- Empirical comparisons: Experiments report that SGP is consistently faster than D-PSGD while finding solutions of comparable accuracy.The paper also compares synchronous Overlap SGP with asynchronous gossip-based methods.
6. Experiments
Experiments compare SGP with ALLREDUCE-SGD, D-PSGD, and asynchronous methods across image classification and machine translation. SGP reduces training time and communication overhead, while topology and runtime budget shape its accuracy tradeoffs.
- Scaling and convergence: SGP completes 90 ImageNet epochs faster than ALLREDUCE-SGD and D-PSGD, with roughly 1.5× faster training than D-PSGD over 10 Gbps Ethernet.At 32 nodes, SGP takes approximately 1.7 hours, versus roughly 2.6 hours for D-PSGD and 5.1 hours for ALLREDUCE-SGD.
- Scaling and convergence: SGP and D-PSGD show strong scaling, whereas ALLREDUCE-SGD’s per-iteration time increases substantially as the number of nodes grows.On InfiniBand, all methods exhibit near-linear scaling; on Ethernet, SGP and D-PSGD maintain nearly constant average iteration time.
- Scaling and convergence: SGP consistently outperforms D-PSGD and ALLREDUCE-SGD in total training time across the tested node counts.The experiments use 4, 8, 16, and 32 servers, corresponding to 32, 64, 128, and 256 GPUs, over 10 Gbps Ethernet.
- Parameter deviations: Larger communication topologies reduce SGP and D-PSGD accuracy, consistent with greater divergence between independently maintained node models under approximate averaging.The reported explanation is a hypothesis tied to model differences increasing in larger networks.
- Parameter deviations: Parameter deviations scale with the learning-rate magnitude, drop by an order of magnitude at epochs 30, 60, and 80, and are larger for sparse topologies.Deviations also increase during the first five warmup epochs.
- Communication and the speed-accuracy tradeoff: Communicating more during the first 30 epochs provides a speed-accuracy balance and can mitigate later accuracy tradeoffs.The evaluated hybrids use either ALLREDUCE or 2-peer SGP initially, followed by 1-peer SGP.
- Overlap SGP: Overlap SGP greatly speeds training without accuracy degradation relative to non-overlap SGP, while omitting the push-sum bias correction substantially worsens accuracy.Synchronous 1-OSGP also runs faster than asynchronous AD-PSGD and achieves better training and validation accuracy.
- Fixed runtime budget: Under a fixed runtime budget, SGP reaches 77.1% top-1 validation accuracy in 4.6 hours, compared with 76.2% for ALLREDUCE-SGD after 5.1 hours.Overlap SGP also reaches 77.1% in 2.7 hours on 32 nodes and 256 GPUs over 10 Gbps Ethernet.
7. Conclusion
The paper proposes SGP and OSGP to accelerate distributed deep neural network training, with theoretical guarantees and empirical evidence of faster training and better models in less time.
- SGP and OSGP accelerate distributed DNN training while providing convergence guarantees in the smooth non-convex setting.
- The methods match known convergence rates for parallel SGD.
- Experiments span ImageNet classification with ResNet-50 and WMT16 EN-DE translation with Transformers across several computing infrastructures.
- SGP and OSGP run significantly faster than parallel SGD in communication-bound settings and train better models in less time.
A. Communication Topology
The communication topology uses directed, time-varying exponential graphs and column-stochastic PushSum mixing, enabling sparse communication while preserving favorable averaging properties. Overlap SGP hides communication through delayed, asynchronous communication, but imperfect consensus can inject optimization noise.
- Communication topology: Each node cycles through exponentially distant peers, sending one message per iteration in the one-peer topology.
- Mixing matrices: The mixing matrices are column-stochastic, and each node independently chooses uniform weights for its out-neighbors.
- Averaging: After k = floor(log2(n − 1)) iterations, the directed exponential graph has λ2(P(k−1:0)) = 0, so all nodes exactly obtain the average.
- Optimization impact: Slower averaging leaves nodes less synchronized, causing gradients to be evaluated at different parameter points and injecting additional optimization noise.
- Overlap SGP: Overlap SGP is vanilla SGP with delayed communication, sending messages every τ iterations while allowing messages to arrive between intervals.
C. Implementation Details
The implementation separates local GPU computation from inter-node PushSum communication, using biased parameters with de-biasing for gradient computation and optional momentum or Adam adaptations. Experiments also examine node variability and consensus.
- System implementation: Each machine is treated as one node, with local mini-batch computation across eight GPUs using ALLREDUCE and inter-node averaging through PUSHSUM.
- System implementation: Separate gossip and computation threads overlap inter-node network I/O with local model training.
- SGP update flow: SGP aggregates model parameters and PushSum weights, de-biases parameters by dividing by the weight, applies gradients, and mixes both into the send buffer.
- Optimization settings: ImageNet experiments use Nesterov momentum with the SGD learning-rate schedule, while machine translation combines SGP with the Adam preconditioner.
- Consensus: During early ResNet-50 training, node validation errors vary, but the variability diminishes as learning rates decrease and nodes converge to similar errors.
D.3. SGP Scaling Analysis
SGP scales efficiently on both Ethernet and InfiniBand, with reported scaling efficiencies above 88%, and its throughput is more robust to high-latency interconnects than SGD.
- Scaling efficiency: SGP achieves 88.6% scaling efficiency on 10Gbit/s Ethernet and 92.4% on InfiniBand.
- Throughput comparison: SGP exhibits better scaling and greater robustness to high-latency interconnects than SGD.
- Delay modeling: Virtual nodes and edges model message delays, including a two-iteration delay from node 3 to node 0.
E. Proofs of Theoretical Guarantees
The proofs model bounded communication delays with virtual nodes and analyze SGP through augmented, column-stochastic mixing matrices. This matrix formulation supports bounds on consensus disagreement and the later convergence theorems.
- Modeling message delays: Bounded message delays are represented by adding virtual nodes that store transmitted but undelivered messages.With maximum delay τ, the augmented graph contains n(τ + 1) nodes.
- Modeling message delays: Virtual nodes forward messages along delay chains and are initialized with zero parameters and push-sum weights, so they do not change the final consensus value.They model in-transit communication rather than computation.
- Mixing matrices: The augmented mixing matrices remain column stochastic, preserving the key property required by the analysis.Column stochasticity means each column sums to 1.
- Matrix representation: The global SGP representation concatenates parameters, stochastic seeds, and gradients across virtual and non-virtual nodes.Virtual nodes use zero stochastic seeds and gradients because they do not perform gradient updates.
- Consensus bounds: The analysis bounds the distance between each non-virtual node’s de-biased parameters and the node-wise average under connectivity and bounded-delay assumptions.The assumptions include B-strong connectivity, delays bounded by τ, and lower-bounded nonzero mixing weights.
E.2. Towards the proof of the main Theorems
The convergence proof is organized around Lemma 8, which combines preliminary bounds and conditional-expectation arguments to establish the main theorem results. The final steps substitute these bounds, choose a stepsize, and simplify the resulting rates.
- Proof strategy: Lemma 8 is the central intermediate result used to prove Theorems 1 and 2.The convergence analysis builds the main theorem proofs from this lemma.
- Intermediate bounds: Preliminary lemmas bound quantities involving stochastic gradients, consensus terms, and expectations under Assumptions 1–3.The proof uses independence of nodes’ stochastic samples and the stated problem assumptions.
- Proof strategy: The proof derives bounds by conditioning on the history, taking expectations, and summing the resulting inequalities over iterations.The summation runs from k = 0 to K − 1.
- Main theorem proofs: The theorem proofs substitute the bounds from Lemma 8, remove a nonnegative left-hand-side term, and apply the prescribed stepsize and iteration lower bound.The remaining terms are absorbed into the leading term under the lower-bound condition on K.
- Main theorem proofs: The final expressions are reported using Big O notation, with constants absorbed into the asymptotic bound.The proof concludes after deriving these simplified expressions.