[论文解读] SuperNeurons: Dynamic GPU Memory Management for Training Deep Neural Networks

Linnan Wang, Jinmian Ye+6 more|arXiv (Cornell University)|Jan 13, 2018

Advanced Neural Network Applications参考文献 26被引用 23

一句话总结

SuperNeurons 引入了一种动态 GPU 内存调度运行时，使训练深度神经网络能够突破 GPU DRAM 的物理限制。通过结合活性分析（Liveness Analysis）、统一张量池（Unified Tensor Pool）和成本感知重计算（Cost-Aware Recomputation），它将峰值内存使用量降低至最大层内存使用量，从而实现在 12GB 的 K40c GPU 上训练 ResNet2500（10⁴ 层）——相比 TensorFlow 最多可训练深达 3.24 倍的模型，且在各类框架中均表现出领先性能。

ABSTRACT

Going deeper and wider in neural architectures improves the accuracy, while the limited GPU DRAM places an undesired restriction on the network design domain. Deep Learning (DL) practitioners either need change to less desired network architectures, or nontrivially dissect a network across multiGPUs. These distract DL practitioners from concentrating on their original machine learning tasks. We present SuperNeurons: a dynamic GPU memory scheduling runtime to enable the network training far beyond the GPU DRAM capacity. SuperNeurons features 3 memory optimizations, extit{Liveness Analysis}, extit{Unified Tensor Pool}, and extit{Cost-Aware Recomputation}, all together they effectively reduce the network-wide peak memory usage down to the maximal memory usage among layers. We also address the performance issues in those memory saving techniques. Given the limited GPU DRAM, SuperNeurons not only provisions the necessary memory for the training, but also dynamically allocates the memory for convolution workspaces to achieve the high performance. Evaluations against Caffe, Torch, MXNet and TensorFlow have demonstrated that SuperNeurons trains at least 3.2432 deeper network than current ones with the leading performance. Particularly, SuperNeurons can train ResNet2500 that has $10^4$ basic network layers on a 12GB K40c.

研究动机与目标

• 解决由于 GPU DRAM 容量有限而导致的训练更深层、更宽的深度神经网络时日益严重的内存瓶颈问题。
• 克服现有框架（如 Caffe、Torch、MXNet 和 TensorFlow）中静态内存优化技术的局限性，这些技术无法适应非线性架构中动态数据与依赖关系的变化。
• 在无需手动模型分割或多 GPU 分解的情况下，实现远超现代 GPU 物理内存容量的极深网络的端到端训练。
• 通过集成性能感知调度与通信重叠机制，在实现激进内存优化的同时保持高训练性能。
• 提供一种透明、用户无感知的运行时环境，动态管理张量的分配、释放与卸载，同时保证在层粒度上的最小峰值内存使用量。

提出的方法

• 活性分析（Liveness Analysis）识别并回收前向与反向传播过程中不再需要的张量，将峰值内存使用量从 ∑(l_i^f + l_i^b) 降低至 ∑(l_i^f) + l_N^b。
• 统一张量池（UTP）将计算量大、内存占用高的层（检查点）卸载至外部内存，进一步将峰值内存使用量降低至 ∑(l_i^f ∉ checkpoints) + l_N^b。
• 成本感知重计算（Cost-Aware Recomputation）选择性地丢弃廉价或非检查点层的前向结果，并在反向传播期间重建，将峰值内存最小化至 max(l_i)，即任意单层的最大内存使用量。
• 通过预分配的 GPU 内存池与基于堆的分配器，重用内存段以摊销频繁张量操作的分配/释放开销，降低运行时成本。
• 基于 GPU DRAM 的 LRU 张量缓存通过重用频繁访问的张量，实现通信与计算的重叠，减少 UTP 操作期间 CPU 与 GPU 之间的数据移动。
• 该运行时通过以张量为调度单位，与深度学习框架无缝集成，契合 cuDNN 的逐层计算模型，对用户完全透明。

实验结果

研究问题

• RQ1 动态内存调度如何在突破 GPU DRAM 容量限制的前提下，减少深度神经网络训练中的峰值内存使用量？
• RQ2 活性分析与张量重用在具有可变内存访问模式的非线性、深层架构中，能在多大程度上缓解内存压力？
• RQ3 在数据并行训练中，通过将张量从 GPU DRAM 卸载至外部内存并实现通信与计算的重叠，是否能维持高训练吞吐量？
• RQ4 成本感知重计算在不降低训练性能的前提下，通过牺牲计算来换取内存节省的效率如何？
• RQ5 统一的、动态的内存管理系统是否能在多种深度学习框架中，同时优于静态内存优化技术的内存效率与训练速度？

主要发现

• SuperNeurons 实现了在 12GB 的 K40c GPU 上训练 ResNet2500（约 10⁴ 个基本层）的训练，该模型已超出当前框架的内存容量。
• 其训练深度至少是 TensorFlow 的 3.2432 倍，MXNet 的 4.0000 倍，Torch 的 12.6316 倍，以及 Caffe 的 12.9730 倍。
• 峰值内存使用量被降低至 max(l_i)，即任意单层所需的最大内存，实现了在层粒度上的理论最小值。
• SuperNeurons 在各类框架中均保持领先的训练性能，端到端评估显示其在线性网络（AlexNet、VGG16）和非线性网络（ResNet50→152、Inception V4）中均保持一致的速度优势。
• 由于高效的通信重叠与缓存机制，随着批量大小增加，性能下降幅度极小；但当 GPU 内存仅能容纳单层时，性能下降有所增加。
• 该系统支持的模型规模比 Caffe 支持的最大模型大 19.8 倍，展现出显著超越当前最先进框架的可扩展性。