ExCP: Extreme LLM Checkpoint Compression via Weight-Momentum Joint Shrinking

大型语言模型（LLM）最近在人工智能领域引起了重大关注。然而，这些模型的训练过程在计算和存储容量方面存在重大挑战，因此压缩检查点已成为一个紧迫的问题。在本文中，我们提出了一种新颖的极端检查点压缩（ExCP）框架，它可以显著减少训练检查点所需的存储空间，同时实现几乎无损的性能。我们首先计算相邻检查点的残差，以获得更高压缩比所需的基本但稀疏信息。为了进一步挖掘检查点中的冗余参数，我们提出了一种权重-动量联合收缩方法，利用模型优化期间的另一个重要信息，即动量。特别地，我们利用模型和优化器的信息，丢弃尽可能多的参数，同时保留关键信息以确保最佳性能。此外，我们利用非均匀量化进一步压缩检查点的存储。我们广泛评估了我们提出的ExCP框架，涵盖了从410M到7B参数的几个模型，并展示了显著的存储减少，同时保持了强大的性能。例如，我们在Pythia-410M模型中实现了约70倍的压缩，最终性能在各种下游任务中与原始模型一样准确。代码将在https://github.com/Gaffey/ExCP上提供。

解决问题：论文旨在解决大型语言模型训练过程中的存储和计算挑战，提出一种新的极端检查点压缩（ExCP）框架。

关键思路：论文提出了一种利用相邻检查点的残差和权重-动量联合收缩的方法来压缩检查点，同时利用非均匀量化进一步压缩存储空间。

亮点：论文在多个模型上进行了广泛的评估，证明了ExCP框架可以显著减少存储空间的同时保持强大的性能。他们在Pythia-410M模型上实现了近70倍的压缩，并在多个下游任务上保持了原始模型的准确性。此外，作者提供了开源代码。

相关研究：最近的相关研究包括《Efficient Transformer-based Large-scale Language Models Using Layer-wise Coordination》、《Compressing Large-scale Transformer-based Language Models》等。