通过近似和内存共享反向传播降低精调内存开销
机器学习
2024-06-25 v1 人工智能
摘要
对预训练的大模型进行下游任务的精调是一个重要问题，但由于大规模参数的存在，面临巨大的内存开销。本文致力于从激活函数和层归一化两个角度降低精调过程中的内存开销。为此，我们提出了近似反向传播（Approx-BP）理论，提供了正向传播与反向传播解耦的理论可行性。我们将Approx-BP理论应用于反向传播训练，推导出GELU和SiLU激活函数的内存高效替代方案，这些替代方案在反向传播中使用ReLU的导数函数，但保持正向传播不变。此外，我们引入了内存共享反向传播策略，使两个相邻图层的激活内存可以共享，从而消除激活内存使用的冗余。我们的方法既不会引入额外计算，也不会降低训练效率。我们针对预训练的视觉和语言模型进行了大量实验，结果表明，我们的方案可减少最高可达约30%的峰值内存使用。我们的代码已发布于https://github.com/yyyyychen/LowMemoryBP。
引用
@article{arxiv.2406.16282,
title = {Reducing Fine-Tuning Memory Overhead by Approximate and Memory-Sharing Backpropagation},
author = {Yuchen Yang and Yingdong Shi and Cheems Wang and Xiantong Zhen and Yuxuan Shi and Jun Xu},
journal= {arXiv preprint arXiv:2406.16282},
year = {2024}
}
备注
25 pages, ICML 2024 Accepted