[论文解读] Echo: Compiler-based GPU Memory Footprint Reduction for LSTM RNN Training

Bojian Zheng, Abhishek Tiwari+2 more|arXiv (Cornell University)|May 22, 2018

Advanced Neural Network Applications参考文献 72被引用 6

一句话总结

Echo 是一种基于编译器的优化技术，通过选择性地重新计算特征图而非存储它们，减少了 LSTM RNN 训练中的 GPU 内存占用。该技术利用层特定的运行时估计和全局内存收益分析，实现了平均 1.89 倍、最大 3.13 倍的内存占用降低，且运行时开销极小，从而在相同硬件上实现更快的训练速度、更大的批量大小和更深的模型。

ABSTRACT

The Long-Short-Term-Memory Recurrent Neural Networks (LSTM RNNs) are a popular class of machine learning models for analyzing sequential data. Their training on modern GPUs, however, is limited by the GPU memory capacity. Our profiling results of the LSTM RNN-based Neural Machine Translation (NMT) model reveal that feature maps of the attention and RNN layers form the memory bottleneck and runtime is unevenly distributed across different layers when training on GPUs. Based on these two observations, we propose to recompute the feature maps rather than stashing them persistently in the GPU memory. While the idea of feature map recomputation has been considered before, existing solutions fail to deliver satisfactory footprint reduction, as they do not address two key challenges. For each feature map recomputation to be effective and efficient, its effect on (1) the total memory footprint, and (2) the total execution time has to be carefully estimated. To this end, we propose *Echo*, a new compiler-based optimization scheme that addresses the first challenge with a practical mechanism that estimates the memory benefits of recomputation over the entire computation graph, and the second challenge by non-conservatively estimating the recomputation overhead leveraging layer specifics. *Echo* reduces the GPU memory footprint automatically and transparently without any changes required to the training source code, and is effective for models beyond LSTM RNNs. We evaluate *Echo* on numerous state-of-the-art machine learning workloads on real systems with modern GPUs and observe footprint reduction ratios of 1.89X on average and 3.13X maximum. Such reduction can be converted into faster training with a larger batch size, savings in GPU energy consumption (e.g., training with one GPU as fast as with four), and/or an increase in the maximum number of layers under the same GPU memory budget.

研究动机与目标

• 解决 LSTM RNN 训练中 GPU 内存消耗过高的问题，该问题限制了批量大小和计算利用率。
• 克服先前重新计算技术在 LSTM 环境下无法准确估计内存收益和运行时开销的局限性。
• 在不修改训练源代码的前提下，实现自动且透明的内存占用降低。
• 在相同的 GPU 内存预算下，支持除 LSTM RNN 外的多种模型，包括 Transformer 和 ResNet。
• 通过内存占用降低，实现更快的训练速度、更大的批量大小或更高的模型深度。

提出的方法

• Echo 采用基于编译器的方法，分析完整的计算图，并估算通过重新计算特征图而非存储它们所带来的净内存收益。
• 它引入了一种实用机制，通过同时考虑释放的特征图和重新计算所需的新分配数据，评估全局内存节省效果。
• Echo 通过利用层特定的特性（如非全连接层的计算成本较低）非保守地估计重新计算的开销。
• 该方法基于内存节省与运行时成本之间的权衡，识别出最优的重新计算层，避免对计算密集型层进行过度简化的排除。
• 它可无缝集成到现有的深度学习框架中，无需修改模型源代码或手动调整内核。
• Echo 在真实 GPU 上对多个最先进的模型（包括 NMT、DeepSpeech2、Transformer 和 ResNet）进行了评估。

Figure 1 : Left: A single-layer LSTM RNN that scans through an input sequence. Right: Zoom-in view of one LSTM cell. Both diagrams are simplified for clarity.

实验结果

研究问题

• RQ1 选择性重新计算是否能在不引入显著运行时开销的前提下，有效降低 LSTM RNN 训练中的 GPU 内存占用？
• RQ2 如何在考虑存储释放和新分配数据的前提下，准确估算整个计算图中重新计算的净内存收益？
• RQ3 基于编译器的系统能否自动识别复杂模型（如 Transformer 和 ResNet）中最佳的重新计算目标？
• RQ4 内存占用降低在多大程度上可转化为实际优势，如更大的批量大小或更深的模型？
• RQ5 在多种 DNN 工作负载下，Echo 与先前技术相比，在内存减少和运行时效率方面表现如何？

主要发现

• Echo 在包括 NMT、DeepSpeech2、Transformer 和 ResNet 在内的多种模型上，平均将 GPU 内存占用降低了 1.89 倍，最大降低达 3.13 倍。
• 该方法使单张 GPU 的训练速度达到基线配置下四张 GPU 的水平，显著降低了能耗。
• 在批量大小固定的情况下，Echo 使单张 RTX 2080 Ti（11 GB 内存）上的 Transformer 模型最大层数增加 1.83 倍，ResNet 模型增加 4.0 倍。
• 该方法仅带来极小的运行时开销，证明了其在真实训练工作负载中的实用性。
• Echo 不仅适用于 LSTM RNN，对 Transformer 和 ResNet 等 CNN 模型也表现出强大的适用性。
• 该技术可实现更快的训练收敛速度，并在相同内存预算下支持更深的模型，与深度学习架构设计的发展趋势保持一致。

Figure 2 : NMT (LSTM RNN Encoder (Middle)-Decoder (Right) with Attention (Left))