PipeOptim: Ensuring Effective 1F1B Schedule With Optimizer-Dependent Weight Prediction

Asynchronous pipeline model parallelism with a “1F1B” (one forward, one backward) schedule generates little bubble overhead and always provides quite a high throughput. However, the “1F1B” schedule inevitably leads to weight inconsistency and weight staleness issues due to the cross-training of different mini-batches across GPUs. To simultaneously address these two problems, in this paper, we propose an optimizer-dependent weight prediction strategy (a.k.a PipeOptim) for asynchronous pipeline training. The key insight of our proposal is that we employ a weight prediction strategy in the forward pass to approximately ensure that each mini-batch uses consistent and staleness-free weights to compute the forward pass of the “1F1B” schedule. To be concrete, we first construct the weight prediction scheme based on the update rule of the used optimizer when training the deep neural network models. Then throughout the “1F1B” pipeline training, each mini-batch is mandated to execute weight prediction, subsequently employing the predicted weights to perform the forward pass. As a result, PipeOptim 1) inherits the advantage of the “1F1B” schedule and generates high throughput, and 2) can ensure effective parameter learning regardless of the type of the used optimizer. We conducted extensive experimental evaluations using nine different deep-learning models to verify the effectiveness of our proposal. The experiment results demonstrate that PipeOptim outperforms the other five popular pipeline approaches including GPipe, PipeDream, PipeDream-2BW, SpecTrain, and XPipe.