显存不够,框架来凑:两行代码显存翻倍,2080Ti也能当V100来用-阿里云开发者社区

开发者社区> 超努力的写代码> 正文

显存不够,框架来凑:两行代码显存翻倍,2080Ti也能当V100来用

简介: 2080Ti 竟然可以当 V100 来用,这个功能有点儿厉害。
+关注继续查看

自深度学习大潮兴起,模型就朝着越来越大、越来越 「深」 的方向发展。

2012 年,拥有 5 个卷积层的 AlexNet 第一次在视觉任务上展现出强大的能力。在此之后,基础模型就开始「深」化起来:2014 年的 VGG-Net 达到了 19 层;2015 年的 ResNet、2017 年的 DenseNet 更是将深度提升到了上百层。


模型大小的提升极大地提高了性能。因此,各大视觉任务都将 ResNet、DenseNet 等当做基本的 BackBone。但与此同时,模型的增大也意味着对显存的需求随之变高。


为什么 GPU 显存如此重要?


九年前,Hinton 等人率先用两张 3GB 显存的 GTX 580 GPU 高效训练 AlexNet。在此之后,显存需求与模型大小就一直同步增长。打比赛想要取到好成绩、做实验想要超越 State of the art 效果、做工程想要拟合庞大的业务数据等等,这些都离不开显存的加持。


模型加一层,显存涨一分


在深度学习模型中,占用显存的总是那些特别大的张量,比如各层的权重矩阵、计算出来的张量(激活值)、反向传播需要的张量等。在视觉任务中,占据绝大多数的是中间计算出来的张量。随着模型变得更深更大,每一层的激活值张量都需要保留在显存中。


以 ResNet50 为例,在模型的训练中,前向传播中 50 层的计算结果都需要保存在显存中,以便让反向传播利用这些张量计算梯度。如果使用 ResNet108,需要的显存就会比 ResNet50 多出一倍多。显存的增加,带来的当然是模型效果的提升。另一方面,如果显存不够,许多工作也必将无法实现。


显存不够,写论文、打比赛屡遭掣肘


在实验室跑模型、写论文的过程中,显存不够用也是常有的事。一般实验室的显卡都是大家共用的,可能分配到每个人的手上已经所剩无几。甚至于,随着顶尖模型越来越大,所有人都没有足够的算力、显存去复现终极实验,更不用说超越其 SOTA 结果。


遇到这种情况,学生无非只有两种选择:向导师申请新的 GPU 资源,或者缩减模型做一个 Mini 版的实验。前者并不总是能够成功,后者则可能会有种种不完美。如果能用有限的显存跑顶尖的大模型,做实验、写论文都会变得更加简单。


此外,无论是在学校还是在公司打比赛,算力不够、显存不足都是常有的事。顶尖竞争者的模型结构可能相差无几,区别就在于谁的模型更大、更有能力去处理复杂的样本。更直观地说,排行榜领先者的模型也许就只差十几层,但也正是因为显存受限少了那十几层,有些模型才与冠军失之交臂。


显存:约束算法工程师的瓶颈


再举一个常见的例子,企业中的算法工程师拥有足够的算力,显存没那么重要。然而,只使用并行策略分担显存,还是可能会出现显存足够、但每张 GPU 的计算负载又不足的情况。

6.jpg

4 张 V100,显存占满,而 GPU 利用率很低。

即使是 V100 这样强大的算力,训练大模型时也很容易占满 16GB 显存。然而由于批量不够大,上图每张 V100 GPU 的利用率只有 20% 到 30%。只有继续增大每次迭代的数据吞吐量,才能增加 GPU 的利用率。


MegEngine:显存需要优化


其实对于深度学习从业者来说,日常应用中出现的情况远不止上面三种。做深度学习,不论是研究还是工程,时不时就会遇到显存问题。但这个问题优化起来又很复杂,需要利用大量的工程实现来缓解。显然,这样的优化应该由深度学习框架来完成。不过,在实际应用中不难发现,TensorFlow、PyTorch 似乎都没有提供完善的官方解决方案。


但如果把目光投向新生势力,情况可能就不一样了。在旷视开源深度学习框架 MegEngine 最近发布的 1.4 版本中,该框架首次引入了动态图显存优化技术,大大降低了显存占用问题。


具体来说,通过复现并优化 ICLR 2021 Spotlight 论文《Dynamic Tensor Rematerialization》(以下简称 DTR),MegEngine 实现了「用计算换取更多显存」。有了这项技术的加持,模型的显存占用大大降低,同样的硬件可以训练更大的模型、承载更大的 BatchSize。如此一来,学生的小显卡也能开始训练大模型,而工程师们的服务器也经得起更充分的应用。


7.jpg


原本需要 16GB 显存的模型,优化后使用的显存峰值就降到了 4GB。
MegEngine 这种显存优化技术,让 1060 这样的入门级显卡也能训练原本 2080Ti 才能加载得上的模型;而 11GB 显存的 2080Ti,更能挑战原本 32GB V100 才能训练的模型。要知道,V100 的价格可是 2080Ti 的 9 倍还多。


两行代码,显存「翻倍」


如要需要自己去优化显存, 可能 99% 的算法工程师都会放弃。最好的办法是告诉深度学习框架,这次训练就分配多少显存,剩下的就交给框架自己去优化。MegEngine 的动态图显存优化就是基于这一逻辑。


通过两行代码,框架可以全自动地完成显存优化,将所有优化逻辑与复杂的工程实现都隐藏在 MegEngine 内部。


8.jpg


如上图所示,在动态计算图中导入 DTR 显存优化模块,并配置显存释放阈值为 5GB。训练时,因为显存已经「翻倍」了,Batch Size 翻四倍也能装到 GPU 中。


显存扩增带来的收益


很多时候,提高显存的利用率,最显著的作用就是能训练更大的模型。从一定程度上来说,参数量越大就意味着效果越好;而批大小越大,梯度更新方向就越准确,模型性能也就越优异。MegEngine 开发团队做了很多实验,以确保提高显存利用率的同时训练是优质的。


最简单的验证方法就是不断增加批大小,看看显卡到底能坚持到什么程度。下面两张表分别展示了在 PyTorch 及 MegEngine 上加载或不加载动态图显存优化(DTR)技术的效果。


9.jpg


如果不使用动态图显存优化技术,PyTorch 上的模型一次训练迭代最多只能处理 64 个样本,MegEngine 能处理 100 个样本。只要加上 DTR,PyTorch 模型一次迭代就能处理 140 个样本,MegEngine 能尝试处理 300 个样本。


如果换算成模型大小,加上动态图显存优化技术的 MegEngine,在相同的 GPU 及批大小情况下,能高效训练增大近乎 5 倍的模型。


MegEngine 动态图显存优化技术


深度学习模型的显存占用一般分为权重矩阵、前向传播的中间张量、反向传播的梯度矩阵(Adam 优化器)三部分。


权重矩阵和梯度矩阵占的内存很难优化,各个模型基本上都有一个定值。前向传播的中间计算结果则不然:随着 Batch Size 的增加以及模型层和数量的增加,显存必然跟着增加。如果模型比较大,中间计算结果将占据最主要的显存。


10.gif


如上图所示,在前向传播中(第一行从左到右),蓝色圆圈表示模型的中间计算结果开始占用显存。一直到前向传播完成,第一行完全变为蓝色圆圈,前面计算所占用的显存都不能释放。


等到反向传播开始(第二行从右到左),随着梯度的计算与完成应用,前向传播保留在显存中的张量才可以释放。


很明显,如果要降低显存占用,就要拿前向传播保存的中间计算结果开刀,这也正是 MegEngine 动态图显存优化的主要方向。


用计算换显存


对于动态计算图,最直接的方法就是用计算或内存换显存。因此,MegEngine 首先要决定到底使用哪种技术。


MegEngine 团队通过实验发现,用计算耗时远比交换耗时少。例如从显存中节省 612.5MB 空间,用带宽换显存要比用计算换显存慢了几十上百倍。

11.jpg


因此很明确,动态计算图中也应该使用梯度检查点技术,用计算换显存。

如下为梯度检查点技术原理示意,前向传播中第三个点为检查点,它会一直保存在显存中。第四个点在完成计算后即可释放显存,在反向传播中如果需要第四个点的值,可以从第三个点重新计算出第四个点的值。


12.gif


虽然大致原理不难理解,但具体怎么做还是比较复杂的,MegEngine 团队借鉴了论文《Dynamic Tensor Rematerialization》,将其优化并实现到 MegEngine 中。


DTR,最前沿的显存优化技术


DTR 是一种完全动态的启发式策略,核心思想是当显存超过某个阈值时,动态地释放一些合适的张量,直到显存低于阈值。一般而言,释放张量的标准有三个:重新计算出该张量的开销越小越好;占用的显存越大越好;在显存中停留的时间越长越好。


13.jpg

除去从检查点恢复前向传播结果张量带来的主要开销,DTR 的额外开销在于寻找应该被释放的最优张量,即计算上图张量 t 的 f(t)值。为了降低这一部分的计算量,MegEngine 还采用了两种运行时优化:

  • 不考虑小的张量,它们不加入候选集
  • 每次在需要释放张量的时候,随机采样并遍历少部分张量,以节省计算开销


最难的是工程实现


虽然 DTR 看上去原理也不复杂,但真正的难题在于提高易用性,即将所有细节都隐藏到框架的底层,只为开发者提供最简单的接口。


在此就用一个最简单的计算例子,跟着框架演算一遍,看看 MegEngine 是如何利用动态图的计算历史恢复与释放张量的。



14.jpg


现在假设输入有 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++ 处理完毕,用户完全不需要考虑。

image.png

以上为 MegEngine 底层用于追踪计算路径信息的结构体。其中 op 表示产生该张量的算子;inputs 和 outputs 分别表示这个算子需要的输入与输出张量;compute_time 表示该算子实际的运行时间。


实际上,在使用 MegEngine 的过程中,全都是用 Python 接口创建张量,只不过框架会对应追踪每个张量的具体信息。每当需要访问张量,不用考虑张量是否在显存中时,没有也能立刻恢复出来。所有这些复杂的工程化的操作与运算逻辑都隐藏在了 MegEngine C++ 底层。

15.jpg

Python 代码会翻译成 C++ 底层实现,C++ 代码会通过指针管理显卡内存中真正的张量(右图绿色部分)。

幸好这样的复杂操作不需要算法工程师完成,都交给 MegEngine 好了。


MegEngine 能做的事情远不止于此,只不过大多是像动态图显存优化这种技术一样,润物细无声地把用户的实际问题解决于无形。2020 年 3 月开源的 MegEngine 在以肉眼可见的速度快速成长,从静态计算图到动态计算图,再到持续提升的训练能力、移动端推理性能优化、动态显存优化…… 这也许就是开源的魅力。只有不断优化和创新,才能吸引和满足「挑剔」的开发者。MegEngine 下一个推出的功能会是什么?让我们拭目以待。

版权声明:本文内容由阿里云实名注册用户自发贡献,版权归原作者所有,阿里云开发者社区不拥有其著作权,亦不承担相应法律责任。具体规则请查看《阿里云开发者社区用户服务协议》和《阿里云开发者社区知识产权保护指引》。如果您发现本社区中有涉嫌抄袭的内容,填写侵权投诉表单进行举报,一经查实,本社区将立刻删除涉嫌侵权内容。

相关文章
阿里云服务器怎么设置密码?怎么停机?怎么重启服务器?
如果在创建实例时没有设置密码,或者密码丢失,您可以在控制台上重新设置实例的登录密码。本文仅描述如何在 ECS 管理控制台上修改实例登录密码。
10081 0
Fundebug前端JavaScript插件更新至1.7.1,拆分录屏代码,还原部分Script error.
摘要: BUG监控插件压缩至18K。 1.7.0拆分了录屏代码,BUG监控插件压缩至18K,另外我们还原了部分Script error,帮助用户更方便地Debug。请大家及时更新哈~ 拆分录屏代码 从1.7.0版本开始,我们拆分了录屏代码。
1349 0
不改一行代码,轻松拥有企业级微服务治理|MSE微服务治理专业版重磅发布
随着业务的发展,微服务拆分越来越复杂,微服务的治理也成了一个比较令人头疼的问题。有没有更加简单且高效的方法来解决微服务治理的难题?
444 0
程序员自己的svn代码库 同步公司和家里的代码
最近试用了google code, 是一个类似source forge网站的东西, 你可以自己建项目,并且将项目checkin到google code的SVN里面 .
884 0
怎么设置阿里云服务器安全组?阿里云安全组规则详细解说
阿里云服务器安全组设置规则分享,阿里云服务器安全组如何放行端口设置教程
8467 0
汇编(三)段地址、偏移地址、内存单元、寄存器、CS、IP、代码段、debug、实验题
段地址、偏移地址、内存单元、寄存器、CS、IP、代码段、debug、实验题、汇编、汇编编程、win10怎么使用debug
4259 0
阿里云服务器如何登录?阿里云服务器的三种登录方法
购买阿里云ECS云服务器后如何登录?场景不同,阿里云优惠总结大概有三种登录方式: 登录到ECS云服务器控制台 在ECS云服务器控制台用户可以更改密码、更换系.
13885 0
1946
文章
0
问答
文章排行榜
最热
最新
相关电子书
更多
《2021云上架构与运维峰会演讲合集》
立即下载
《零基础CSS入门教程》
立即下载
《零基础HTML入门教程》
立即下载