[论文解读] 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 Modeling被引用 4

一句话总结

PipeFisher 提出了一种新颖的训练方案，利用流水线并行大语言模型（LLM）训练中的空闲时间（气泡期）执行 K-FAC，一种基于费舍尔信息矩阵的二阶优化方法。通过将 K-FAC 的曲率计算和矩阵求逆操作卸载到这些空闲时段，PipeFisher 将 GPU 利用率从 59.8% 提升至 97.6%，并在模拟中将 BERT-Large 的 Phase 1 预训练时间减少了 75.7%，同时通过更优的优化提升了收敛速度。

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.

研究动机与目标

• 为解决大语言模型（LLM）流水线并行训练中因启动和终止阶段产生的空闲‘气泡’导致的加速器利用率低下问题。
• 通过在这些空闲时段分配有意义的计算任务，实现超越单纯利用率提升的附加效益。
• 证明通过 K-FAC 实现的二阶优化可高效集成到流水线调度中，且不增加通信开销。
• 展示通过提升利用率与 K-FAC 加速收敛的协同效应，可显著缩短训练时间。

提出的方法

• PipeFisher 将 K-FAC 的曲率计算与矩阵求逆操作整合到流水线并行训练调度的空闲气泡中，例如 GPipe 或 Chimera。
• 它在流水线气泡期间调度 K-FAC 任务——具体为费舍尔信息矩阵（A_l, B_l）的计算及其求逆——以避免阻塞前向和反向传播。
• 在早期训练中，使用前一阶段的过时逆矩阵进行首次预条件处理，以确保收敛稳定性。
• 该方法复用现有的流水线调度框架，仅在标准流水线训练基础上增加约 6.5% 的计算开销（每步）。
• 该方法适用于任何流水线并行方案，可应用于 Transformer 及非 Transformer 架构，但非均匀层的负载均衡更具挑战。
• 通过缩短训练周期，实现高效的超参数搜索，降低评估不同学习率调度方案的成本。

实验结果

研究问题

• RQ1 流水线并行 LLM 训练中的空闲时间（气泡期）能否被有效利用来运行计算密集型优化方法？
• RQ2 将 K-FAC 计算卸载到流水线气泡中是否能提升整体训练效率与收敛速度？
• RQ3 在不增加通信成本的前提下，通过向气泡期分配辅助计算任务，GPU 利用率最多可提升多少？
• RQ4 通过 PipeFisher 集成 K-FAC 对 BERT 预训练的训练时间与模型性能有何影响？

主要发现

• 在使用 Chimera 流水线调度的 BERT-Large 预训练中，PipeFisher 将 GPU 利用率从 59.8% 提升至 97.6%，显著减少了空闲时间。
• 由于更高的利用率和 K-FAC 带来的更快收敛，BERT-Large Phase 1 的模拟训练时间从 275.1 分钟减少至 208.3 分钟，实现 75.7% 的性能提升。
• 对于 BERT-Base，PipeFisher 将训练时间减少至基线的 50%，表明在不同模型规模下均具有一致的性能增益。
• PipeFisher 的计算开销极低（每步约 6.5%），且仅在预条件处理阶段引入额外成本，其余部分与标准流水线训练一致。
• 该方法通过缩短评估不同学习率调度方案所需的时间，实现了更快的超参数搜索。
• PipeFisher 可推广至其他二阶优化器（如 Shampoo），并可扩展至其他辅助任务（如 Sharpness-Aware Minimization, SAM）。