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

2023-12-01

arXiv.org

AI Summary

In asynchronous pipeline parallelism, the 1F1B (one-forward-one-backward) scheduling incurs weight inconsistency and staleness due to interleaved mini-batch execution across GPUs. To address this, we propose an optimizer-aware forward-weight prediction mechanism: leveraging optimizer-specific update rules to dynamically model weight evolution and accurately predict—prior to forward propagation—the most up-to-date weights for each mini-batch. This is the first approach to strictly guarantee weight consistency and zero staleness per mini-batch under 1F1B, while remaining compatible with arbitrary optimizers. Our method integrates asynchronous scheduling, forward-weight precomputation, and delayed gradient synchronization. Extensive experiments across eight models and three task categories demonstrate throughput comparable to GPipe and PipeDream, while achieving convergence accuracy on par with fully synchronous baselines.

Abstract

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 ensure that each mini-batch uses consistent and staleness-free weights to compute the forward pass. 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"pipelined training, each mini-batch is mandated to execute weight prediction ahead of the forward pass, subsequently employing the predicted weights to perform the forward pass. As a result, PipeOptim 1) inherits the advantage of the"1F1B"schedule and generates pretty high throughput, and 2) can ensure effective parameter learning regardless of the type of the used optimizer. To verify the effectiveness of our proposal, we conducted extensive experimental evaluations using eight different deep-learning models spanning three machine-learning tasks including image classification, sentiment analysis, and machine translation. The experiment results demonstrate that PipeOptim outperforms the popular pipelined approaches including GPipe, PipeDream, PipeDream-2BW, and SpecTrain. The code of PipeOptim can be accessible at https://github.com/guanleics/PipeOptim.

Problem

• Addresses weight inconsistency in 1F1B schedules
• Prevents weight staleness across GPUs in pipelines
• Ensures effective parameter learning in asynchronous training

Innovation

• Weight prediction strategy
• Asynchronous pipeline training
• Consistent weight usage