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

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

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

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

1.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

2.节省训练内存的方法

目前，节省训练内存的方法主要有三类：1. 重计算（Gradient checkpointing/Rematerialization) 2. 使用 CPU 内存进行交换 (swapping) 和 3. 使用分布式训练将 Tensor 分散存储在多个 GPU 上。这三类方法互相不冲突，可以结合使用。大部分机器学习框架对这些方法都提供了一些支持，也有不少相关的论文。但是，想要高效、自动化地实现这些策略并不容易。与已有方法不同，我们提出了 ActNN，一个新的基于压缩的内存节省框架。在提供理论证明的同时