如何在单块GPU上训练超大型深度学习模型

问题:GPU 内存限制


GPU 在深度神经网络训练之中的强大表现无需我赘言。通过现在流行的深度学习框架将计算分配给 GPU 来执行,要比自己从头开始便捷很多。然而,有一件事你会避之唯恐不及,即 GPU 的动态随机存取内存(DRAM(Dynamic Random Access Memory))限制。


在给定模型和批量大小的情况下,事实上你可以计算出训练所需的 GPU 内存而无需实际运行它。例如,使用 128 的批量训练 AlexNet 需要 1.1GB 的全局内存,而这仅是 5 个卷积层加上 2 个全连接层。如果我们观察一个更大的模型,比如 VGG-16,使用 128 的批量将需要大约 14GB 全局内存。目前最先进的 NVIDIA Titan X 内存为 12GB。VGG-16 只有 16 个卷积层和 3 个全连接层,这比大约有 100 个层的 resnet 模型小很多。


640.jpeg

图 1: 当使用基线、全网分配策略时 GPU 内存的使用情况(左轴)。(Minsoo Rhu et al. 2016)


现在,如果你想要训练一个大于 VGG-16 的模型,而你又没有一个具有内存的 GPU,你也许有几个解决内存限制问题的选择。


  • 减小你的批量大小,但这可能会妨碍你的训练速度和精确度。

  • 在多 GPU 环境下做模型并行,这是另一个复杂的事情。

  • 缩小你的模型,如果你不情愿做出上述两个选择,或者已经尝试但效果不好。


或者你也可以等待下一代更强大的 GPU 出现。然而将来的网络变得更大更深是大势所趋,我们并不希望物理内存的限制成为算法开发的一个障碍。


观察:什么占用内存


我们可以根据功能性把 GPU 内存中的数据分为 4 个部分:


  • 模型参数

  • 特征图

  • 梯度图

  • 工作区


前 3 个功能容易理解。模型参数的意义对于所有人来说都很熟悉了。特征图是正向过程中生成的中间结果。梯度图是反向过程中生成的中间结果。而工作区是 cuDNN 库函数所使用的临时变量/矩阵的一个缓冲区。对于一些 cuDNN 函数,用户需要将缓冲区作为函数参数传给 GPU 内核函数。一旦函数返回,该缓冲区将被释放。


640-2.jpeg



我们可以看到,一般而言,我们有的层越多,分配给特征图的内存比重就越多(由上图中的三角形表示)。我们也可以看到对于像 VGG-16 这样的更大模型来说,这一比例基本上要超过 50%。


想法:使用 CPU 内存作为临时容器


有一个关于特征图的事实:它们在正向过程中生成之后立刻被用于下一层,并在稍后的反向过程中仅再用一次。每个 GPU 内核函数仅使用与当前层(通常只有 1 个张量)相关的特征映射。这将导致绝大多数内存在几乎所有的时间上出现空置的情况(它们保有数据但不使用)。


这一想法是:如果 GPU 内存中的大部分数据出现空置,为什么不把它们保存在更便宜的 CPU 内存上呢?下图更清晰地展现了这一想法。



640-11.jpeg



左侧部分所示的间隙表明特征图如何在内存之中被空置。右侧部分表明这一节省 GPU 内存使用的想法:使用 CPU 内存作为那些特征图的临时容器。


权衡:时间 vs 空间


根据论文,vDNN(虚拟化 DNN,也就是原文的系统设计的名称)把 AlexNet 的平均 GPU 内存使用成功降低了 91%,GoogLeNet 的降低了 95%。但你很可能已经看到,这样做的代价是训练会变慢。例如,vDNN 可以在 12GB 的 GPU 上使用 256 的批量训练 VGG-16,但是假设我们在一块拥有足够内存的 GPU 上训练同样的模型而不使用 vDNN 来优化内存使用,我们可以避免 18% 的性能损失。


当使用 cuDNN 核时,工作区大小也会出现权衡的情况。一般而言,你有的工作区越多,你可以使用的算法就越快。如果有兴趣请查阅 cuDNN 库的参考。在后面的整个讨论中我们都将会看到有关时间空间的这一权衡。


优化策略:在前向过程中卸载,在后向过程中预取


你应该已经知道 vDNN 是如何在正向过程中优化内存分配的。基本的策略是在生成特征图后将其从 GPU 上卸下传给 CPU,当它将在反向过程中被重新使用时再从 CPU 预取回 GPU 内存。这个存储空间可被释放以作他用。这样做的一个风险是如果网络拓扑是非线性的,特征图的一个张量可能被应用于数个层,从而导致它们不能被立刻卸载。当然,这个问题可以通过简单的完善优化策略来解决。


在后向过程中,vDNN 采用一种更具侵略性的策略,这是由于相较于特征图,梯度图不存在「稍后重用」的问题。因此一旦生成了相关的权值更新,它们就可以被释放(相较于那些映射,这些权值更新是很小的)。


优化策略:内存管理器 CUDA 流


vDNN 实现的关键是使用一个 cuda 流来管理内存分配/释放、卸载和预取。以下是一些细节:


传统的 cuda 内存分配/释放(cudaMalloc & cudaFree)是同步性 API。由于这两个操作随着训练过程需要不断地被执行,同步性 API 并不够高效。


如同分配/释放操作,卸载 API 也需要是异步性的。当 vDNN 选择卸载特征图的一个张量时,vDNN 的内存管理器流(memory manager stream of vDNN)将在主机上分配一个页锁定的内存区域,并且通过 PCIe 发出一个非阻塞传输(a non-blocking transfer)将其传输至 CPU。这些特征图张量在正向过程中为只读的,因此这个传输过程可以在计算中被安全地覆盖。当依赖于其传输和计算都完成后,相关的张量可以被释放。只有当前层的卸载完成时,才能进行下一层的处理。


预取操作是在反向处理中从 CPU 返回到 GPU 以得到之前被卸载的特征图。和上面的操作类似,预取操作也需要是异步性的。由于在预取和相同层的计算之间存在数据依赖,vDNN 将同时异步开始当前层的计算以及前一层的预取。


成本:为了节省内存的性能损失在哪里?


最显著的潜在性能损失来自于由卸载/预取引入的隐式依赖(implicit dependency)。我们考虑这样一种情况:当数据传输比正向计算需要花费更长的时间时,卸载/预取操作将会带来性能损失。下图表清晰地表明这种情况(图 9:卸载和预取的性能影响 (Minsoo Rhu et al. 2016))


640-12.jpeg


相似情形也可能在反向过程中发生。


问题的形式化:在限制内存预算的情况下,如何获得最佳性能?


如上所述,在时间和空间之间有一个权衡,并且在前的章节中我们已经看到这个权衡是产生作用的。想象一下你正在 12GB 的 GPU 上使用 128 的批量(这需要 14GB 内存如果没有使用卸载/预取)训练 VGG-16。假设我们可以使用上面提到的想法,把内存使用压缩到只有 2GB,那么仅使用了大约 2GB 的内存也许很浪费,因为你本可以利用空闲的 10GB 空间来减少性能损失(这里忽略能耗方面的成本)。因此,我们可以这种方式重新形式化这个问题:在限制内存预算的情况下,如何获得最佳性能?



配置时间-空间权衡:决定一个层是否应该被卸载/预取(offloaded/prefetched),以及应该选择哪种卷积算法。


从这个问题的形式化中我们可以看出,我们往往并不需要把内存使用压缩到极致,有些层完全可以不被卸载/预取。为了取得最佳配置,我们需要为每个层决定两件事:一个是我们是否需要卸载/预取,一个是我们在前向/反向传播过程中该选择哪个算法(更快收敛的算法需要更大的存储空间)。


通常,靠近输入的层有较长重用度(reuse)距离。因此,首先卸载/预取这些层比较好。于是,我们并不需要为每一层决定是否使用(这个选择总量随层级数成指数增长)。我们只需要选择一层作为基准,与之相比更加靠近输入层的都进行卸载/预取,其余层将其张量保留在 GPU 上。


为每层决定算法也并不可行(同样选择总量随层级数成指数增长)。我们可以通过例如强制多层使用相同的算法(gemm 或 fft 等算法)来缩小我们的可选备用空间。这将原本指数级大小的选择数降为线性大小。


经过上面的处理,现在我们的配置选择空间已经足够小,这样我们可以尽情搜索以确定最佳配置方案。下面的图表说明了我们的配置空间:


640-13.jpeg


左上角表示最节省内存的配置(卸载/预取每一层,并使用最慢的算法),右下角表示性能最优的配置(当然,真实的配置空间应该是网格,可行模型与不可行模型之间的边界应改成阶梯状,不过这一图表足以表明整个搜索空间的大致分布特征)。


下一步即找到具有最佳性能的可行配置。我们可以沿可行性边界搜索配置。如果你对该搜索过程的实现细节感兴趣,请参阅原论文(搜索过程在第 3.C 节中进行了描述)。


论文连接:https://arxiv.org/pdf/1602.08124.pdf640.png


推荐者介绍:


Hongyu 毕业于上海交通大学 09 级 ACM 班,多伦多大学在读博士,现从事关于深度神经网络训练过程的一些 profiling 的工作。


加入机器之心 ML 系统与架构小组:


近些年人工智能的突破不仅仅是机器学习算法的努力,更有包含了其所依赖的系统和架构领域的进步。如果你想更加全面的了解机器学习这个领域,对其依赖的上下游领域的了解会很有帮助。系统和架构就是这样一个下游领域。这个方向的工作极大方便了其上游算法的开发,或许你也是现在诸多 ML 系统产品的用户之一。但面对一个这样重要的跨领域的方向,你可能会感到这样一些困境:


  1. 找不到合适的学习资料

  2. 有学习动力,但无法坚持

  3. 学习效果无法评估

  4. 遇到问题缺乏讨论和解答的途径


不论你是想要获得相关跨领域的更全面大局观,还是你只是想对手中的 ML 系统工具更加了解,你都是机器之心想要帮助和挖掘的对象。基于这个方向现在越来越成为众多研究人员关注的焦点,机器之心发起了一个互助式学习小组——「人工智能研学社· ML 系统与架构小组」。本小组将通过优质资料分享、教材研习、论文阅读、群组讨论、专家答疑、讲座与分享等形式加强参与者对强化学习和深度学习的理解和认知。


  • 面向人群:有一定的机器学习算法基础,并且对分布式计算、并行计算等也有所了解,同时想掌握此方向最新成果的人

  • 学习形式:学习资料推荐、统一进度学习(教材或论文)、群组讨论、专家答疑、讲座等。

  • 加入方式:


1)添加机器之心小助手微信,并注明:加入系统与架构学习组


640-14.jpeg


2)完成小助手发送的入群测试(题目会根据每期内容变化),并提交答案,以及其他相关资料(教育背景 、从事行业和职务 、人工智能学习经历等)

3)小助手将邀请成功通过测试的朋友进入「人工智能研学社· ML 系统与架构学习组」


入群测试 QUIZ


1)教育背景 2)从事行业和职务 3)人工智能经历 4)人工智能系统或架构经历


  1. 数据并行和模型并行的区别在于哪里?

  2. 在训练过程中,数据并行和模型并行很难达到线性的 scale-up,你认为瓶颈会在哪?有没有办法突破这样的瓶颈?这些突破的方法又会造成什么样的副作用?

工程深度学习硬件深度研学社GPU工程
暂无评论
暂无评论~