[Paper Review] PipeFisher: Efficient Training of Large Language Models Using Pipelining and Fisher Information Matrices

Kazuki Osawa, Shigang Li+1 more|arXiv (Cornell University)|Nov 25, 2022

Topic Modeling4 citations

TL;DR

PipeFisher proposes a novel training scheme that leverages idle time (bubbles) in pipeline-parallel LLM training to execute K-FAC, a second-order optimization method based on Fisher information matrices. By offloading K-FAC's curvature computation and matrix inversion to these idle periods, PipeFisher boosts GPU utilization from 59.8% to 97.6% and reduces simulated Phase 1 pretraining time for BERT-Large by 75.7%, while also improving convergence through better optimization.

ABSTRACT

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.

Motivation & Objective

• To address low accelerator utilization in pipeline-parallel training of large language models (LLMs), which results from idle 'bubbles' during startup and teardown.
• To exploit these idle periods by assigning meaningful computational work that provides auxiliary benefits beyond mere utilization gains.
• To demonstrate that second-order optimization via K-FAC can be efficiently integrated into pipeline schedules without increasing communication overhead.
• To show that the combined effect of improved utilization and faster convergence via K-FAC leads to significant reductions in training time.

Proposed method

• PipeFisher integrates K-FAC's curvature computation and matrix inversion into the idle bubbles of a pipeline-parallel training schedule, such as GPipe or Chimera.
• It schedules K-FAC work—specifically, the computation of Fisher information matrices (A_l, B_l) and their inversion—during pipeline bubbles to avoid blocking forward and backward passes.
• The method uses stale inverse matrices from previous steps for the first preconditioning, ensuring convergence stability during early training.
• It leverages existing pipeline scheduling frameworks and requires only a minimal computational overhead (~6.5% per step) beyond standard pipeline training.
• The approach is compatible with any pipeline-parallel scheme and can be applied to both Transformer and non-Transformer architectures, though load balancing is more complex for non-uniform layers.
• It enables efficient hyperparameter search by reducing the cost of evaluating different learning rate schedules due to faster training cycles.

Experimental results

Research questions

• RQ1 Can idle time (bubbles) in pipeline-parallel LLM training be effectively utilized to run computationally intensive optimization methods?
• RQ2 Does offloading K-FAC computation to pipeline bubbles improve overall training efficiency and convergence speed?
• RQ3 To what extent can GPU utilization be increased by assigning auxiliary work to bubbles without increasing communication costs?
• RQ4 How does the integration of K-FAC via PipeFisher affect training time and model performance in BERT pretraining?

Key findings

• PipeFisher increases GPU utilization from 59.8% to 97.6% in BERT-Large pretraining using Chimera pipeline scheduling, significantly reducing idle time.
• Simulated training time for BERT-Large Phase 1 is reduced from 275.1 to 208.3 minutes—a 75.7% improvement—due to higher utilization and faster convergence from K-FAC.
• For BERT-Base, PipeFisher reduces training time to 50% of the baseline, demonstrating consistent gains across model sizes.
• The computational overhead of PipeFisher is minimal (~6.5% per step), with preconditioning being the only additional cost beyond standard pipeline training.
• The method enables faster hyperparameter search by reducing the time required to evaluate different learning rate schedules.
• PipeFisher is generalizable to other second-order optimizers like Shampoo and can be extended to other auxiliary tasks such as Sharpness-Aware Minimization (SAM).