torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models
By Chiheon Kim, Heungsub Lee, Myungryong Jeong, Woonhyuk Baek, Boogeon Yoon, Ildoo Kim, Sungbin Lim, Sungwoong Kim. 2020
63 citations recorded by the source corpus.
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 this https URL .
Subfields: Computer Vision and Pattern Recognition, Efficiency & Compression
How researchers see this paper
The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.
In Computer Vision and Pattern Recognition
- Sam Shleifer
- S. Avestimehr
- Chien-chin Huang
- Chuan Wu
- Min Xu
- M. Soltanolkotabi
- Alban Desmaison
- Chaoyang He
In Efficiency & Compression
Citations in the corpus
- PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023, background): “Pipeline parallelism [5, 7, 10, 14, 19] partitions a model instance into stages and distributes stages across multiple devices, where activations and gradients are communicated across stage boundaries.”
- DAPPLE: a pipelined data parallel approach for training large models (2020, methodology): “GPipe[24, 31] explores synchronous pipeline approach to train largemodels with limited GPUmemory.”
- PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers (2021, methodology): “To support PipeTransformer ’s elastic pipelining, we maintain a customized version of PyTorch Pipe (Kim et al., 2020).”