低比特优化器

论文《具有4位状态的内存高效优化器》的官方实现。

优化器状态是训练神经网络时内存消耗的主要来源,限制了在给定内存预算内可训练的最大模型。将优化器状态从32位浮点数压缩到更低位宽是减少训练内存占用的有效方法,而当前可实现的最低位宽是8位。在这项工作中,我们通过对一阶和二阶动量进行详细的实证分析,将优化器状态位宽降低到4位。具体来说,我们发现动量具有复杂的离群值模式,当前的分块量化无法准确近似。我们使用更小的块大小,并提出同时利用行和列信息来实现更好的量化。我们进一步发现了量化二阶动量时的零点问题,并通过排除零点的线性量化器解决了这个问题。我们的4位优化器在广泛的基准测试中进行了评估,包括自然语言理解、机器翻译、图像分类和指令微调。在所有任务中,我们的优化器都能达到与全精度对应版本相当的精度,同时具有更好的内存效率。

安装

要求 Python >= 3.7 + CUDA >= 11.0 + torch >= 1.13.0

安装运行:

git clone https://github.com/thu-ml/low-bit-optimizers.git
pip install -v -e .

使用方法

使用4位优化器

要开始使用4位优化器,只需将现有的优化器替换为我们的4位优化器之一:4位AdamW、4位Factor或4位AdamW(融合版)。

import lpmm

# 注释掉或删除旧的优化器
# optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, betas=(0.9, 0.999))

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

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

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

目前支持的优化器有Adam(AdamW)和SGD。

修改量化超参数

要修改非融合优化器的量化配置(例如,归一化函数、量化映射、位数等),创建一个新的配置文件,并使用qconfig参数将其文件路径传递给优化器。示例配置可以在lpmm/configs目录中找到。默认情况下,非融合优化器的量化配置在lpmm/configs/default.yml中指定,而融合优化器的配置在lpmm/configs/2nd_moment_group_128.yml中指定。融合优化器的配置目前是固定的,无法更改。

要使用新的配置文件,请参照以下示例:

config_path = f"configs/default.yml" # 配置文件路径
optimizer = lpmm.optim.AdamW(model.parameters(), lr=1e-3, betas=(0.9, 0.999), qconfig=config_path)

常用的超参数及其可能的值包括:

SCALE_TYPE(归一化函数): tensor, dim0, dim1, group, rank1
QUANT_TYPE(量化映射): nonlinear, power-1, power-2
BITS: 4, 5, 6, 7, 8
ENABLE(是否量化状态): True, False

我们建议使用BITS = 4或8。

覆盖特定参数的量化启用

要使用32位精度而不是量化来优化某些参数,请使用override_quantize_enable方法,如下所示:

optimizer = lpmm.optim.AdamW(model.parameters(), lr=1e-3, betas=(0.9, 0.999))
optimizer.override_quantize_enable(module, param_name, enable=False)

在这个例子中,module是包含参数的模块,param_name是你希望用32位精度优化的参数名称。设置enable=False将防止对指定参数进行量化。