通过动态规划优化反向传播的检查点选择,显著降低GPU内存占用。
GPU Memory Usage Optimization for Backward Propagation in Deep Network Training
- 基于数学建模识别前向与反向传播所需关键数据
- 提出O(n)时间复杂度算法,实现峰值内存最低的检查点配置
- 适用于大规模模型训练,尤其适合显存受限场景
在现代深度学习中,设计更大的神经网络以应对更复杂的任务并提升精度已成为趋势。然而,卷积层中间数据的内存分配在训练过程中常引发严重内存压力。现有解决方案包括依赖硬件的方法和通用的重计算(rematerialization)策略——即在前向传播阶段仅保存部分中间结果作为检查点,反向传播时按需重新计算。该方法以增加计算时间为代价,避免存储全部中间数据,从而降低内存占用。本文聚焦于高效寻找最优检查点子集,以最小化训练过程中的峰值内存使用。首先通过数学方程描述网络训练机制,识别前后向传播所需的全部必要数据。基于理论分析,修正了检查点选择问题的目标函数,并结合追踪数据提出一种O(n)时间复杂度的算法,显著优于原有O(n³)的动态规划方法。大量实验验证了新方法在降低内存消耗方面的有效性。
原文摘要 · Abstract (English)
In modern Deep Learning, it has been a trend to design larger Deep Neural Networks (DNNs) for the execution of more complex tasks and better accuracy. On the other hand, Convolutional Neural Networks (CNNs) have become the standard method for most of computer vision tasks. However, the memory allocation for the intermediate data in convolution layers can cause severe memory pressure during model training. Many solutions have been proposed to resolve the problem. Besides hardware-dependent solutions, a general methodology rematerialization can reduce GPU memory usage by trading computation for memory efficiently. The idea is to select a set of intermediate results during the forward phase as checkpoints, and only save them in memory to reduce memory usage. The backward phase recomputes the intermediate data from the closest checkpoints in memory as needed. This recomputation increases execution time but saves memory by not storing all intermediate results in memory during the forward phase. In this paper, we will focus on efficiently finding the optimal checkpoint subset to achieve the least peak memory usage during the model training. We first describe the theoretical background of the training of a neural network using mathematical equations. We use these equations to identify all essential data required during both forward and backward phases to compute the gradient of weights of the model. We first identify the checkpoint selection problem and propose a dynamic programming algorithm with time complexity O(n3) to solve the problem of finding the optimal checkpoint subset. With extensive experiments, we formulate a more accurate description of the problem using our theoretical analysis and revise the objective function based on the tracing, and propose an O(n)-time algorithm for finding the optimal checkpoint subset.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。