ICML 2024 | 梯度檢查點太慢？不降速、省顯存，LowMemoryBP大幅提升反向傳播顯存效率
本文論文一作是南開大學統計與數據科學學院研二碩士生楊雨辰，指導老師為南開大學統計與數據科學學院的徐君副教授。徐君老師團隊的研究重點是計算機視覺、生成式 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 倉庫上已經開源相關的模塊。
