DawnPiper: A Memory-scablable Pipeline Parallel Training Framework

TL;DR

DawnPiper is a novel pipeline parallel training framework that optimizes memory usage and model partitioning, enabling larger models and faster training speeds for large-scale neural networks.

Contribution

It introduces a DL compilation-based profiling method and a performance-optimal theorem for memory-efficient pipeline partitioning, improving scalability and efficiency.

Findings

Up to 4x increase in trainable batch size over vPipe

Up to 11x increase over PipeDream

1.5x performance speedup compared to vPipe

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.