PredTOP: Latency Predictor Utilizing DAG Transformers for Distributed Deep Learning Training with Operator Parallelism

With the increasing sizes of deep learning (DL) models, distributed training with various parallelization techniques, such as pipeline, model, and tensor parallelism, has become an absolute necessity. Accurately and efficiently predicting iteration latency for distributed DL training is critical for high-quality application scheduling, resource allocation, and automatic performance optimization. Unfortunately, existing latency predictors for distributed DL training mainly consider data parallelism, but fail to be applied for DL training optimized by other parallelization techniques. Also, the latency prediction of large-scale distributed DL training with hybrid parallelism is very challenging due to the extremely high measurement cost for one-time execution and the prohibitively huge search space of combinations from multiple parallelization techniques. In this paper, we propose PredTOP, a framework to accurately and efficiently predict the iteration latency of distributed DL training with different types of parallelization techniques. First, we numerically formulate the end-to-end time of pipeline parallelism using stage latency to leverage the low cost of white-box modeling. Then, we adopt a black-box modeling technique to predict the latency of each stage in pipeline parallelism. Specifically, we utilize Transformers over directed acyclic graphs to improve prediction accuracy. The performance superiority of our PredTOP is validated by extensive experimental results with two real-world benchmarks in comparison with state-of-the-art approaches and illustrated in a realistic automatic parallelization plan generation system.