大模型检查点无损压缩70倍,华为诺亚开源新技术ExCP

原文标题:70倍极致压缩!大模型的检查点再多也不怕

原文作者:机器之心

冷月清谈:

**关键要点:**

具体方法:

  1. 检查点残差压缩:利用检查点之间的梯度更新的稀疏性,压缩检查点残差信息。
  2. 权重-优化器动量联合压缩:将优化器中动量与权重结合压缩,提高总体压缩率。
  3. 非均匀量化和编码压缩:进一步降低压缩文件大小。

实验结果:

  • 在大语言模型和ViT-L32视觉模型上均取得了显著的压缩效果,无损压缩率达70%。
  • 消融实验表明,残差剪枝方法有效减少了剪枝损失。
  • 压缩前后模型的问答能力保持良好。



怜星夜思:


1、ExCP方法中,是如何利用剪枝策略提升压缩比例的?
2、ExCP方法与其他模型压缩技术相比,优势和劣势是什么?
3、ExCP方法在实际应用中有哪些潜在的挑战?

原文内容

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

该论文的作者均来自于华为诺亚实验室,第一作者为李文硕,通讯作者为王云鹤和陈醒濠。相关团队团队近年来在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 这样较大的视觉模型上这一方法也能取得很好的效果。


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


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



© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

ExCP方法提出了一种自适应的剪枝策略,根据每个训练步骤中优化器动量的值动态调整剪枝比例。这种动态剪枝机制可以显著提高压缩效果,同时保持模型的性能。

不过,ExCP方法也存在一些劣势。与量化等其他压缩技术相比,ExCP方法的压缩过程更加复杂,需要对模型进行额外的处理。此外,ExCP方法可能会增加训练时间,因为需要对检查点进行额外的压缩和解压缩。

ExCP方法与其他模型压缩技术相比,优势在于实现了更高的无损压缩率(高达70%),同时又不损害模型的性能。此外,ExCP方法还可以联合压缩权重和优化器动量,这在其他技术中并不常见。

期待该团队未来有更多创新性的工作,为模型压缩领域做出更大的贡献。

在实际应用中,ExCP方法可能会面临一些挑战。例如,如何将该方法与其他优化技术(如混合精度训练)相结合,以进一步提高模型训练的效率和压缩率。

期待未来有更多关于ExCP方法的研究和应用,相信它将在解决大模型训练中的存储和计算成本问题方面发挥重要作用。

顺便说一句,ExCP方法的代码已经开源啦!大家可以去体验一下它的强大威力。GitHub - Gaffey/ExCP: Official implementation of ICML 2024 paper "ExCP: Extreme LLM Checkpoint Compression via Weight-Momentum Joint Shrinking".

通过剪枝权重残差,根据优化器动量的二阶矩作为指标,确定不同层的剪枝比例。这一策略充分利用了权重和优化器动量的关联性,提升了整体的压缩效果。

此外,ExCP方法可能需要调整才能应用于不同的模型架构和数据集。需要进行广泛的实验和分析,以确定ExCP方法在不同场景下的最佳实践。

看到ExCP方法在社区中引起热烈讨论,感觉很欣慰。这种开放式的交流和协作有助于推动模型压缩领域的发展。

加油,科研工作者们!

感觉ExCP方法有点像一个高级版的剪枝技术,既利用了权重的稀疏性,又考虑了优化器的信息,实现了更高的压缩效果。

另外一个挑战是ExCP方法的推理速度。压缩后的模型可能需要更多的计算资源来进行推理,这可能会影响实际应用中的部署和使用。

总的来说,ExCP方法是一种很有前景的模型压缩技术,特别适用于需要高压缩率和大模型训练的场景。

从论文中可以看出,ExCP方法的研究团队在模型压缩领域颇有建树,他们在ICML、CVPR等顶会上发表了许多高质量的研究成果。

值得一提的是,ExCP方法不仅适用于大语言模型,在ViT-L32等视觉模型上也取得了良好的压缩效果。这表明该方法的适用范围较广,具有很强的通用性。

通过剪枝优化器动量,使用一阶矩作为指示器,并同步剪枝对应的权重残差,进一步提升了压缩率。这充分考虑了优化器动量和权重之间的关系,实现了更加高效的压缩。

论文中还深入分析了剪枝策略与模型性能之间的关系,为未来改进剪枝策略提供了理论基础。