Mind Bubbles and Memory: Bounds on Scheduling Pipeline Parallelism with Rematerialization
Résumé
Training large neural networks, especially Transformer-based Large Language Models (LLMs), requires massive high-performance computing (HPC) resources. Within each microbatch, computations follow a strictly sequential flow through a stack of transformer blocks: a forward pass to compute the loss, and a backward pass to propagate gradients. This sequential structure limits intrinsic parallelism. To improve performance, several complementary strategies have been developed: data, tensor, sequence, and pipeline parallelism, typically combined to achieve scalability over tens of thousands of GPUs.
This paper presents a formal analysis of pipeline parallelism (PP) for large-scale training. In PP, the model is partitioned into multiple stages, and microbatches are injected into the pipeline to overlap computation. The main challenge is to minimize idle periods (pipeline bubbles) while managing memory usage, since each GPU must store intermediate activations from multiple in-flight microbatches. Existing scheduling algorithms such as GPIPE, 1F1B, HANAYO, and MEGATRON reduce idle time but lack formal lower bounds or explicit modeling of memory constraints.
We develop a unified analytical approach for PP scheduling, deriving lower bounds on completion time for both single-wave and multi-wave regimes. Our analysis explicitly incorporates a memory constraint K, denoting the number of activations that can be stored per GPU. Exact results are provided for two extreme cases (minimal memory (K = 1) and large memory (K ≥ m)), while general lower bounds are established for intermediate configurations. Our analysis highlights the intrinsic coupling between pipeline utilization and memory footprint, providing a foundation for evaluating and comparing pipeline scheduling algorithms under realistic memory constraints.
| Origine | Fichiers produits par l'(les) auteur(s) |
|---|---|
| Licence |