Adaptive DNN Partitioning Method for Pipeline Parallel Training with Time-Varying Resources

Pipeline parallelism provides a flexible and scalable approach for training large deep neural networks (DNNs). To realize efficient pipeline parallel training, deploying model partitions that adapt to the resource configurations is crucial. However, due to the time-varying nature prevalent in real-world applications, model partitions can quickly become outdated, resulting in imbalanced workload distribution. Consequently, maintaining load-balanced model partitions under time-varying resource scenarios poses a challenge. In this paper, we propose an adaptive partitioning method that dynamically adjusts model partitions in response to variations in available resources, ensuring high training efficiency. First, we derive the execution duration model and redistribution cost model of pipeline training systems. Then, we perform a joint optimization of the two factors with the goal of minimizing the training time. Based on this, we implement a prototype system called ARPipe, in which a Fast Partitioning Decision (FPD) algorithm is proposed to quickly determine optimized model partitions. Experimental evaluations show that under real-world circumstances with time-varying hardware resources, ARPipe can reduce training time by up to 30% in comparison to the state-of-the-art method GPipe when training the ResNet50 model on the CIFAR10 dataset.