本文論文一作是南開(kāi)大學(xué)統(tǒng)計(jì)與數(shù)據(jù)科學(xué)學(xué)院研二碩士生楊雨辰，指導(dǎo)老師為南開(kāi)大學(xué)統(tǒng)計(jì)與數(shù)據(jù)科學(xué)學(xué)院的徐君副教授。徐君老師團(tuán)隊(duì)的研究重點(diǎn)是計(jì)算機(jī)視覺(jué)、生成式 AI 和高效機(jī)器學(xué)習(xí)，并在頂級(jí)會(huì)議和期刊上發(fā)表了多篇論文，谷歌學(xué)術(shù)引用超過(guò) 4700 次。

自從大型 Transformer 模型逐漸成為各個(gè)領(lǐng)域的統(tǒng)一架構(gòu)，微調(diào)就成為了將預(yù)訓(xùn)練大模型應(yīng)用到下游任務(wù)的重要手段。然而，由于模型的尺寸日益增大，微調(diào)所需要的顯存也逐漸增加，如何高效地降低微調(diào)顯存就成了一個(gè)重要的問(wèn)題。此前，微調(diào) Transformer 模型時(shí)，為了節(jié)省顯存開(kāi)銷，通常的做法是使用梯度檢查點(diǎn)（gradient checkpointing，也叫作激活重算），以犧牲訓(xùn)練速度為代價(jià)降低反向傳播（Backpropagation, BP）過(guò)程中的激活顯存占用。

最近，由南開(kāi)大學(xué)統(tǒng)計(jì)與數(shù)據(jù)科學(xué)學(xué)院徐君老師團(tuán)隊(duì)發(fā)表在 ICML 2024 上的論文《Reducing Fine-Tuning Memory Overhead by Approximate and Memory-Sharing Backpropagation》提出通過(guò)更改反向傳播（BP）過(guò)程，在不增加計(jì)算量的情況下，顯著減少峰值激活顯存占用。

• 論文：Reducing Fine-Tuning Memory Overhead by Approximate and Memory-Sharing Backpropagation
• 論文鏈接：https://arxiv.org/abs/2406.16282
• 項(xiàng)目鏈接：https://github.com/yyyyychen/LowMemoryBP

文章提出了兩種反向傳播改進(jìn)策略，分別是 Approximate Backpropagation（Approx-BP）和 Memory-Sharing Backpropagation（MS-BP）。Approx-BP 和 MS-BP 分別代表了兩種提升反向傳播中內(nèi)存效率的方案，可以將其統(tǒng)稱為 LowMemoryBP。無(wú)論是在理論還是實(shí)踐意義上，文章都對(duì)更高效的反向傳播訓(xùn)練提供了開(kāi)創(chuàng)性的指導(dǎo)。

在理論顯存分析中，LowMemoryBP 可以大幅降低來(lái)自激活函數(shù)和標(biāo)準(zhǔn)化層的激活顯存占用，以 ViT 和 LLaMA 為例，可以對(duì) ViT 微調(diào)降低 39.47% 的激活顯存，可以對(duì) LLaMA 微調(diào)降低 29.19% 的激活顯存。

在實(shí)際實(shí)驗(yàn)中，LowMemoryBP 可以有效地使包括 ViT, LLaMA, RoBERTa, BERT, Swin 在內(nèi)的 Transformer 模型微調(diào)峰值顯存占用降低 20%~30%，并且不會(huì)帶來(lái)訓(xùn)練吞吐量和測(cè)試精度的損失。

Approx-BP

在傳統(tǒng)反向傳播訓(xùn)練中，激活函數(shù)梯度的反向回傳是嚴(yán)格對(duì)應(yīng)其導(dǎo)函數(shù)的，對(duì)于 Transformer 模型中常用的 GELU 和 SiLU 函數(shù)，這意味著需要將輸入特征張量完整地存入激活顯存中。而本文的作者提出了一套反向傳播近似理論，即 Approx-BP 理論。在該理論的指導(dǎo)下，作者使用分段線性函數(shù)逼近激活函數(shù)，并用分段線性函數(shù)的導(dǎo)數(shù)（階梯函數(shù)）替代 GELU/SiLU 梯度的反向回傳。這個(gè)方法導(dǎo)出了兩個(gè)非對(duì)稱的內(nèi)存高效激活函數(shù)：ReGELU2 和 ReSiLU2。這類激活函數(shù)由于使用 4 段階梯函數(shù)進(jìn)行反向回傳，從而使得激活存儲(chǔ)只需要使用 2bit 數(shù)據(jù)類型。

MS-BP

BP 網(wǎng)絡(luò)每一層通常都會(huì)將輸入張量存入激活顯存以用作反向傳播計(jì)算。作者指出如果可以將某一層的反向傳播改寫成依賴輸出的形式，那么這一層和后一層就可以共享同一個(gè)激活張量，從而降低激活存儲(chǔ)的冗余。

而文章指出 Transformer 模型中常用的 LayerNorm 和 RMSNorm，在將仿射參數(shù)合并到后一層的線性層之后，可以很好地符合 MS-BP 策略的要求。經(jīng)過(guò)重新設(shè)計(jì)的 MS-LayerNorm 和 MS-RMSNorm 不再產(chǎn)生獨(dú)立的激活顯存。

實(shí)驗(yàn)結(jié)果

作者對(duì)計(jì)算機(jī)視覺(jué)和自然語(yǔ)言處理領(lǐng)域的若干個(gè)代表模型進(jìn)行了微調(diào)實(shí)驗(yàn)。其中，在 ViT，LLaMA 和 RoBERTa 的微調(diào)實(shí)驗(yàn)中，文章提出的方法分別將峰值顯存占用降低了 27%，29% 和 21%，并且沒(méi)有帶來(lái)訓(xùn)練效果和訓(xùn)練速度的損失。注意到，作為對(duì)比的 Mesa（一個(gè) 8-bit Activation Compressed Training 方法）使訓(xùn)練速度降低了約 20%，而文章提出的 LowMemoryBP 方法則完全保持了訓(xùn)練速度。

結(jié)論及意義

文章提出的兩種 BP 改進(jìn)策略，Approx-BP 和 MS-BP，均在保持訓(xùn)練效果和訓(xùn)練速度的同時(shí)，實(shí)現(xiàn)了激活顯存的顯著節(jié)省。這意味著從 BP 原理上進(jìn)行優(yōu)化是非常有前景的顯存節(jié)省方案。此外，文章提出的 Approx-BP 理論突破了傳統(tǒng)神經(jīng)網(wǎng)絡(luò)的優(yōu)化框架，為使用非配對(duì)導(dǎo)數(shù)提供了理論可行性。其導(dǎo)出的 ReGELU2 和 ReSiLU2 展現(xiàn)了這一做法的重要實(shí)踐價(jià)值。

歡迎大家閱讀論文或者代碼去了解算法的詳細(xì)細(xì)節(jié)，LowMemoryBP 項(xiàng)目的 github 倉(cāng)庫(kù)上已經(jīng)開(kāi)源相關(guān)的模塊。