揭秘流水线并行:PipeDream 的首个理论
Demystifying Pipeline Parallelism: First Theory for PipeDream
AI总结:
本文通过引入随机化 PipeDream (RPD) 抽象,首次为 PipeDream 风格方法提供了非凸收敛保证,并分析了其稳态延迟与阶段数的缩放关系,同时与 LocalSGD 进行了比较。
摘要:
训练现代机器学习模型越来越需要跨多个加速器进行分布式计算。数据并行仍然是默认选择,并且通常与张量并行分片相结合,但一旦参数、激活或优化器状态不再适合单个设备,模型并行就变得不可避免。本文通过 PipeDream (PD) (Harlap et al., 2018) 的视角研究流水线模型并行。我们的第一个贡献是理论性的:我们引入了随机化 PipeDream (RPD),一种陈旧块-SGD 抽象,据我们所知,这为 PD 风格方法提供了第一个干净的非凸收敛保证。我们的第二个贡献是扩展诊断:我们证明了稳态 PD 引起的延迟随阶段数 S 增长为 $S^2 - S/2 + O(1)$,因此收敛定理中的陈旧读取贡献缩放为 $\Theta(\gamma^2 S^4)$,在调谐速率形式中等价于 $\Theta(S^4/K)$。我们的第三个贡献是与 LocalSGD 的比较,后者通过周期性模型平均来权衡权重陈旧性与同步气泡。在我们报告的模拟时间实验中,表现更好的方法取决于目标:PD 在二次目标和小型语言建模训练损失任务上表现更好,而对于逻辑回归,随着阶段数增加,LocalSGD 变得优越。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。