本文将介绍来自加州伯克利大学的 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,一个新的基于压缩的内存节省框架。在提供理论证明的同时,我们基于