4位优化器：推动内存效率的新境界

low-bit-optimizers

突破内存瓶颈：4位优化器的崛起

在深度学习领域，模型规模的不断扩大带来了巨大的计算和存储挑战。随着神经网络变得越来越复杂，优化器状态的内存消耗已成为限制可训练模型大小的主要因素。为了突破这一瓶颈，研究人员一直在寻求创新的方法来压缩优化器状态，而最新的突破性研究将优化器状态的位宽从32位浮点数降低到了仅4位，为大规模模型训练开辟了新的可能性。

从32位到4位：优化器状态压缩的演进

传统上，神经网络训练中使用的优化器状态通常以32位浮点数存储。这种高精度虽然确保了计算的准确性，但也导致了巨大的内存开销。随着模型规模的增长，这种开销变得越来越难以承受。近年来，研究人员开始探索将优化器状态压缩到更低位宽的可能性，8位优化器的成功实现是这一方向的重要里程碑。

然而，清华大学机器学习组的研究人员并未止步于此。他们通过对一阶和二阶动量进行深入的经验分析，成功将优化器状态的位宽进一步降低到4位。这一突破性成果不仅大幅减少了内存占用，还在保持模型性能的同时，为更大规模的模型训练提供了可能。

技术创新：克服4位量化的挑战

将优化器状态压缩到4位并非易事。研究人员发现，动量值存在复杂的离群值模式，这使得传统的块级量化方法难以准确近似。为了解决这个问题，他们采用了更小的块大小，并创新性地利用了行级和列级信息来改进量化效果。

此外，研究人员还识别并解决了二阶动量量化中的零点问题。他们提出了一种不包含零点的线性量化器，进一步提高了量化的准确性。这些技术创新的结合，使得4位优化器能够在各种基准测试中实现与全精度对应物相当的准确性。

4位优化器示意图

广泛的应用与验证

为了验证4位优化器的有效性，研究团队在多个领域的基准测试中进行了全面评估，包括自然语言理解、机器翻译、图像分类和指令调优等任务。结果令人振奋：在所有这些任务中，4位优化器都能够达到与全精度对应物相当的准确性，同时显著降低了内存占用。

这一成果的意义不仅限于学术研究。它为工业界提供了一种实用的解决方案，有助于在有限的内存预算内训练更大的模型。特别是在边缘设备和资源受限的环境中，4位优化器的应用可能会带来革命性的变化，使得在这些设备上部署和训练复杂模型成为可能。

实际应用：Low-bit Optimizers 库

为了让更多研究者和开发者能够便捷地使用这项技术，研究团队开发并开源了 Low-bit Optimizers 库。这个 PyTorch 库提供了多种 4 位优化器的实现，包括 4-bit AdamW、4-bit Factor 和 4-bit AdamW (fused) 等。

使用这个库非常简单，只需要几行代码就可以将现有的优化器替换为 4 位版本：

import lpmm

# 使用 4-bit AdamW
optimizer = lpmm.optim.AdamW(model.parameters(), lr=1e-3, betas=(0.9, 0.999))

# 或者使用 4-bit Factor
optimizer = lpmm.optim.AdamW(model.parameters(), lr=1e-3, betas=(0.9, 0.999), factor_second_moment=True)

# 或者使用 4-bit AdamW (fused)
optimizer = lpmm.optim.AdamW(model.parameters(), lr=1e-3, betas=(0.9, 0.999), fused=True)

这个库不仅提供了开箱即用的 4 位优化器，还允许用户通过修改配置文件来调整量化参数，以适应不同的应用场景。例如，用户可以更改标准化函数、量化映射、位数等参数，以在内存效率和计算精度之间找到最佳平衡点。

Low-bit Optimizers 库使用示例