QLoRA 的双重量化：对量化常数再次量化以节省显存的实现
大家好，今天我们要深入探讨 QLoRA 中一项至关重要的技术：双重量化。这项技术通过对量化常数进行再次量化，能够显著降低模型在微调期间的显存占用，使得在资源有限的环境下进行大规模语言模型的微调成为可能。
1. 背景：QLoRA 和量化
QLoRA (Quantization-aware Low-Rank Adaptation) 是一种高效的参数高效微调方法，它通过以下关键技术来实现显存节省：
- 4-bit NormalFloat (NF4) 量化: 将预训练模型的权重从 FP16 或 FP32 量化为 4-bit NF4 格式。NF4 是一种针对正态分布数据优化的量化格式，可以更好地保留模型的性能。
- LoRA (Low-Rank Adaptation): 只训练少量（通常是低秩）的适配器参数，并将这些参数添加到冻结的预训练模型中。这大大减少了需要更新的参数数量。
- 分页优化器: 使用 NVIDIA unified memory 将优化器状态分页到 CPU 内存，从而释放 GPU 显存。
- 双重量化: 对量化常数进行再次量化，进一步减少显存占用。
我们今天的重点就是最后一个：双重量化。为了理解它的作用，我们先回顾一下量化过程。
2. 量化的基本原理
量化的目标是将浮点数（例如 FP32 或 FP16）映射到一组离散的整数值，从而减少存储空间。一个简单的线性量化过程可以表示为：
# 浮点数 -> 整数
quantized_value = round( (original_value - zero_point) / scale )
# 整数 -> 浮点数 (反量化)
reconstructed_value = quantized_value * scale + zero_point其中：
original_value是原始的浮点数值。quantized_value是量化后的整数值。scale是一个缩放因子，用于将浮点数范围映射到整数范围。zero_point是零点，用于将浮点数中的零映射到整数范围的中心。
例如，如果我们想将 FP32 数值量化为 8-bit 整数，整数范围就是 [-128, 127]。scale 和 zero_point 的选择至关重要，它们决定了量化误差的大小。
3. QLoRA 中的量化
在 QLoRA 中，预训练模型的权重被量化为 4-bit NF4 格式。这意味着每个权重只需要 4 位来存储，大大减少了显存占用。然而，每个量化后的权重块（例如，一个卷积层的所有权重）都需要存储 scale 和 zero_point。这些 scale 和 zero_point 通常以 FP16 或 FP32 的格式存储，它们仍然会占用相当可观的显存，尤其是在模型非常大的情况下。
4. 双重量化的概念
双重量化的核心思想是：既然 scale 和 zero_point 也是数值，那么我们也可以对它们进行量化，进一步减少显存占用。具体来说，QLoRA 对 scale 和 zero_point 再次进行量化，将它们量化为更低精度的格式，例如 8-bit 整数。
5. 双重量化的实现
以下是一个简化的 Python 代码示例，展示了双重量化的过程：
import torch
def quantize(tensor, num_bits=8):
"""
对张量进行量化.
Args:
tensor: 要量化的张量.
num_bits: 量化的位数.
Returns:
量化后的张量, scale, zero_point
"""
min_val = torch.min(tensor)
max_val = torch.max(tensor)
qmin = 0
qmax = 2**num_bits - 1
scale = (max_val - min_val) / (qmax - qmin)
zero_point = qmin - min_val / scale
zero_point = torch.clamp(zero_point, qmin, qmax).round()
quantized_tensor = torch.round((tensor - min_val) / scale)
quantized_tensor = torch.clamp(quantized_tensor, qmin, qmax).to(torch.uint8) # 使用 uint8 存储
return quantized_tensor, scale, zero_point
def dequantize(quantized_tensor, scale, zero_point):
"""
对量化后的张量进行反量化.
Args:
quantized_tensor: 量化后的张量.
scale: 缩放因子.
zero_point: 零点.
Returns:
反量化后的张量.
"""
return quantized_tensor.float() * scale + zero_point
def double_quantize(scale, zero_point, num_bits=8):
"""
对 scale 和 zero_point 进行二次量化.
Args:
scale: 缩放因子.
zero_point: 零点.
num_bits: 量化的位数.
Returns:
二次量化后的 scale, scale_scale, scale_zero_point, 二次量化后的 zero_point, zero_point_scale, zero_point_zero_point
"""
# 对 scale 进行量化
quantized_scale, scale_scale, scale_zero_point = quantize(scale, num_bits)
# 对 zero_point 进行量化
quantized_zero_point, zero_point_scale, zero_point_zero_point = quantize(zero_point, num_bits)
return quantized_scale, scale_scale, scale_zero_point, quantized_zero_point, zero_point_scale, zero_point_zero_point
def double_dequantize(quantized_scale, scale_scale, scale_zero_point, quantized_zero_point, zero_point_scale, zero_point_zero_point):
"""
对二次量化后的 scale 和 zero_point 进行反量化.
Args:
quantized_scale: 二次量化后的 scale.
scale_scale: scale 的 scale.
scale_zero_point: scale 的 zero_point.
quantized_zero_point: 二次量化后的 zero_point.
zero_point_scale: zero_point 的 scale.
zero_point_zero_point: zero_point 的 zero_point.
Returns:
反量化后的 scale 和 zero_point.
"""
# 对 scale 进行反量化
scale = dequantize(quantized_scale, scale_scale, scale_zero_point)
# 对 zero_point 进行反量化
zero_point = dequantize(quantized_zero_point, zero_point_scale, zero_point_zero_point)
return scale, zero_point
# 示例
if __name__ == '__main__':
# 假设我们有一个 tensor
original_tensor = torch.randn(10)
# 第一次量化
quantized_tensor, scale, zero_point = quantize(original_tensor)
# 双重量化
quantized_scale, scale_scale, scale_zero_point, quantized_zero_point, zero_point_scale, zero_point_zero_point = double_quantize(scale, zero_point)
# 双重反量化
scale_reconstructed, zero_point_reconstructed = double_dequantize(quantized_scale, scale_scale, scale_zero_point, quantized_zero_point, zero_point_scale, zero_point_zero_point)
# 第一次反量化
reconstructed_tensor = dequantize(quantized_tensor, scale_reconstructed, zero_point_reconstructed)
# 打印原始 tensor 和反量化后的 tensor
print("Original Tensor:", original_tensor)
print("Reconstructed Tensor:", reconstructed_tensor)
# 计算均方误差
mse = torch.mean((original_tensor - reconstructed_tensor)**2)
print("Mean Squared Error:", mse)代码解释：
-
quantize(tensor, num_bits): 这个函数将输入的张量tensor量化为num_bits位整数。它计算张量的最小值和最大值，并使用它们来计算scale和zero_point。 然后，它将张量量化为整数，并将结果存储为torch.uint8类型，以节省内存。 -
dequantize(quantized_tensor, scale, zero_point): 这个函数将量化后的张量反量化回浮点数。它使用scale和zero_point将整数值转换回浮点数值。 -
double_quantize(scale, zero_point, num_bits): 这是双重量化的关键函数。它接收scale和zero_point作为输入，并对它们分别进行量化。 这意味着它会为scale计算一个新的scale(称为scale_scale) 和一个zero_point(称为scale_zero_point)，并将scale量化为整数。 同样的操作也适用于zero_point。 此函数返回量化后的scale(quantized_scale)，scale_scale，scale_zero_point, 量化后的zero_point(quantized_zero_point)，zero_point_scale和zero_point_zero_point。 -
double_dequantize(quantized_scale, scale_scale, scale_zero_point, quantized_zero_point, zero_point_scale, zero_point_zero_point): 这个函数将双重量化后的scale和zero_point反量化回其原始值。它首先使用scale_scale和scale_zero_point将量化后的scale反量化，然后使用zero_point_scale和zero_point_zero_point将量化后的zero_point反量化。 -
示例： 示例代码演示了如何使用这些函数对一个随机张量进行量化、双重量化和反量化。 它还计算了原始张量和反量化后的张量之间的均方误差，以衡量量化过程造成的精度损失。
6. 双重量化的优势
- 显著减少显存占用: 通过对
scale和zero_point进行量化，可以进一步减少模型在微调期间的显存占用，使得在资源有限的环境下训练更大的模型成为可能。 - 可忽略的精度损失: 只要选择合适的量化位数和量化方法，双重量化引入的精度损失通常可以忽略不计。
7. 双重量化的挑战
- 额外的计算开销: 双重量化引入了额外的量化和反量化操作，这会增加计算开销。然而，在实际应用中，这些开销通常可以忽略不计，因为量化和反量化操作相对较快。
- 选择合适的量化位数: 为
scale和zero_point选择合适的量化位数是一个需要仔细考虑的问题。 如果量化位数太低，可能会导致较大的精度损失。 如果量化位数太高，则可能无法实现预期的显存节省效果。
8. QLoRA 中的具体实现细节
在 QLoRA 的实际实现中，双重量化通常使用 8-bit 整数来量化 scale 和 zero_point。 此外，QLoRA 还使用一些技巧来进一步优化量化过程，例如：
- 分组量化: 将权重分成多个组，并为每个组计算一个
scale和zero_point。 这可以减少量化误差，并提高模型的性能。 - 动态量化: 在训练过程中动态地调整
scale和zero_point，以适应权重的分布变化。
9. 双重量化与其他显存优化技术的比较
以下是一个表格，比较了双重量化与其他一些常见的显存优化技术：
| 技术 | 优点 | 缺点 |
|---|---|---|
| 量化 | 显著减少显存占用，提高推理速度 | 可能会导致精度损失，需要仔细选择量化方法和位数 |
| 双重量化 | 在量化的基础上进一步减少显存占用，对量化参数进行压缩 | 引入额外的量化和反量化操作，增加计算开销，需要选择合适的量化位数 |
| 混合精度训练 | 使用 FP16 或 BF16 替代 FP32 进行训练，减少显存占用，提高训练速度 | 需要仔细管理精度，可能会导致训练不稳定 |
| 梯度累积 | 通过累积多个 batch 的梯度来模拟更大的 batch size，减少显存占用 | 可能会降低训练速度 |
| 参数共享 | 减少模型中的参数数量，降低显存占用 | 可能会降低模型性能 |
| LoRA | 只训练少量适配器参数，显著减少显存占用 | 需要选择合适的秩和适配器位置 |
10. 未来发展趋势
双重量化作为一种有效的显存优化技术，在未来将继续发挥重要作用。 未来的发展趋势可能包括：
- 更高级的量化方法: 开发更高级的量化方法，例如非对称量化和混合精度量化，以进一步提高量化精度。
- 自适应量化: 根据模型的不同层和权重的分布特点，自适应地选择量化位数和量化方法。
- 硬件加速: 在硬件层面支持双重量化，以提高量化和反量化的速度。
让大模型微调不再是难题
今天我们详细讨论了 QLoRA 中的双重量化技术，了解了它的原理、实现方式、优势和挑战。双重量化通过对量化常数进行再次量化，显著减少了显存占用，使得在资源有限的环境下进行大规模语言模型的微调成为可能。希望今天的讲解能够帮助大家更好地理解 QLoRA 和双重量化，并在实际应用中发挥它们的作用。