节约显存新思绪，正在 PyTorch 里利用 2 bit 激活紧缩练习神经网络

呆板之心专栏

作者：郑恻隐、陈键飞

随着超大范围深度进修模子渐渐成为 AI 的趋向，怎样在有限的 GPU 内存下练习这些模子成为了一个困难。

本文将先容来自加州伯克利大学的 ActNN，一个基于 PyTorch 的激活压缩练习框架。在同样的内存限定下，ActNN 通过利用 2 bit 激活压缩，可以将 batch size 扩大 6-14 倍，将模子尺寸大概输入图片扩大 6-10 倍。ActNN 相干论文已被 ICML 2021 吸收为 Long Talk，代码开源于 github。

论文 http://arxiv.org/abs/2104.14129

代码 http://github.com/ucbrise/actnn

AI 练习撞上「内存墙」

从 AlexNet，ResNet 到 GPT-3，深度进修性能的突破都离不开模子范围的猖獗增进。大模子有更好的性能已经成为业界的共鸣。已往几年，不但练习一个开始进模子必要的算力在指数增进，练习一个开始进模子必要的内存也在指数增进。如下图所示，大型 Transformer 模子的参数目以每两年翻 240 倍的速率指数增进。但是，单个 GPU 的内存却只以每两年翻 2 倍的速率在迟钝增进。别的，在练习模子时，不但要存储模子参数，还要存储中心效果激活值和优化器状态，所必要的内存更多。怎样在有限的 GPU 内存下练习这些大范围模子成为了挑衅。

source：Gholami A, Yao Z, Kim S, Mahoney MW, Keutzer K. AI and Memory Wall. RiseLab Medium Blog Post, University of California Berkeley

节约练习内存的要领

现在，节约练习内存的要领重要有三类：1. 重盘算（Gradient checkpointing/Rematerialization) 2. 利用 CPU 内存举行互换 (swapping) 和 3. 利用漫衍式练习将 Tensor 疏散存储在多个 GPU 上。这三类要领相互不辩论，可以联合利用。大部门呆板进修框架对这些要领都提供了一些支持，也有不少相干的论文。但是，想要高效、主动化地实现这些计谋并不简单。与已有要领差别，我们提出了 ActNN，一个新的基于压缩的内存节约框架。在提供理论证明的同时，我们基于 PyTorch 提供了一个高效易用的实现。Table.1 比力了 ActNN 和已有的一些内存节约体系。ActNN 支持 PyTorch 的动态图实行模式，而且不必要预先辈行庞大的计谋搜刮。ActNN 作为一个独立的 Python 库，利用时 import 即可，不必要修改或重新编译 PyTorch。与已有的事情相比，ActNN 机动且易于利用。同时，ActNN 在理论上也可以和已有的技能相互叠加。

ActNN：2 bit 激活压缩练习

在练习一个多层神经网络时，在前向流传中，每一层的中心效果都要被存下来用于盘算反向流传的梯度。这些中心效果，又被叫做「激活值」（activation)，现实上据有了大部门的内存斲丧，尤其是在 batch size 较大大概输入图片较大的时间。ActNN 的道理是便是压缩这些激活值来节约内存。如下图所示，左图表现的是一般的前向流传和反向流传，前向流传时会存下全部层的 fp32 激活值用于反向流传，内存利用在盘算 loss 的时间到达峰值。右图表现的是 ActNN 的练习要领：在前向流传时，通过一个压缩操纵 Q 将激活值压缩后再存储；反向流传时，通过解压缩操纵 Q^-1 将激活值解压再盘算梯度。

假如只是为了节约内存，这里可以利用种种压缩算法，但是大部门现有的压缩算法并不克不及高效地运行在 GPU 上，会引入较大的开销。ActNN 选择了利用 2-bit 量化作为这里的压缩算法。量化操纵的价钱较小，并且有一些好的数学性子许可我们利用有损压缩到达较大的压缩比。

把 fp32 浮点数目化为 2-bit 整数是一个有损压缩，会引入一些偏差。论文从理论上阐发了量化引入的偏差是怎样影响练习的收敛性的。

第一，存在一个随机化的量化计谋，使得利用有损量化压缩后，估量出的有损梯度是原梯度的一个无偏估量。

在这一条件下，我们套用已有的随机梯度降落收敛性定理，得出最终收敛时的偏差会被梯度的方差所限定。

第二，我们推导出了利用量化压缩之后，随机梯度降落盘算出的梯度的方差。

等号右边的第一项是随机梯度降落在 minibatch 采样时孕育发生的方差，等号右边的第二项是有损压缩分外引入的方差。这条公式表现地描画了有损压缩带来的影响。细致到，当有损量化压缩带来的方差远小于本来随机梯度降落自带的方差时，ActNN 引入的有损压缩就不会影响练习的收敛性。更多关于公式的推导和可视化拜见文末的论文链接。论文对差别的算子（conv2d，batch norm，linear等）都提供了细致的阐发。

由上述公式开导，我们提出了一些新的量化本领用于低落有损压缩引入的分外方差。我们引入了新的量化本领 （ Per-group Quantization，Fine-Grained Mixed-Precision，Runtime Adaptation） 来使用梯度在差别样本，差别纬度，差别层之间的异构特性。最终的压缩算法会安排更多的 bit 给更紧张的激活值。均匀每个浮点数安排到 2 bit。

在详细实现压缩算法时，另有许多可以调治的参数。这里孕育发生了一个内存节约和练习速率的弃取。普通来说，利用更庞大的压缩算法可以节约更多的内存，但是也会引入更多分外的开销，使练习速率变慢。为了给用户较大的机动性，ActNN 提供了 5 个优化品级 L1-L5 供用户选择。低的优化品级节约的内存较少，但是运行速率快。高的优化品级节约的内存多，但是运行也更慢。在最高优化品级 L5 下，ActNN 会联合一个简洁的内存互换计谋，将压缩后的激活值移到 CPU 内存上，进一步节约内存。

实现

要在 PyTorch 实现 ActNN 算法非常简洁。对付一个 PyTorch nn Module，我们只必要在其 forward 函数里参加量化压缩，在其 backward 函数里参加解压缩操纵。全部的盘算照旧在 fp32 下举行，与本来一样，伪代码如下图所示。

ActNN 为大部门常用的 PyTorch nn.Module 实现了利用量化压缩的版本。用户只需将模子里的全部 PyTorch nn.Module 更换成 ActNN 对应的 Module （如把 nn.Conv2d 更换成 actnn.Conv2d)，即可节约内存，不必要变动其他代码。ActNN 同时也提供了一个 wrapper 实现一行代码主动更换。

试验效果

由于 ActNN 举行的是有损压缩，以是最紧张的一点是先验证 ActNN 是否会影响模子的精度。下图是利用 ActNN 在 ImageNet 上练习 ResNet-50 的效果。FP 代表一般的 fp32 练习， BLPA 是来自 NeurIPS 2019 的一个相干事情。可以看到，在 ActNN 的 2-bit 压缩模式下，模子险些没有丧失精度。在更极限的 1.25 bit 的情形下，ActNN 也能收敛，只不外会丧失一些精度。而之前的事情 BLPA 在小于 4 bit 的情形就下无法收敛。

我们还在图像支解，物体检测，以及自监视进修等多个使命上举行了试验。ActNN 都能在 2-bit 压缩模式下到达和一般 fp32 险些一样的效果。在部门使命上，由于 ActNN 可以利用更大的 batch size，乃至可以取得更好的测试效果。细致的试验效果和练习记载拜见文末的论文与 github 链接。

之后，我们比拟了 ActNN 与一般 fp32 练习的现实内存利用情形。如下表所示，ActNN 可以将激活值占用的内存压缩 12 倍，将练习利用的总内存压缩 4 - 7 倍。这一现实内存压缩成效切合理论推导。为什么激活值压缩倍率是 12 而不是 32 bit / 2 bit = 16？重要是由于 ActNN 不克不及利用 inplace 的 ReLU，以及必要存储少量分外的 min 和 scale 用于解压缩。

最终，我们测试了 ActNN 的练习速率。由于 ActNN 在练习历程中举行了压缩，这些压缩在节约内存的同时也会引入分外的盘算开销。普通来说，免得内存越多，进入的分外开销就越多，练习也就越慢。我们在 NVIDIA T4 (16 GB 内存) 上比拟了 ActNN 和已有内存节约体系的练习速率。如下图所示，DTR （ICLR 2020），BLPA （NeurIPS 2019）和 swap 分别是基于重盘算，压缩和内存互换的三种要领，红叉代表 Out-of-memory。y 轴是练习吞吐量 (images per second)，越高越好。绿色的曲线是综合 ActNN 在差别优化品级下的最优效果。可以看到，ActNN 不但能开到最大的 batch size（即最省内存），同时在全部 batch size 下都比 baseline 的练习速率更快。

我们还对更多的网络举行了测试。在同样的内存限定下，ActNN 可以将 batch size 扩大 6-14 倍，将模子尺寸大概输入图片扩大 6-10 倍。细致的试验设置和效果拜见文末的论文链接。

两行代码即可在 PyTorch 中利用

import actnn
model = actnn.QModule(model)

ActNN 提供了一个主动模子转换封装。只需在练习剧本里插入两行代码，即可将一般的 PyTorch 模子转换为利用 ActNN 的模子。同时，ActNN 也提供了更高级的 API 支持定制化的利用场景。

更多的例子拜见 github 链接。我们提供了在图像辨认、图像支解、物体检测，以及自监视进修等多个使命上利用 actnn 的完备例子和练习记载，接待试用！