PipeFisher: Efficient Training of Large Language Models Using Pipelining and Fisher Information Matrices
This work addresses efficiency issues in distributed LLM training, offering incremental improvements by repurposing idle time for second-order optimization.
The paper tackles the problem of pipeline bubbles reducing accelerator utilization during large language model training by proposing PipeFisher, which uses Fisher information matrix-based optimization in idle bubbles to accelerate convergence, achieving 25-50% reduction in simulated training time for BERT models.
Pipeline parallelism enables efficient training of Large Language Models (LLMs) on large-scale distributed accelerator clusters. Yet, pipeline bubbles during startup and tear-down reduce the utilization of accelerators. Although efficient pipeline schemes with micro-batching and bidirectional pipelines have been proposed to maximize utilization, a significant number of bubbles cannot be filled using synchronous forward and backward passes. To address this problem, we suggest that extra work be assigned to the bubbles to gain auxiliary benefits in LLM training. As an example in this direction, we propose PipeFisher, which assigns the work of K-FAC, a second-order optimization method based on the Fisher information matrix, to the bubbles to accelerate convergence. In Phase 1 pretraining of BERT-Base and -Large models, PipeFisher reduces the (simulated) training time to 50-75% compared to training with a first-order optimizer by greatly improving the accelerator utilization and benefiting from the improved convergence by K-FAC.