Search papers, labs, and topics across Lattice.
This paper introduces Randomized PipeDream (RPD), a novel theoretical framework for pipeline model parallelism that provides the first nonconvex convergence guarantee for PipeDream-style methods. It also analyzes the scaling behavior of steady-state PipeDream, revealing that the delay grows quadratically with the number of stages, which has implications for convergence rates. In simulated experiments, the performance of RPD and LocalSGD varies based on the training objective, with RPD excelling in certain scenarios while LocalSGD outperforms it in others as stages increase.
The first clean nonconvex convergence guarantee for pipeline parallelism reveals surprising trade-offs between methods depending on the training objective.
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.
