阿里NIPS 2017论文解读：如何降低TensorFlow训练的显存消耗？

Posted 2021-04-24 阿里技术

阿里妹导读：被誉为国际顶级的机器学习和神经网络学术会议NIPS 2017于12月4日-9日在美国加州长滩市举行。在本届会议上，阿里巴巴除有两篇论文入选Workshop并进行Oral和Poster形式报告外，三大技术事业部连续3天在阿里展区举行多场技术研讨会，向5000余名参会人员介绍阿里在机器学习、人工智能领域的技术研究、产品与落地应用。 

NIPS 2017阿里巴巴展台-阿里巴巴iDST院长金榕在做演讲

 

这篇介绍深度模型训练GPU显存优化的论文《TrainingDeeper Models by GPU Memory Optimization on TensorFlow》在NIPS 2017 ML Systems Workshop 中由作者做口头报告。这篇论文聚焦特征图，提出两种方法减少深度神经网络训练过程中的显存消耗，并且把这些方法的实现无缝整合到TensorFlow中，克服了TensorFlow训练大模型时无法有效优化显存的缺点。

 

近期深度学习在不同应用中发挥的作用越来越重要。训练深度学习模型的必要逻辑包括适合GPU的并行线性代数计算。但是，由于物理限制，GPU的设备内存（即显存）通常比主机内存小。最新的高端NVIDIA GPU P100具备12–16 GB的显存，而一个CPU服务器有128GB的主机内存。然而，深度学习模型的趋势是「更深更宽」的架构。例如，ResNet 包含多达1001个神经元层，神经网络机器翻译（NMT）模型包含8个使用注意力机制的层 ，且NMT模型中的大部分的单个层是按顺序水平循环展开的，难以避免地带来大量显存消耗。

简言之，有限的GPU显存与不断增长的模型复杂度之间的差距使显存优化成为必    然。下面将介绍深度学习训练流程中GPU显存使用的主要组成。

 

特征图（feature map）。对于深度学习模型，特征图是一个层在前向传输中生成的中间输出结果，且在后向传输的梯度计算中作为输入。图1是ResNet-50在ImageNet数据集上进行一次小批量训练迭代的GPU显存占用曲线。随着特征图的不断累积，曲线到达最高点。特征图的大小通常由批尺寸（batchsize）和模型架构决定（如CNN架构的卷积步幅大小、输出通道数量；RNN架构的门数量、时间步长和隐层大小）。不再需要作为输入的特征图占用的显存将会被释放，导致图1中显存占用曲线的下降。对于复杂的模型训练，用户必须通过调整批尺寸，甚至重新设计模型架构来避免「内存不足」的问题。尽管在分布式训练的情况下 ，训练任务可以分配到多个设备上来缓解内存不足的问题，但是这也导致了额外的通信开销。设备的带宽限制也可能显著拖慢训练过程。

图1：ResNet-50的显存占用在一个训练步中的变化曲线。横轴代表分配／释放次数，纵轴代表当前显存占用的总比特数。

 

权重。与特征图相比，权重占用内存相对较少 。在这篇论文中，权重作为GPU内存中的持久内存，只有整个训练任务完成后才可以被释放。

 

临时显存（Temporary memory）。一些算法（如基于Fast-Fourier-Transform（FFT）的卷积算法）需要大量的额外显存。这些显存占用是暂时的，在计算结束后立即得到释放。临时显存的大小可以通过在GPU软件库（如cuDNN）中列举每个算法来自动调整，因此可以被忽略。

很明显，特征图是GPU显存使用的主要组成部分。论文作者聚焦特征图，提出了两种方法来解决GPU显存限制问题，即通用的「swap-out/in」方法以及适用于Seq2Seq模型的内存高效注意力层。所有这些优化都基于TensorFlow 。TensorFlow具备内置内存分配器，实现了「best-fit with coalescing」的算法。该分配器旨在通过coalescing支持碎片整理（de-fragmentation）。但是，它的内置内存管理策略未考虑大模型训练时的显存优化。

 

该论文的贡献如下。聚焦于特征图，提出两种方法减少深度神经网络训练过程中的GPU显存消耗。基于数据流图的「swap-out/in」方法使用主机内存作为更大的内存池，从而放宽GPU显存上限的限制；而内存高效的注意力层可用来优化显存消耗量大的Seq2Seq模型。这些方法的实现被无缝整合到TensorFlow中，且可透明地应用于所有模型，无需对现有模型架构的描述作任何改变。

 

论文： Training Deeper Models by GPUMemory Optimization on TensorFlow

 

https://github.com/LearningSys/nips17/blob/9ee207c054cf109bc4a068b1064b644d75d0381f/assets/papers/paper_18.pdf

 

摘要：随着大数据时代的到来、GPGPU的获取成本降低以及神经网络建模技术的进步，在GPU上训练深度学习模型变得越来越流行。然而，由于深度学习模型的内在复杂性和现代GPU的显存资源限制，训练深度模型仍然是一个困难的任务，尤其是当模型大小对于单个GPU而言太大的时候。在这篇论文中，我们提出了一种基于通用数据流图的GPU显存优化策略，即「swap-out/in」，将主机内存当做一个更大的内存池来克服GPU的内存限制。同时，为了优化内存消耗大的Seq2Seq模型，我们还提出了专用的优化策略。我们将这些策略无缝整合到TensorFlow中，且优化不会造成准确率的损失。我们在大量的实验中观察到了显著的显存使用降低。给定一个固定的模型和系统配置，最大训练批尺寸可以增加2到30倍。

 

图2：引用计数（referencecount）

图3：swapout/in 优化的原子操作（Atomic operation）。删除从节点e到节点b的引用边，并添加了红色和蓝色的节点和边。

图4：注意力操作（Attentionoperation）优化。d指梯度。图左未经优化，图右经过了显存优化。  

表1：对swap out/in 的评估，GPU的显存上限是12GB。

 表2：对显存高效序列模型的评估。

你可能还喜欢

 

关注「阿里技术」

把握前沿技术脉搏