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

TL;DR

This paper introduces torchgpipe, a PyTorch library enabling efficient micro-batch pipeline parallelism with checkpointing, facilitating training of large models in define-by-run environments.

Contribution

The paper presents a set of design components for pipeline-parallel gradient computation in PyTorch's eager mode, enabling practical large-model training.

Findings

Efficient pipeline parallelism demonstrated on AmoebaNet-D and U-Net.

Library is ready-to-use and publicly available.

Achieves improved training efficiency for giant models.

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 .