torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models

• DOI: 10.48550/arxiv.2004.09910
• Publication Date: 2020-01-01
• Authors: Chiheon Kim; Heungsub Lee; Myungryong Jeong; Woonhyuk Baek; Boogeon Yoon; Ildoo Kim; Sungbin Lim; Sungwoong Kim
• Keywords: On the fly

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 set of components to enable pipeline-parallel gradient computation PyTorch's define-by-run eager execution environment. show that each component is necessary fully benefit from such environment, demonstrate the efficiency applying it various network architectures including AmoebaNet-D U-Net. Our available at https://github.com/kakaobrain/torchgpipe .