节省显存新思路，在 PyTorch 里使用 2 bit 激活压缩训练神经网络

机器之心专栏
作者：郑怜悯、陈键飞
随着超大规模深度学习模型逐渐成为AI的趋势，如何在有限的GPU内存下训练这些模型成为了一个难题。
本文将介绍来自加州伯克利大学的ActNN ，一个基于PyTorch的激活压缩训练框架。在同样的内存限制下， ActNN通过使用2bit激活压缩，可以将batchsize扩大6-14倍，将模型尺寸或者输入图片扩大6-10倍。 ActNN相关论文已被ICML2021接收为LongTalk ，代码开源于github 。
论文https://arxiv.org/abs/2104.14129
【节省显存新思路，在 PyTorch 里使用 2 bit 激活压缩训练神经网络】代码https://github.com/ucbrise/actnn
AI训练撞上「内存墙」
从AlexNet ， ResNet到GPT-3 ，深度学习性能的突破都离不开模型规模的疯狂增长。大模型有更好的性能已经成为业界的共识。过去几年，不仅训练一个最先进模型需要的算力在指数增长，训练一个最先进模型需要的内存也在指数增长。如下图所示，大型Transformer模型的参数量以每两年翻240倍的速度指数增长。但是，单个GPU的内存却只以每两年翻2倍的速度在缓慢增长。另外，在训练模型时，不光要存储模型参数，还要存储中间结果激活值和优化器状态，所需要的内存更多。如何在有限的GPU内存下训练这些大规模模型成为了挑战。

文章图片
source：GholamiA,YaoZ,KimS,MahoneyMW,KeutzerK.AIandMemoryWall.RiseLabMediumBlogPost,UniversityofCaliforniaBerkeley
节省训练内存的方法
目前，节省训练内存的方法主要有三类：1.重计算（Gradientcheckpointing/Rematerialization)2.使用CPU内存进行交换(swapping)和3.使用分布式训练将Tensor分散存储在多个GPU上。这三类方法互相不冲突，可以结合使用。大部分机器学习框架对这些方法都提供了一些支持，也有不少相关的论文。但是，想要高效、自动化地实现这些策略并不容易。与已有方法不同，我们提出了ActNN ，一个新的基于压缩的内存节省框架。在提供理论证明的同时，我们基于PyTorch提供了一个高效易用的实现。 Table.1比较了ActNN和已有的一些内存节省系统。 ActNN支持PyTorch的动态图执行模式，并且不需要预先进行复杂的策略搜索。 ActNN作为一个独立的Python库，使用时import即可，不需要修改或重新编译PyTorch 。与已有的工作相比， ActNN灵活且易于使用。同时， ActNN在理论上也可以和已有的技术相互叠加。

文章图片
ActNN：2bit激活压缩训练
在训练一个多层神经网络时，在前向传播中，每一层的中间结果都要被存下来用于计算反向传播的梯度。这些中间结果，又被叫做「激活值」（activation) ，实际上占据了大部分的内存消耗，尤其是在batchsize较大或者输入图片较大的时候。 ActNN的原理是就是压缩这些激活值来节省内存。如下图所示，左图表示的是普通的前向传播和反向传播，前向传播时会存下所有层的FP32激活值用于反向传播，内存使用在计算loss的时候达到峰值。右图表示的是ActNN的训练方法：在前向传播时，通过一个压缩操作Q将激活值压缩后再存储；反向传播时，通过解压缩操作Q^-1将激活值解压再计算梯度。

文章图片
如果只是为了节省内存，这里可以使用各种压缩算法，但是大部分现有的压缩算法并不能高效地运行在GPU上，会引入较大的开销。 ActNN选择了使用2-bit量化作为这里的压缩算法。量化操作的代价较小，而且有一些好的数学性质允许我们使用有损压缩达到较大的压缩比。