2024/08/05 14:25

70倍极致压缩！大模型的检查点再多也不怕

AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

该论文的作者均来自于华为诺亚实验室，第一作者为李文硕，通讯作者为王云鹤和陈醒濠。相关团队团队近年来在ICML、CVPR、NeurIPS、ICCV、ECCV等顶会上有多项代表性工作发表，在高效大语言模型、视觉模型等领域都有丰富的成果产出，和知名高校和科研机构合作广泛。

大模型作为当下 AI 工业界和学术界当之无愧的「流量之王」，吸引了大批学者和企业投入资源去研究与训练。随着规模越做越大，系统和工程问题已经成了大模型训练中绕不开的难题。例如在 Llama3.1 54 天的训练里，系统会崩溃 466 次，平均 2.78 小时一次！

那么，频繁存储检查点就显得十分必要。但存储检查点本身也是一个大工程。

Meta 做了很多努力来加速存储检查点时间和增加存储频率，来对抗频繁出现的系统失效。但频繁存储也意味着大量的存储资源开销，其训练集群配备了 240PB 的 SSD 来应对这一挑战，光存储这一项的耗费就要亿元！

华为诺亚的 ExCP 方法也就应运而生，为了应对存储带来的巨大开销，他们提出了极致压缩检查点技术，能够无损压缩模型 70 倍，大幅降低训练中的存储开销。

代码目前已经开源，在 Apache 2.0 框架下发布，issue 中已经有小伙伴成功复现了结果。

文章地址：https://arxiv.org/abs/2406.11257
仓库地址：https://github.com/Gaffey/ExCP

方法也很有创新性，文章中提到了两个重要的概念，一个是利用训练中检查点的残差信息，通过时间序列上信息的稀疏性实现更高的剪枝比例；另一个是将优化器和权重联合起来进行压缩，实现整体的高压缩率。

具体方法

1. 检查点残差

在训练过程中，当前的参数可以看作上一个检查点存储的权重加上逐次迭代时梯度更新的总和，这部分是相对稀疏的，包含的信息量较少，因此对这一残差进行压缩，可以获得更好的压缩比例。而与此相反的，优化器中存储的动量是梯度一阶矩和二阶矩的滑动平均值，对于一阶矩来说，它的滑动平均默认的参数是 0.9，在数百到数千个迭代之后与上一次检查点存储的内容已经没有太大的关联，所以对于优化器直接压缩其本身的值而非残差。最终待压缩的检查点表示为

2. 权重 - 优化器动量联合压缩

目前已有的模型压缩相关的工作一般只关注于模型的推理性能，或者是模型最终存储检查点的大小，而不关注模型在整个训练过程中对储存空间的开销。因而已有工作只对权重进行压缩，而忽略了 Adam 等常见优化器中实际上存储了两倍于权重数量的动量。这一工作一方面将两者一起进行了压缩，显著提升了整体的压缩比例；另一方面也利用了权重和优化器动量的关联性，进一步提升彼此的压缩比例。

权重剪枝：由于剪枝的权重是残差值，优化器动量的二阶矩可以大致表示在过去一段时间内权重残差值的变化幅度，所以可以使用优化器动量的二阶矩作为指标来确定不同层的剪枝比例。剪枝策略如下文公式所示

式中，W 和分别表示权重和二阶矩。

优化器动量剪枝：对于动量剪枝，可以使用一阶矩作为指示器来进行剪枝，论文中有关于可收敛性的一个简要证明。同时，如果一个位置的权重已经被剪枝，那么对应位置的优化器动量也应该同步被处理，所以剪枝策略如下文公式所示

式中，

表示一阶矩。

3. 整体压缩流程

整体压缩流程如 Algorithm 1 所示，依次进行计算权重残差 / 联合压缩 / 非均匀量化 / 编码压缩等步骤，得到最终的压缩结果。

而恢复出检查点完整文件的流程则如 Algorithm 2 所示，进行解压缩之后，首先从非均匀量化后存储的码本和下标中恢复出浮点结果，然后再与基准权重（上一个检查点的原始权重或恢复出的重建权重）相加，得到检查点完整文件。而恢复出整个训练流程中的检查点文件的流程如 Algorithm 3 所示，在完成训练后只保存初始化权重的随机种子和每个检查点存储的压缩结果，然后依次对检查点进行恢复以得到完整的检查点序列，以供从其中选择某个或多个检查点恢复训练 / 进行测试等。

实验结果

文章中不仅对于大语言模型做了评估，在 ViT-L32 这样较大的视觉模型上这一方法也能取得很好的效果。

从消融实验里也可以看出，采用残差剪枝的方法大大减少了剪枝带来的损失。

文章中还提供了大语言模型压缩前后问答的样例，可以看到压缩本身对于模型的问答能力也没有造成损害。

工程ExCP华为诺亚方舟实验室

相关技术

机器学习物体识别深度学习人脸识别图像分割

权重技术

线性模型中特征的系数，或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0，则相应的特征对模型来说没有任何贡献。

来源：Google AI Glossary

基准技术

一种简单的模型或启发法，用作比较模型效果时的参考点。基准有助于模型开发者针对特定问题量化最低预期效果。

来源：Google ML Glossary

参数技术

在数学和统计学裡，参数（英语：parameter）是使用通用变量来建立函数和变量之间关系（当这种关系很难用方程来阐述时）的一个数量。

来源：维基百科

剪枝技术

剪枝顾名思义，就是删去一些不重要的节点，来减小计算或搜索的复杂度。剪枝在很多算法中都有很好的应用，如：决策树，神经网络，搜索算法,数据库的设计等。在决策树和神经网络中，剪枝可以有效缓解过拟合问题并减小计算复杂度；在搜索算法中，可以减小搜索范围，提高搜索效率。

来源：Wikipedia

收敛技术

在数学，计算机科学和逻辑学中，收敛指的是不同的变换序列在有限的时间内达到一个结论（变换终止），并且得出的结论是独立于达到它的路径（他们是融合的）。通俗来说，收敛通常是指在训练期间达到的一种状态，即经过一定次数的迭代之后，训练损失和验证损失在每次迭代中的变化都非常小或根本没有变化。也就是说，如果采用当前数据进行额外的训练将无法改进模型，模型即达到收敛状态。在深度学习中，损失值有时会在最终下降之前的多次迭代中保持不变或几乎保持不变，暂时形成收敛的假象。

来源：Wikipedia Google ML glossary

SSD技术

一种计算机视觉模型。论文发表于 2015 年（Wei Liu et al.）

来源：介绍文章

动量技术

优化器的一种，是模拟物理里动量的概念，其在相关方向可以加速SGD，抑制振荡，从而加快收敛

来源：An overview of gradient descent optimization algorithms

优化器技术

优化器基类提供了计算梯度loss的方法，并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法，如梯度下降和Adagrad。优化器是提供了一个可以使用各种优化算法的接口，可以让用户直接调用一些经典的优化算法，如梯度下降法等等。优化器（optimizers）类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类，但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer（tensorflow下的优化器包）等等这些算法。

来源：维基百科

法大大机构

深圳法大大网络科技有限公司（www.fadada.com）是国内领先的第三方电子合同平台，主要为金融、房地产、汽车、人力资源服务、教育、保险、第三方支付、旅游、医疗、物流、供应链、B2B、B2C线上交易平台等行业以及政府机构提供电子合同、电子文件签署及存证服务，同时整合提供司法鉴定和律师服务等增值服务。

https://www.fadada.com

机器之心机构

机器之心，成立于2014年，是国内最具影响力、最专业、唯一用于国际品牌的人工智能信息服务与产业服务平台。目前机器之心已经建立起涵盖媒体、数据、活动、研究及咨询、线下物理空间于一体的业务体系，为各类人工智能从业者提供综合信息服务和产业服务。

https://www.jiqizhixin.com/