机器之心原创
作者:思
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 的计算负载又不足的情况。
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。如此一来,学生的小显卡也能开始训练大模型,而工程师们的服务器也经得起更充分的应用。
原本需要 16GB 显存的模型,优化后使用的显存峰值就降到了 4GB。
MegEngine 这种显存优化技术,让 1060 这样的入门级显卡也能训练原本 2080Ti 才能加载得上的模型;而 11GB 显存的 2080Ti,更能挑战原本 32GB V100 才能训练的模型。要知道,V100 的价格可是 2080Ti 的 9 倍还多。
两行代码,显存「翻倍」
如要需要自己去优化显存, 可能 99% 的算法工程师都会放弃。最好的办法是告诉深度学习框架,这次训练就分配多少显存,剩下的就交给框架自己去优化。MegEngine 的动态图显存优化就是基于这一逻辑。
通过两行代码,框架可以全自动地完成显存优化,将所有优化逻辑与复杂的工程实现都隐藏在 MegEngine 内部。
如上图所示,在动态计算图中导入 DTR 显存优化模块,并配置显存释放阈值为 5GB。训练时,因为显存已经「翻倍」了,Batch Size 翻四倍也能装到 GPU 中。
显存扩增带来的收益
很多时候,提高显存的利用率,最显著的作用就是能训练更大的模型。从一定程度上来说,参数量越大就意味着效果越好;而批大小越大,梯度更新方向就越准确,模型性能也就越优异。MegEngine 开发团队做了很多实验,以确保提高显存利用率的同时训练是优质的。
最简单的验证方法就是不断增加批大小,看看显卡到底能坚持到什么程度。下面两张表分别展示了在 PyTorch 及 MegEngine 上加载或不加载动态图显存优化(DTR)技术的效果。
如果不使用动态图显存优化技术,PyTorch 上的模型一次训练迭代最多只能处理 64 个样本,MegEngine 能处理 100 个样本。只要加上 DTR,PyTorch 模型一次迭代就能处理 140 个样本,MegEngine 能尝试处理 300 个样本。
如果换算成模型大小,加上动态图显存优化技术的 MegEngine,在相同的 GPU 及批大小情况下,能高效训练增大近乎 5 倍的模型。
MegEngine 动态图显存优化技术
深度学习模型的显存占用一般分为权重矩阵、前向传播的中间张量、反向传播的梯度矩阵(Adam 优化器)三部分。
权重矩阵和梯度矩阵占的内存很难优化,各个模型基本上都有一个定值。前向传播的中间计算结果则不然:随着 Batch Size 的增加以及模型层和数量的增加,显存必然跟着增加。如果模型比较大,中间计算结果将占据最主要的显存。
如上图所示,在前向传播中(第一行从左到右),蓝色圆圈表示模型的中间计算结果开始占用显存。一直到前向传播完成,第一行完全变为蓝色圆圈,前面计算所占用的显存都不能释放。
等到反向传播开始(第二行从右到左),随着梯度的计算与完成应用,前向传播保留在显存中的张量才可以释放。
很明显,如果要降低显存占用,就要拿前向传播保存的中间计算结果开刀,这也正是 MegEngine 动态图显存优化的主要方向。
用计算换显存
对于动态计算图,最直接的方法就是用计算或内存换显存。因此,MegEngine 首先要决定到底使用哪种技术。
MegEngine 团队通过实验发现,用计算耗时远比交换耗时少。例如从显存中节省 612.5MB 空间,用带宽换显存要比用计算换显存慢了几十上百倍。
因此很明确,动态计算图中也应该使用梯度检查点技术,用计算换显存。
如下为梯度检查点技术原理示意,前向传播中第三个点为检查点,它会一直保存在显存中。第四个点在完成计算后即可释放显存,在反向传播中如果需要第四个点的值,可以从第三个点重新计算出第四个点的值。
虽然大致原理不难理解,但具体怎么做还是比较复杂的,MegEngine 团队借鉴了论文《Dynamic Tensor Rematerialization》,将其优化并实现到 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 表示产生该张量的算子;inputs 和 outputs 分别表示这个算子需要的输入与输出张量;compute_time 表示该算子实际的运行时间。
实际上,在使用 MegEngine 的过程中,全都是用 Python 接口创建张量,只不过框架会对应追踪每个张量的具体信息。每当需要访问张量,不用考虑张量是否在显存中时,没有也能立刻恢复出来。所有这些复杂的工程化的操作与运算逻辑都隐藏在了 MegEngine C++ 底层。
Python 代码会翻译成 C++ 底层实现,C++ 代码会通过指针管理显卡内存中真正的张量(右图绿色部分)。
幸好这样的复杂操作不需要算法工程师完成,都交给 MegEngine 好了。
MegEngine 能做的事情远不止于此,只不过大多是像动态图显存优化这种技术一样,润物细无声地把用户的实际问题解决于无形。2020 年 3 月开源的 MegEngine 在以肉眼可见的速度快速成长,从静态计算图到动态计算图,再到持续提升的训练能力、移动端推理性能优化、动态显存优化…… 这也许就是开源的魅力。只有不断优化和创新,才能吸引和满足「挑剔」的开发者。MegEngine 下一个推出的功能会是什么?让我们拭目以待。
欢迎体验MegEngine DTR 功能:MegStudio-提供免费算力的在线深度学习开发平台