ICML 2024 | 梯度检查点太慢？不降速、省显存，LowMemoryBP大幅提升反向传播显存效率

AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。

---

如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本文论文一作是南开大学统计与数据科学学院研二硕士生杨雨辰，指导老师为南开大学统计与数据科学学院的徐君副教授。徐君老师团队的研究重点是计算机视觉、生成式AI和高效机器学习，并在顶级会议和期刊上发表了多篇论文，谷歌学术引用超过4700次。

自从大型Transformer模型逐渐成为各个领域的统一架构，微调就成为了将预训练大模型应用到下游任务的重要手段。然而，由于模型的尺寸日益增大，微调所需要的显存也逐渐增加，如何高效地降低微调显存就成了一个重要的问题。此前，微调Transformer模型时，为了节省显存开销，通常的做法是使用梯度检查点（gradient checkpointing，也叫作激活重算），以牺牲训练速度为代价降低反向传播（Backpropagation, BP）过程中的激活显存占用。

最近，由南开大学统计与数据科学学院徐君老师团队发表在ICML 2024上的论文《Reducing Fine-Tuning Memory Overhead by Approximate and Memory-Sharing Backpropagation》提出通过更改反向传播（BP）过程，在不增加计算量的情况下，显著减少峰值激活显存占用。

• 论文：Reducing Fine-Tuning Memory Overhead by Approximate and Memory-Sharing Backpropagation
• 论文链接：https://arxiv.org/abs/2406.16282
• 项目链接：https://github.com/yyyyychen/LowMemoryBP

文章提出了两种反向传播改进策略，分别是Approximate Backpropagation（Approx-BP）和Memory-Sharing Backpropagation（MS-BP）。Approx-BP和MS-BP分别代表了两种提升反向传播中内存效率的方案，可以将其统称为LowMemoryBP。无论是在理论还是实践意义上，文章都对更高效的反向传播训练提供了开创性的指导。

在理论显存分析中，LowMemoryBP可以大幅降低来自激活函数和标准化层的激活显存占用，以ViT和LLaMA为例，可以对ViT微调降低39.47%的激活显存，可以对LLaMA微调降低29.19%的激活显存。

在实际实验中，LowMemoryBP可以有效地使包括ViT, LLaMA, RoBERTa, BERT, Swin在内的Transformer模型微调峰值显存占用降低20%~30%，并且不会带来训练吞吐量和测试精度的损失。

Approx-BP

在传统反向传播训练中，激活函数梯度的反向回传是严…

---

---

感谢阅读！如果您对AI的更多资讯感兴趣，可以查看更多AI文章：GPTNB。