A Neural Network-Based Pipeline Parallel Strategy Solver for Heterogeneous Environments

Abstract The escalating interest in language models has accentuated the significance of distributed training approaches, specifically pipeline parallelism. Most investigations primarily employ dynamic programming, these methods face inherent limitations in computational efficiency and do not offer a straightforward end-to-end resolution pathway. Furthermore, their applicability is significantly curtailed in heterogeneous environments. Addressing this challenge, we introduce a novel Neural Network-based Pipeline Parallel strategy solver (NNPiper) tailored for heterogeneous environments. NNPiper can perceive computational and communication costs, the number of stages to be divided, and the number of micro-batches, and it can directly provide the strategy for allocating specific devices to each pipeline stage, without requiring online training based on actual cluster data. Our unique neural network model and training method, the Virtual Contrastive Training Algorithm (VCTA), enables efficient training of NNPiper without collecting large amounts of real data. Compared with the state-of-the-art method, NNPiper can improve the theoretical training speed on average by 17.975%, in the completely heterogeneous environment for the Generative Pre-trained Transformer (GPT) models.