๐ค AI Summary
This work addresses the lack of non-convex convergence guarantees in PipeDream-style pipeline parallelism by proposing the Randomized PipeDream (RPD) framework, for which it establishes the first rigorous non-convex convergence theory. By introducing a randomized block SGD abstraction coupled with explicit modeling of communication delays, the analysis reveals that under steady-state conditions, the delay grows quadratically with the number of pipeline stages \(S\), leading to stale gradient terms scaling as \(\Theta(S^4)\). Empirical evaluations demonstrate that RPD outperforms LocalSGD in quadratic optimization and small-scale language model training, whereas LocalSGD exhibits superior performance as \(S\) increases in logistic regression tasks, highlighting a nuanced trade-off between the two methods across different problem settings.
๐ Abstract
Training modern machine learning models increasingly requires computation to be distributed across many accelerators. Data parallelism remains the default choice and is often paired with tensor-parallel sharding, but model parallelism becomes unavoidable once parameters, activations, or optimizer states no longer fit on a single device. This paper studies pipeline model parallelism through the lens of PipeDream (PD) (Harlap et al., 2018). Our first contribution is theoretical: we introduce Randomized PipeDream (RPD), a stale block-SGD abstraction that yields, to our knowledge, the first clean nonconvex convergence guarantee for a PD-style method. Our second contribution is a scaling diagnosis: we prove that the delay induced by steady-state PD grows as $S^2 - S/2 + O(1)$ for $S$ stages, so the stale-read contribution in the convergence theorem scales as $ฮ(ฮณ^2 S^4)$, equivalently as $ฮ(S^4/K)$ in the tuned-rate form. Our third contribution is a comparison with LocalSGD, whose periodic model averaging trades weight staleness for synchronization bubbles. In our reported simulated-time experiments, the better-performing method depends on the objective: PD performs better on the quadratic objective and on a small language-modeling training-loss task, while for logistic regression LocalSGD becomes superior as the number of stages increases.