[论文解读] torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models

Chiheon Kim, Heungsub Lee+6 more|arXiv (Cornell University)|Apr 21, 2020

Advanced Neural Network Applications参考文献 26被引用 32

一句话总结

一个 PyTorch 库，实现 GPipe 风格的微批处理流水线并行性，带有检查点和若干优化组件，以在 eager 执行环境中实现对超大模型的高效训练。

ABSTRACT

We design and implement a ready-to-use library in PyTorch for performing micro-batch pipeline parallelism with checkpointing proposed by GPipe (Huang et al., 2019). In particular, we develop a set of design components to enable pipeline-parallel gradient computation in PyTorch's define-by-run and eager execution environment. We show that each component is necessary to fully benefit from pipeline parallelism in such environment, and demonstrate the efficiency of the library by applying it to various network architectures including AmoebaNet-D and U-Net. Our library is available at https://github.com/kakaobrain/torchgpipe .

研究动机与目标

• Motivate and enable training of massive neural networks beyond single-device capacity.
• Provide a ready-to-use PyTorch library for GPipe-style pipeline parallelism with checkpointing.
• Demonstrate the necessity of several optimization components to achieve efficient pipeline parallelism in define-by-run/eager PyTorch.
• Show performance gains on architectures like AmoebaNet-D and U-Net.
• Discuss handling of non-sequential models and skip connections in pipeline parallelism.

提出的方法

• Represent the network as a sequence of partitions with disjoint parameter sets.
• Apply micro-batch pipeline parallelism with gradient checkpointing to reduce memory usage.
• Introduce a deterministic clock-cycle to schedule forward tasks across devices.
• Provide backward dependency via Fork and Join to enforce correct autograd graph for backward pass.
• Use non-default CUDA streams to enable concurrent copy and computation.
• Introduce portals and shared-memory autograd functions for skip connections and checkpointed recomputation.

实验结果

研究问题

• RQ1 How can GPipe-like pipeline parallelism be realized efficiently in PyTorch's define-by-run/eager mode?
• RQ2 What components are essential to achieve high throughput and memory efficiency when training giant models with pipeline parallelism?
• RQ3 How can non-sequential models and skip connections be accommodated in pipeline-parallel training?
• RQ4 What is the impact of the proposed components on throughput, device utilization, and memory usage?
• RQ5 Can torchgpipe achieve competitiveness with GPipe on large architectures like AmoebaNet-D and U-Net in PyTorch?

主要发现

• Each optimization component (deterministic clock-cycle, backward Fork/Join dependencies, non-default streams, and portals) provides measurable speedups.
• Combined components yield nearly a twofold speedup over the baseline in the reported U-Net experiments.
• Non-default streams enable concurrent copy and computation, improving utilization.
• Portals reduce unnecessary copies when skip connections are present, lowering memory pressure and improving timeline efficiency.
• AmoebaNet-D and U‑Net benchmarks show throughput and memory characteristics consistent with effective pipeline parallelism in PyTorch.