BitPipe: Bidirectional Interleaved Pipeline Parallelism for Accelerating Large Models Training

TL;DR

Experiments show that BitPipe improves the training throughput of GPT-style and BERT-style models by 1.05x-1.28x compared to the state-of-the-art synchronous approaches.

Abstract

With the increasing scale of models, the need for efficient distributed training has become increasingly urgent. Recently, many synchronous pipeline parallelism approaches have been proposed to improve training throughput. However, these approaches still suffer from two major issues, i.e., pipeline bubbles caused by periodic flushing and extra communication due to the increasing number of pipeline stages. To this end, we propose BitPipe, a bidirectional interleaved pipeline parallelism for accelerating large models training. Specifically, a hybrid scheme of fusing interleaved pipelines with bidirectional pipelines is proposed to reduce the computational time of each single micro-batch and multiply the number of devices executing simultaneously. A V-shaped schedule with eager gradient synchronization is introduced to reduce and overlap the communication between devices. Experiments conducted on up to 32 GPUs show that BitPipe improves the training throughput of GPT-style and BERT-style models by 1.05x-1.28x compared to the state-of-the-art synchronous approaches. The code of our implementation is available at https://github.com/wuhouming/BitPipe.

Figures (13)

• Figure 1: Classic synchronous pipeline schedules, with 4 pipeline devices and 8 micro-batches within a training iteration. Both schedules have the same bubble overhead and weights memory consumption ($M_\uptheta$). The activations memory consumption ($M_{\rm a}$) of the 1F1B schedule exhibits better efficiency but existing imbalance.
• Figure 2: Synchronous approaches considered in this paper, with 4 pipeline devices and 4 micro-batches within a training iteration. Dark colors show the first stage and light colors show the second stage. In Chimera, each device is responsible for 2 pipelines in different directions (black text colors represent the down pipeline and white text colors for the up pipeline).
• Figure 3: Model chunks and bidirectional interleaved pipelines scheduling of BitPipe, with 4 pipeline devices and 4 micro-batches within a training iteration.
• Figure 4: Interleaved pipeline schedules, with 2 pipeline devices and 2 micro-batches for simplicity.
• Figure 5: Overlap communication by eager gradient synchronization.