本文论文一作是南开大学统计与数据科学学院研二硕士生杨雨辰,教学诚实为南开大学统计与数据科学学院的徐君副栽培。徐君诚实团队的参议重心是臆度机视觉、生成式 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
在传统反向传播磨练中,激活函数梯度的反向回传是严格对应其导函数的,关于 Transformer 模子中常用的 GELU 和 SiLU 函数,这意味着需要将输入特征张量竣工地存入激活显存中。而本文的作家提议了一套反向传播一样表面,即 Approx-BP 表面。在该表面的教学下,作家使用分段线性函数贴近激活函数,并用分段线性函数的导数(路线函数)替代 GELU/SiLU 梯度的反向回传。这个关节导出了两个非对称的内存高效激活函数:ReGELU2 和 ReSiLU2。这类激活函数由于使用 4 段路线函数进行反向回传,从而使得激活存储只需要使用 2bit 数据类型。
MS-BP
BP 网罗每一层常常齐会将输入张量存入激活显存以用作反向传播臆度。作家指出要是不错将某一层的反向传播改写成依赖输出的表情,那么这一层和后一层就不错分享吞并个激活张量,从而裁汰激活存储的冗余。
而著述指出 Transformer 模子中常用的 LayerNorm 和 RMSNorm,在将仿射参数合并到后一层的线性层之后,不错很好地允洽 MS-BP 战略的条件。经过再行策画的 MS-LayerNorm 和 MS-RMSNorm 不再产生空闲的激活显存。
本质放肆
作家对臆度机视觉和当然道话搞定限制的多少个代表模子进行了微调本质。其中,在 ViT,LLaMA 和 RoBERTa 的微调本质中,著述提议的关节分辩将峰值显存占用裁汰了 27%,29% 和 21%,何况莫得带来磨练效力和磨练速率的赔本。详确到,看成对比的 Mesa(一个 8-bit Activation Compressed Training 关节)使磨练速率裁汰了约 20%,而著述提议的 LowMemoryBP 关节则弥漫保执了磨练速率。
论断及有趣
著述提议的两种 BP 调动战略,Approx-BP 和 MS-BP,均在保执磨练效力和磨练速率的同期,达成了激活显存的显耀大要。这意味着从 BP 旨趣上进行优化曲直常有远景的显存大要决策。此外,著述提议的 Approx-BP 表面打破了传统神经网罗的优化框架,为使用非配对导数提供了表面可行性。其导出的 ReGELU2 和 ReSiLU2 展现了这一作念法的热切实施价值。
接待寰宇阅读论文或者代码去了解算法的详备细节,LowMemoryBP 名堂标 github 仓库上还是开源联系的模块。