Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DawnPiper: A Memory-scablable Pipeline Parallel Training Framework

Xuan Peng, Xuanhua Shi, Haolin Zhang, Yunfei Zhao, Xuehai Qian

TL;DR

DawnPiper tackles memory waste in pipeline parallel training by addressing memory imbalance across stages. It combines DL compilation-based profiling to produce a fine-grained computation graph, a partition theorem to constrain search space, a binary partitioning algorithm, and Capuchin-inspired memory optimization to swap or recompute within memory limits. The framework achieves substantial improvements in trainable batch size (up to 4× vs vPipe and 11× vs PipeDream) and up to 1.5× speedups over vPipe on various models and configurations, including CNNs and transformers, in both synchronous and asynchronous modes. These results demonstrate memory-efficient, scalable pipeline training that better utilizes GPU memory across diverse architectures.

Abstract

Pipeline parallelism is a crucial paradigm for large-scale model training. However, imbalances in memory footprint across stages can lead to significant GPU memory wastage, limiting the model sizes that pipeline parallelism can effectively support. In this paper, we introduce DawnPiper, a memory-scalable pipeline parallel training framework. Firstly, we develop a DL compilation-based profiling method that transforms the model into a fine-grained computation graph. This refinement gives us a finer granularity of model partitioning and memory optimization while facilitating automatic code generation. Based on observed memory usage characteristics, we derive a performance-optimal theorem for pipeline parallel partitioning that substantially reduces the partition search space. Secondly, we propose a binary pipeline partitioning algorithm and utilize a cost-model based memory optimization approach to efficiently identify nearly optimal pipeline parallel strategy. DawnPiper achieves up to a 4x and 11x increase in trainable maximum batch size compared to vPipe and PipeDream, respectively, and provides up to a 1.5x performance speedup compared to vPipe.

DawnPiper: A Memory-scablable Pipeline Parallel Training Framework

TL;DR

DawnPiper tackles memory waste in pipeline parallel training by addressing memory imbalance across stages. It combines DL compilation-based profiling to produce a fine-grained computation graph, a partition theorem to constrain search space, a binary partitioning algorithm, and Capuchin-inspired memory optimization to swap or recompute within memory limits. The framework achieves substantial improvements in trainable batch size (up to 4× vs vPipe and 11× vs PipeDream) and up to 1.5× speedups over vPipe on various models and configurations, including CNNs and transformers, in both synchronous and asynchronous modes. These results demonstrate memory-efficient, scalable pipeline training that better utilizes GPU memory across diverse architectures.

Abstract

Pipeline parallelism is a crucial paradigm for large-scale model training. However, imbalances in memory footprint across stages can lead to significant GPU memory wastage, limiting the model sizes that pipeline parallelism can effectively support. In this paper, we introduce DawnPiper, a memory-scalable pipeline parallel training framework. Firstly, we develop a DL compilation-based profiling method that transforms the model into a fine-grained computation graph. This refinement gives us a finer granularity of model partitioning and memory optimization while facilitating automatic code generation. Based on observed memory usage characteristics, we derive a performance-optimal theorem for pipeline parallel partitioning that substantially reduces the partition search space. Secondly, we propose a binary pipeline partitioning algorithm and utilize a cost-model based memory optimization approach to efficiently identify nearly optimal pipeline parallel strategy. DawnPiper achieves up to a 4x and 11x increase in trainable maximum batch size compared to vPipe and PipeDream, respectively, and provides up to a 1.5x performance speedup compared to vPipe.
Paper Structure (26 sections, 1 theorem, 2 equations, 9 figures, 2 tables, 2 algorithms)

This paper contains 26 sections, 1 theorem, 2 equations, 9 figures, 2 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Data Parallelism
  4. Pipeline and Hybrid Parallelism
  5. Challeges and Motivations
  6. Memory Imbalance in Pipeline Parallelism
  7. Memory Usage Characteristics in DL Training
  8. Design of DawnPiper
  9. Pipeline Terminology and Partition Theorem
  10. Binary Pipeline Parallel Partition
  11. Memory Optimization
  12. Evaluation
  13. Experimental setup
  14. Memory Efficiency
  15. Synchronous Pipeline Parallel Training
  16. ...and 11 more sections

Key Result

Theorem 4.1

For a model that needs to be divided into two pipeline stages, the optimal partition point will lie within the closed interval between the compute-balanced and memory-balanced positions, provided the following three conditions are met: 1) The computation time and memory usage during forward propagat

Figures (9)

  • Figure 1: Synchronous and Asynchronous Pipeline Parallel
  • Figure 2: Peak Memory Usage on GPipe (4 GPUs)
  • Figure 3: Peak Memory Usage on PipeDream (8 GPUs)
  • Figure 4: CDF of Node's Activation and Consumed Memory
  • Figure 5: DawnPiper System Architecture
  • ...and 4 more figures

Theorems & Definitions (2)

  • Theorem 4.1
  • proof