如何设置内存分给显存 显存不足

2080Ti其实可以当V100用 , 功能有点强大 。
自深度学习浪潮兴起以来 , 该模式一直在向更大、更深的方向发展 。
2012年 , 拥有五个卷积层的AlexNet首次在视觉任务上展现出了超强的能力 。之后 , 基础模式开始“深化”:2014年 , VGG网达到19层;2015年的ResNet和2017年的DenseNet已经将深度提升到了几百层 。
模型大小的增加大大提高了性能 。因此 , ResNet、DenseNet等 。被认为是所有主要视觉任务的基本骨干 。但同时 , 机型的增加也意味着对视频内存的需求变得更高 。
为什么GPU内存如此重要?
九年前 , 辛顿等人率先用两个3GB视频内存的GTX 580 GPU高效训练AlexNet 。此后 , 对视频内存的需求和模型的大小一直同步增长 。想要在游戏中获得好的成绩 , 想要在实验中超越最先进的效果 , 想要在工程中契合庞大的商业数据等等 。 , 这一切都离不开记忆的加持 。
模型增加一层 , 视频内存会增加一个点 。
在深度学习模型中 , 占用内存的总是那些特别大的张量 , 比如各层的权重矩阵、计算出来的张量(激活值)、反向传播所需的张量等 。在视觉任务中 , 中间计算的张量占绝大多数 。随着模型越来越深、越来越大 , 每一层的激活值张量都需要保存在主存中 。
以ResNet50为例 。在模型的训练中 , 前向传播中50层的计算结果需要存储在视频存储器中 , 以便后向传播可以使用这些张量来计算梯度 。如果您使用ResNet108 , 您需要的视频内存将是ResNet50的两倍多 。视频内存的增加 , 当然带来了模型效果的提升 。另一方面 , 如果没有足够的视频内存 , 很多工作将无法完成 。
记忆力不够 , 写论文、玩游戏反复受到约束 。
在实验室运行模型和写论文的过程中 , 内存不足的情况并不少见 。一般实验室的显卡都是大家共享的 , 可能大家手里剩下的不多了 。甚至 , 随着顶级模型变得越来越大 , 没有人有足够的计算能力和内存来重现最终的实验 , 更不用说超越它的SOTA结果了 。
在这种情况下 , 学生只有两种选择:向导师申请新的GPU资源 , 或者缩小模型 , 做一个Mini实验 。前者不一定成功 , 后者可能有各种不完美 。如果能在内存有限的情况下运行顶级模型 , 做实验和写论文会更容易 。
此外 , 无论是在学校还是在公司 , 计算能力和内存不足都是很常见的 。顶级竞争对手的模式结构可能几乎相同 。区别在于谁的模型更大 , 更能处理复杂的样本 。更直观的是 , 排行榜领跑者的模型可能只差十几层 , 但正是因为内存限制少于十几层 , 一些模型才错过了冠军 。
内存:算法工程的瓶颈
再举一个常见的例子 , 企业中的算法工程师有足够的计算能力 , 内存就没那么重要了 。但是 , 如果只采用并行策略来共享视频内存 , 可能仍然有足够的视频内存 , 但是每个GPU的计算负载不足 。

如何设置内存分给显存 显存不足

文章插图
如何设置内存分给显存 显存不足

文章插图
4 V100 , 视频内存满了 , 但是GPU利用率很低 。
即使有V100强大的计算能力 , 在训练大型机型时也很容易占用16GB的视频内存 。但由于批次大小不足 , 上图中每个V100 GPU的利用率只有20%到30% 。只有不断提高每次迭代的数据吞吐量 , 才能提高GPU的利用率 。
MegEngine:视频内存需要优化 。
其实对于深度学习的从业者来说 , 日常应用中远不止以上三种情况 。做深度学习 , 无论是研究还是工程 , 都会时不时遇到记忆问题 。然而 , 这个问题的优化非常复杂 , 需要大量的工程实现来缓解 。显然 , 这样的优化应该由深度学习框架来完成 。但是在实际应用中不难发现 , TensorFlow和PyTorch似乎并没有提供完美的官方解决方案 。
但如果我们看一下新势力 , 情况可能就不一样了 。在最近发布的开源深度学习框架MegEngine 1.4版本中 , 该框架首次引入了动态图形内存优化技术 , 大大减少了内存占用问题 。
具体来说 , MegEngine通过复制和优化ICLR 2021 Spotlight论文《动态张量再物质化》(以下简称DTR)实现了“为更多视频内存而计算” 。有了这项技术的加持 , 模型的内存占用大大减少 , 同样的硬件可以训练更大的模型 , 搭载更大的BatchSize 。这样 , 学生的小显卡也可以开始训练大型号 , 工程师的服务器也可以站得更全应用 。
如何设置内存分给显存 显存不足

文章插图
如何设置内存分给显存 显存不足

文章插图
对于原本需要16GB视频内存的机型 , 优化后使用的视频内存峰值降到了4GB 。
MegEngine是一项内存优化技术 , 它使1060这样的入门级显卡能够训练只能在2080Ti之前加载的型号 。11GB视频内存的2080Ti可以挑战原来只能32GB V100训练的机型 。要知道 , V100的价格是2080Ti的9倍多 。
两行代码 , 记忆“加倍”
如果需要优化自己的视频内存 , 可能99%的算法工程师都会放弃 。最好的方法是告诉深度学习框架本次训练会分配多少视频内存 , 剩下的留给框架优化 。MegEngine的动态图形内存优化就是基于这个逻辑 。
通过两行代码 , 框架可以完全自动优化视频内存 , 并在MegEngine中隐藏所有优化逻辑和复杂的工程实现 。
如何设置内存分给显存 显存不足

文章插图
如何设置内存分给显存 显存不足

文章插图
如上图所示 , 动态计算图中导入了DTR内存优化模块 , 内存释放阈值配置为5GB 。在训练过程中 , 由于视频内存已经“翻倍” , 批处理大小可以在翻两番后加载到GPU中 。
视频内存扩展带来的收益
在许多情况下 , 提高视频内存利用率的最显著效果是训练更大的模型 。在一定程度上 , 参数数量越多 , 效果越好 。批量越大 , 梯度更新方向越准确 , 模型性能越好 。MegEngine开发团队做了很多实验 , 在提高视频内存利用率的同时 , 保证了训练的高质量 。
最简单的验证方法是不断增加批处理大小 , 看看显卡能走多远 。以下两个表格分别显示了DTR技术在PyTorch和MegEngine上加载或不加载的效果 。
如何设置内存分给显存 显存不足

文章插图
如何设置内存分给显存 显存不足

文章插图
如果不使用动态图形内存优化技术 , PyTorch上的模型在一次训练迭代中最多只能处理64个样本 , MegEngine最多可以处理100个样本 。有了DTR , PyTorch模型可以在一次迭代中处理140个样本 , MegEngine可以尝试处理300个样本 。
如果换算成模型大小 , 再加上MegEngine的动态图形内存优化技术 , 在相同的GPU和批处理大小下 , 它可以高效地训练出增加近5倍的模型 。
MegEngine动态图形内存优化技术
深度学习模型的内存占用一般分为三个部分:权重矩阵、前向传播中间张量和后向传播梯度矩阵(Adam optimizer) 。
权重矩阵和梯度矩阵占用的内存很难优化 , 每个模型基本都有固定值 。前向传播的中间计算结果不一样:随着Batch Size和模型层数的增加 , 视频内存必然会增加 。如果模型很大 , 中间的计算结果会占据最重要的内存 。
如何设置内存分给显存 显存不足

文章插图
如何设置内存分给显存 显存不足

文章插图
如上图所示 , 在正向传播中(第一行从左到右) , 蓝色圆圈表示模型的中间计算结果开始占用视频内存 。直到前向传播完成 , 第一行完全变成蓝圈 , 之前计算占用的视频内存无法释放 。
当反向传播开始时(第二行从右向左) , 通过梯度的计算和应用 , 保留在主存储器中的正向传播张量可以被释放 。
显然 , 要想减少对视频内存的占用 , 就必须使用前向传播保存的中间计算结果 , 这也是MegEngine动态图形内存优化的主要方向 。
视频内存的交换计算
动态计算图最直接的方法就是用计算或内存来交换视频内存 。因此 , MegEngine必须首先决定使用哪种技术 。
MegEngine团队通过实验发现 , 计算时间远小于交换时间 。例如 , 它从视频内存中节省了612.5MB 空 , 用带宽交换视频内存比用计算交换视频内存慢几十倍到几百倍 。
如何设置内存分给显存 显存不足

文章插图
如何设置内存分给显存 显存不足

文章插图
因此 , 很明显 , 动态计算图中也应该使用梯度检查点技术 , 并且应该用计算来交换视频存储器 。
下面是梯度检查点的技术原理 。前向传播的第三点是检查点 , 它将始终存储在视频内存中 。第四点可以在计算完成后释放视频内存 。如果反向传播中需要第四点的值 , 可以从第三点重新计算第四点的值 。
如何设置内存分给显存 显存不足

文章插图
如何设置内存分给显存 显存不足

文章插图
虽然一般原理不难理解 , 但怎么做却相当复杂 。MegEngine团队从论文《动态张量再物质化》中学习 , 对其进行了优化 , 并在MegEngine中实现 。
DTR , 最先进的内存优化技术
DTR是一个完全动态的启发式策略 。核心思想是当视频内存超过一定阈值时 , 动态释放一些合适的张量 , 直到视频内存降到阈值以下 。一般来说 , 释放张量有三个标准:重新计算张量的代价越小越好;占用的视频内存越大越好;在视频记忆中停留的时间越长越好 。
如何设置内存分给显存 显存不足

文章插图
如何设置内存分给显存 显存不足

文章插图
除了从检查点恢复前向传播结果张量的主要成本外 , DTR的额外成本在于找到应该释放的最优张量 , 即计算上图中张量T的f(t)值 。为了减少这部分的计算量 , MegEngine还采用了两种运行时优化:
不要考虑小张量 , 它们不加入候选集 。
每次需要释放一个张量时 , 都会随机采样并遍历几个张量 , 以节省计算成本 。
最难的是项目的实现 。
虽然DTR在原则上看起来并不复杂 , 但真正的问题在于提高可用性 , 也就是将所有细节隐藏在框架的底层 , 只为开发人员提供最简单的界面 。
在这里 , 我们使用最简单的计算示例再次遵循框架计算 , 看看MegEngine如何通过使用动态图的计算历史来恢复和释放张量 。
如何设置内存分给显存 显存不足

文章插图
如何设置内存分给显存 显存不足

文章插图
现在 , 假设输入中有两个张量A和B , 你想计算a*b和a+b , 但最大内存只能容纳三个张量 。当黄框计算c=a+b时 , 视频内存仍然可以保留张量c , 但在下一步中 , 当绿框计算d=a*b时 , 只有c可以被释放 , d才能被保存 。
不幸的是 , 在下一步中 , 灰色框需要得到黄色框的计算结果 。不过为了节省视频内存 , C已经发布了 。因此 , MegEngine现在需要做的就是重新运行灰盒的计算图 , 计算c=a+b , 并加载到视频内存中 。显然 , 这将不可避免地需要释放D的视频内存 。
这样 , 鉴于内存的限制 , MegEngine会自动选择合适的张量释放 , 并在必要时重新计算 。如果需要重新计算一个张量的结果 , 比如上图中的D , 则需要具体的历史计算信息(这里是像a+b这样的计算路径) , 同时需要知道两个输入张量A和b 。
所有这样的历史计算信息都是由MegEngine自动获取保存的 , MegEngine的工程师在底层已经用C++进行了处理 , 用户完全不需要考虑 。
struct ComputePath {std::shared_ptr op;SmallVector inputs;SmallVector outputs;double compute_time = 0;} *producer;SmallVector users;size_t ref_cnt = 0;
以上是MegEngine底部用来跟踪计算路径信息的结构 。其中op表示生成张量的算子;输入和输出分别表示该算子所需的输入和输出张量;Compute_time表示操作员的实际运行时间 。
其实在使用MegEngine的过程中 , 所有张量都是用Python接口创建的 , 但是框架会相应地跟踪每个张量的具体信息 。每当需要访问张量时 , 无论张量是否在主存中 , 没有它都可以立即恢复 。所有这些复杂的工程操作和操作逻辑都隐藏在MegEngine C++的底层 。
如何设置内存分给显存 显存不足

文章插图
如何设置内存分给显存 显存不足

文章插图
Python代码将被翻译成C++底层实现 , C++代码将通过指针(右边绿色部分)管理显卡内存中的实张量 。
幸运的是 , 这个复杂的操作不需要由算法工程师来完成 , 而是全部留给MegEngine 。
MegEngine可以做的远不止这些 , 但大部分都像是动态图形内存优化的技术 , 帮助用户无形中解决实际问题 。2020年3月 , 开源的MegEngine以肉眼可见的速度快速成长 , 从静态计算图到动态计算图 , 再到不断提升的训练能力、移动终端推理性能的优化、动态内存的优化...这或许就是开源的魅力所在 。只有不断的优化和创新 , 才能吸引和满足“挑剔”的开发者 。MegEngine的下一个功能会是什么?让我们拭目以待 。
【如何设置内存分给显存 显存不足】

    推荐阅读