QLoRA技术细节：4-bit NormalFloat量化与双重量化（Double Quantization）的实现

大家好，今天我们来深入探讨QLoRA的核心技术：4-bit NormalFloat (NF4) 量化和双重量化 (Double Quantization)。QLoRA通过这些技术，极大地降低了大型语言模型的显存占用，使得在消费级硬件上微调大型模型成为可能。我们将从理论基础入手，逐步剖析NF4量化的原理、双重量化的实现细节，并提供相应的代码示例，帮助大家理解和应用这些技术。

一、量化技术概述

量化是一种模型压缩技术，通过降低模型参数的精度来减少存储空间和计算复杂度。常见的量化方法包括：

线性量化 (Linear Quantization): 将浮点数映射到整数，并使用比例因子和零点进行转换。
对数量化 (Logarithmic Quantization): 将浮点数取对数后再进行线性量化，适用于数值范围跨度较大的情况。
非线性量化 (Non-linear Quantization): 使用非线性函数进行映射，例如 QLoRA 中使用的 NF4 量化。

量化的核心在于找到一种合适的映射关系，尽可能减少量化带来的信息损失，保证模型性能。

二、4-bit NormalFloat (NF4) 量化

NF4 量化是一种针对正态分布数据优化的 4-bit 量化方案。与传统的线性量化相比，NF4 量化能够更好地保留原始数据的分布特征，从而减少量化误差。

2.1 NF4 量化的基本原理

NF4 量化基于以下观察：大型语言模型的权重通常近似服从正态分布。因此，可以利用正态分布的特性来优化量化方案。具体来说，NF4 量化使用以下步骤：

归一化 (Normalization): 将原始浮点数数据归一化到 [-1, 1] 区间。
量化映射 (Quantization Mapping): 将归一化后的数据映射到 4-bit 整数，使用一个特殊设计的映射表，这个表依据正态分布的累积分布函数 (CDF) 来构建。
反量化 (Dequantization): 将 4-bit 整数映射回浮点数，使用与量化映射相反的过程。

关键在于如何构建这个映射表。理想的映射表应该使得量化误差最小化。NF4 量化通过优化量化级别 (quantization level) 的位置来实现这一点。

2.2 NF4 量化的数学描述

假设原始浮点数数据为 x，归一化后的数据为 x'，量化级别为 q，反量化后的数据为 x''。

归一化: x' = x / max(abs(x)) (其中 max(abs(x)) 表示绝对值的最大值)
量化: q = round((x' + 1) / 2 * (2^4 - 1)) (将 [-1, 1] 映射到 [0, 15])
反量化: x'' = 2 * q / (2^4 - 1) - 1

实际上，NF4使用预计算好的量化值表，避免在线计算。

2.3 NF4 量化值的计算

NF4量化值的核心是依据正态分布的累积分布函数(CDF)来确定量化级别。具体步骤如下：

计算分位数: 将区间[0, 1]分成2^4=16个等份，计算每个等份对应的分位数。这些分位数就是标准正态分布下，累积概率分别为i/16 (i=0, 1, …, 15)的数值。
缩放和平移: 将计算出的分位数缩放到[-1, 1]区间。这一步通常会将分位数乘以一个缩放因子，确保量化值的范围与原始数据的范围一致。

以下是用Python代码计算NF4量化值的示例:

import torch
from scipy.stats import norm

def compute_nf4_quant_values(num_bits=4):
    """
    计算NF4量化值。

    Args:
        num_bits: 量化位数，默认为4。

    Returns:
        一个包含量化值的torch.Tensor。
    """
    # 计算量化级别的数量
    num_levels = 2 ** num_bits

    # 计算分位数，使用正态分布的逆累积分布函数 (ppf)
    quantiles = [norm.ppf(i / num_levels) for i in range(1, num_levels)]

    # 将分位数转换为torch.Tensor
    quantiles = torch.tensor(quantiles)

    # 归一化到 [-1, 1] 区间
    quantiles = quantiles / torch.max(torch.abs(quantiles))  # 缩放

    return quantiles

# 示例
nf4_quant_values = compute_nf4_quant_values()
print("NF4 Quantization Values:", nf4_quant_values)

这段代码使用scipy.stats.norm.ppf函数计算标准正态分布的逆累积分布函数，得到分位数，然后将分位数归一化到[-1, 1]区间。

表格：NF4量化级别的示例

量化级别 (Index)	累积概率 (CDF)	标准正态分布分位数	归一化后的量化值
1	1/16	-1.0364	-0.6532
2	2/16	-0.6745	-0.4256
3	3/16	-0.4307	-0.2716
4	4/16	-0.2533	-0.1597
5	5/16	-0.0967	-0.0610
6	6/16	0.0967	0.0610
7	7/16	0.2533	0.1597
8	8/16	0.4307	0.2716
9	9/16	0.6745	0.4256
10	10/16	1.0364	0.6532
11	11/16	1.4051	0.8872
12	12/16	1.8627	1.1753
13	13/16	2.5758	1.6245
14	14/16	4.1214	2.5998
15	15/16	无穷大	无穷大

注意：实际应用中，会进行截断，防止无穷大值。并且，会预先计算好这些量化值，存储在查找表中，加速量化过程。

2.4 NF4 量化的代码实现

以下是使用 PyTorch 实现 NF4 量化的示例代码：

import torch

def quantize_nf4(x, quant_values):
    """
    使用 NF4 量化张量。

    Args:
        x: 要量化的张量。
        quant_values: NF4量化值，由compute_nf4_quant_values函数生成。

    Returns:
        量化后的张量 (torch.int8)。
    """
    # 找到绝对值的最大值
    abs_max = torch.max(torch.abs(x))

    # 归一化
    x_normalized = x / abs_max

    # 量化
    distances = torch.abs(x_normalized.unsqueeze(-1) - quant_values.unsqueeze(0))
    q = torch.argmin(distances, dim=-1).to(torch.int8)

    return q, abs_max

def dequantize_nf4(q, abs_max, quant_values):
    """
    反量化 NF4 张量。

    Args:
        q: 量化后的张量 (torch.int8)。
        abs_max: 绝对值的最大值。
        quant_values: NF4量化值，由compute_nf4_quant_values函数生成。

    Returns:
        反量化后的张量 (torch.float32)。
    """
    x_dequantized = quant_values[q.long()] * abs_max  # 使用long类型索引

    return x_dequantized

# 示例
# 假设有一个权重张量
weights = torch.randn(1024, 1024, dtype=torch.float32)

# 计算NF4量化值
nf4_quant_values = compute_nf4_quant_values()

# 量化
quantized_weights, abs_max = quantize_nf4(weights, nf4_quant_values)

# 反量化
dequantized_weights = dequantize_nf4(quantized_weights, abs_max, nf4_quant_values)

# 打印量化前后的数据类型和形状
print("Original weights dtype:", weights.dtype, "shape:", weights.shape)
print("Quantized weights dtype:", quantized_weights.dtype, "shape:", quantized_weights.shape)
print("Dequantized weights dtype:", dequantized_weights.dtype, "shape:", dequantized_weights.shape)

# 计算量化误差 (可选)
quantization_error = torch.mean(torch.abs(weights - dequantized_weights))
print("Quantization Error:", quantization_error)

这段代码首先定义了 quantize_nf4 函数，该函数将浮点数张量量化为 4-bit 整数张量，并返回量化后的张量和绝对值的最大值。然后，定义了 dequantize_nf4 函数，该函数将 4-bit 整数张量反量化为浮点数张量。最后，代码演示了如何使用这两个函数对一个随机张量进行量化和反量化，并计算量化误差。注意索引必须是Long类型。

三、双重量化 (Double Quantization)

双重量化是一种进一步降低显存占用的技术，它将量化过程中产生的量化参数 (例如，每个 block 的最大绝对值) 也进行量化。

3.1 双重量化的基本原理

在标准的 NF4 量化中，每个 block (例如，一个矩阵的行或列) 都有一个最大绝对值 (abs_max) 用于归一化。这些 abs_max 值本身也是浮点数，需要占用一定的显存空间。双重量化的思想是将这些 abs_max 值也进行量化，从而进一步降低显存占用。

3.2 双重量化的实现细节

双重量化的实现通常包括以下步骤：

Block 划分: 将模型权重划分为多个 block。
一级量化 (First Quantization): 对每个 block 的权重进行 NF4 量化，得到量化后的权重和每个 block 的 abs_max 值。
二级量化 (Second Quantization): 对所有 block 的 abs_max 值进行量化，例如使用 8-bit 线性量化。这将产生量化后的 abs_max 值，以及用于反量化 abs_max 值的比例因子和零点。
存储: 存储量化后的权重、量化后的 abs_max 值、以及用于反量化 abs_max 值的比例因子和零点。

反量化过程则与量化过程相反：

反量化 abs_max: 使用比例因子和零点将量化后的 abs_max 值反量化为浮点数。
反量化权重: 使用反量化后的 abs_max 值将量化后的权重反量化为浮点数。

3.3 双重量化的代码实现

以下是使用 PyTorch 实现双重量化的示例代码：

import torch

def quantize_double(x, n_bits_first=4, n_bits_second=8):
    """
    实现双重量化。

    Args:
        x: 要量化的张量。
        n_bits_first: 第一级量化的位数 (NF4)。
        n_bits_second: 第二级量化的位数 (abs_max)。

    Returns:
        量化后的张量 (torch.int8), 量化后的 abs_max (torch.int8),
        abs_max 的比例因子和零点。
    """

    # 第一级量化 (NF4)
    nf4_quant_values = compute_nf4_quant_values(n_bits_first)
    quantized_x, abs_max = quantize_nf4(x, nf4_quant_values)

    # 第二级量化 (线性量化 abs_max)
    scale, zero_point = linear_quantize_params(abs_max, n_bits_second)
    quantized_abs_max = torch.round((abs_max / scale) + zero_point).to(torch.int8)

    return quantized_x, quantized_abs_max, scale, zero_point

def dequantize_double(quantized_x, quantized_abs_max, scale, zero_point, n_bits_first=4):
    """
    反量化双重量化后的张量。

    Args:
        quantized_x: 量化后的张量 (torch.int8)。
        quantized_abs_max: 量化后的 abs_max (torch.int8)。
        scale: abs_max 的比例因子。
        zero_point: abs_max 的零点。
        n_bits_first: 第一级量化的位数 (NF4)。

    Returns:
        反量化后的张量 (torch.float32)。
    """

    # 反量化 abs_max
    abs_max = (quantized_abs_max.float() - zero_point) * scale

    # 反量化权重
    nf4_quant_values = compute_nf4_quant_values(n_bits_first)
    dequantized_x = dequantize_nf4(quantized_x, abs_max, nf4_quant_values)

    return dequantized_x

def linear_quantize_params(data, n_bits):
  """
  计算线性量化的比例因子和零点。

  Args:
      data: 要量化的数据。
      n_bits: 量化位数。

  Returns:
      比例因子和零点。
  """
  min_val = torch.min(data)
  max_val = torch.max(data)
  qmin = 0.
  qmax = 2.**n_bits - 1.

  scale = (max_val - min_val) / (qmax - qmin)

  zero_point = qmin - min_val / scale

  return scale, zero_point

# 示例
# 假设有一个权重张量
weights = torch.randn(1024, 1024, dtype=torch.float32)

# 量化
quantized_weights, quantized_abs_max, scale, zero_point = quantize_double(weights)

# 反量化
dequantized_weights = dequantize_double(quantized_weights, quantized_abs_max, scale, zero_point)

# 打印量化前后的数据类型和形状
print("Original weights dtype:", weights.dtype, "shape:", weights.shape)
print("Quantized weights dtype:", quantized_weights.dtype, "shape:", quantized_weights.shape)
print("Quantized abs_max dtype:", quantized_abs_max.dtype, "shape:", quantized_abs_max.shape)
print("Dequantized weights dtype:", dequantized_weights.dtype, "shape:", dequantized_weights.shape)

# 计算量化误差 (可选)
quantization_error = torch.mean(torch.abs(weights - dequantized_weights))
print("Quantization Error:", quantization_error)

这段代码首先定义了 quantize_double 函数，该函数实现双重量化，包括 NF4 量化和线性量化 abs_max 值。然后，定义了 dequantize_double 函数，该函数实现双重量化的反量化。 linear_quantize_params 计算线性量化的参数。最后，代码演示了如何使用这两个函数对一个随机张量进行双重量化和反量化，并计算量化误差。注意，这里的双重量化只是一个简单的示例，实际应用中可能需要更复杂的实现方式，例如使用不同的量化方法对 abs_max 值进行量化。

四、QLoRA 中的应用

QLoRA 将 NF4 量化和双重量化应用于大型语言模型的微调过程中。具体来说，QLoRA 使用以下步骤：

冻结预训练模型: 冻结预训练模型的大部分参数，只保留少部分可训练参数 (LoRA 适配器)。
量化权重: 将预训练模型的权重进行 NF4 量化，并将量化后的权重存储在 GPU 显存中。
双重量化: 对 NF4 量化产生的 abs_max 值进行量化，进一步降低显存占用。
微调 LoRA 适配器: 在微调过程中，只更新 LoRA 适配器的参数，而保持量化后的权重不变。

通过这种方式，QLoRA 能够在低显存环境下微调大型语言模型，而不会显著降低模型性能。

五、NF4和双重量化的优势与局限

5.1 优势

显著降低显存占用: NF4 量化可以将模型权重压缩到 4-bit，双重量化可以进一步降低显存占用，使得在消费级硬件上微调大型模型成为可能。
良好的性能保持: NF4 量化针对正态分布数据进行了优化，能够更好地保留原始数据的分布特征，从而减少量化误差，保证模型性能。

5.2 局限

量化误差: 量化是一种有损压缩技术，不可避免地会引入量化误差。虽然 NF4 量化能够减少量化误差，但在某些情况下仍然可能影响模型性能。
实现复杂度: NF4 量化和双重量化的实现相对复杂，需要仔细设计量化方案和反量化过程。
硬件支持: 某些硬件可能对低精度计算的支持不够友好，可能导致性能下降。

六、量化级别和量化误差

NF4通过优化量化级别的位置来最小化量化误差。更细粒度的量化级别(比如更多bit)可以减少量化误差，但会增加存储成本。选择合适的量化级别需要在性能和存储之间进行权衡。

总结

NF4 量化和双重量化是 QLoRA 中的核心技术，它们通过降低模型参数的精度来减少显存占用，使得在消费级硬件上微调大型模型成为可能。 NF4 量化针对正态分布数据进行了优化，能够更好地保留原始数据的分布特征，减少量化误差。双重量化则将量化过程中产生的量化参数也进行量化，进一步降低显存占用。掌握这些技术细节，有助于我们更好地理解和应用 QLoRA，并在实际项目中进行优化。