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

2020/04/21 by Chiheon Kim, Kim, Chiheon, Heungsub Lee +13 · 1 citation
Computer Science · #Advanced Neural Network Applications #Distributed #FOS: Computer and information sciences #Generative Adversarial Networks and Image Synthesis #Machine Learning (cs.LG) #Multimodal Machine Learning Applications #Parallel #and Cluster Computing (cs.DC)

paper · pdf · doi:10.48550/arxiv.2004.09910

openalex publication_date 2020/04/21 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

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 .

Citations

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• ImageNet classification with deep convolutional neural networks
• Training Deep Nets with Sublinear Memory Cost
• MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications
• Device Placement Optimization with Reinforcement Learning
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training
• ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware
• XPipe: Efficient Pipeline Model Parallelism for Multi-GPU DNN Training

Cited by

• Distributed Deep Learning using Stochastic Gradient Staleness

Related

• PyTorch Distributed: Experiences on Accelerating Data Parallel Training
• Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Models
• SimpleFSDP: Simpler Fully Sharded Data Parallel with torch.compile
• BitPipe: Bidirectional Interleaved Pipeline Parallelism for Accelerating Large Models Training
• GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism