简介

LORA (Low-Rank Adaptation) 是一种高效微调大型预训练模型的方法。它通过冻结预训练模型的权重，并在Transformer架构的每一层中引入可训练的秩分解矩阵，显著减少了可训练参数的数量，从而确保了更加高效的适应过程。具体来说，它将一个大矩阵分解为两个低秩矩阵的乘积，即 weight[ho] = w1[hr] @ w2[ro]，其中 r 是秩，是一个关键的超参数。通常，r 的值设置为4、8或12，以平衡表达力和计算效率。

QLoRA 是LoRA的量化版本，它结合了量化技术来进一步减少内存和计算成本。在QLoRA中，LoRA的可训练低秩矩阵 w1 和 w2 保持不量化，以便进行反向传播和优化。然而，原始模型的权重 W 被冻结并量化，以减少内存占用。

low-rank低秩矩阵分解，只对模型的部分权重进行更新，而不是全量微调，可大幅较少参数更新量，减少计算量。

两个低秩矩阵替代一个大矩阵，将 权重 weight 分解为 两个 小张量 weight[ho] = w1[hr]@w2[ro], 权重更新时对 两个低秩 矩阵w1/w2(适配器权重)进行更新。其中 r称为 rank，是一个重要的超参数。r越大，可训练参数越多，表达力越强，不过一般也不需要太大， r = 4,8,12 是常见的设置。

此外，在 Transformer block 中，对哪些参数矩阵的更新量进行模拟？一般是选择自注意力层的 Q, V 矩阵；或者范围更大一些：Q, K, V, O（输出矩阵）。

LoRA : Y = X@W + s * X * w1 * w2 后面为W梯度，

QLoRA : Y = X@DoubleQuant(c1,c2,W_nf4) + s * X * w1 * w2

LoRA的参数 w1 w2 是不量化的，因为它们需要反向传播优化。而原始模型的参数 W 是freeze的，因此可以量化。

量化技术 在QLoRA中扮演了至关重要的角色。特别是，QLoRA采用了非均匀量化（如NF4）来量化权重，这有助于减少量化误差并保留更多的信息。非均匀量化使用量化表将浮点数映射到离散的量化级别，从而实现了更高的压缩率和更低的量化误差。

在QLoRA中，模型的权重以NF4格式存储，但在计算时使用BF16格式。这意味着在每次前向传播之前，需要将NF4格式的权重反量化为BF16格式进行计算。计算完成后，再将结果量化为NF4格式进行存储。这种双重反量化过程进一步节省了量化常数的空间占用。

针对离群值/异常值：分块量化，不同数据块，有不同的量化系数 c，如 分N块单独量化，需要N个量化系数c2。

双重反量化：w存储为 nf4类型，通过两次反量化 dequant(dequant(c1,c2), W_nf4)，将存储数据类型转换为计算数据类型， 进一步节省量化常数的空间占用；

c2为分块量化系数，每块一个量化系数，为降低量化系数存储，利用c1将量化的量化系数c2反量化为浮点，再将量化的权重量化为浮点。

QLoRA 中，模型的权重有两种格式：用 NF4 存储；用 BF16 计算。需要用相应权重计算前向传播时，对 NF4 的权重反量化为 BF16；计算结束后，再量化为 NF4。

QLoRA 步骤：

• 初始量化：首先，大型语言模型 (LLM) 权重 被量化为 4 位，显着减少了内存占用。

  • 量化权重
  • 量化 量化系数

• LoRA微调：然后，执行 LoRA 训练。

  • 反量化 量化系数
  • 反量化 权重
  • 计算Y 梯度更新 前向

bitsandbytes 实现原理

BitsandBytes (bnb) 是一个用于实现QLoRA的库。它提供了高效的4位和8位量化算法，以及相关的量化配置和模型替换功能。通过配置bnb的量化参数，可以轻松地将模型中的nn.Linear层替换为量化的bnb.nn.Linear4bit层。

在bnb的实现中，Linear4bit类继承自nn.Linear，并覆盖了其forward方法以使用4位矩阵乘法。MatMul4Bit是一个自定义的autograd函数，用于执行4位矩阵乘法。它首先反量化权重，然后使用浮点线性层进行矩阵乘法。

import bitsandbytes as bnb
from transformers import  BitsAndBytesConfig

# 量化配置
nf4_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,# 双重量化 嵌套量化
)
# 量化模型
model = AutoModelForCausalLM.from_pretrained(
    args.model_name_or_path,
    cache_dir=args.cache_dir,
    load_in_4bit=args.bits == 4,
    load_in_8bit=args.bits == 8,
    device_map=device_map,
    max_memory=max_memory,
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=args.bits == 4,
        load_in_8bit=args.bits == 8,
        llm_int8_threshold=6.0,
        llm_int8_has_fp16_weight=False,
        bnb_4bit_compute_dtype=compute_dtype,
        bnb_4bit_use_double_quant=args.double_quant, # 双重量化 嵌套量化
        bnb_4bit_quant_type=args.quant_type, # nf4
    ),
    ...)
# 模型训练 推理

上面的bnb量化配置会自动将模型中的nn.Linear层替换为`bnb.nn.Linear4bit量化线性层。

def _replace_with_bnb_linear(
    model,...):
    for name, module in model.named_children():
        if isinstance(module, nn.Linear):
            model._modules[name] = bnb.nn.Linear4bit(
                in_features,
                out_features,
                module.bias is not None,
                quantization_config.bnb_4bit_compute_dtype,
                compress_statistics=quantization_config.bnb_4bit_use_double_quant,
                quant_type=quantization_config.bnb_4bit_quant_type,
                **extra_kwargs,
            )

量化参数 Params4bit 类是一个自定义的torch.nn.Parameter子类，它负责执行4位量化算法并存储量化状态。在量化过程中，它使用量化表将浮点数映射到NF4格式的量化级别。在反量化过程中，它使用相同的量化表将量化级别映射回浮点数。

其中量化线性层依赖量化参数 Params4bit 以及 4比特矩阵乘法matmul_4bit

class Linear4bit(nn.Linear):
    def __init__(self,...):
        # 量化参数
        self.weight = Params4bit(self.weight.data, ...)
    def forward(self, x...):
        out = bnb.matmul_4bit(x, self.weight.t(), bias=bias, quant_state=self.weight.quant_state)
        return out

4比特矩阵乘法matmul_4bit依赖MatMul4Bit。

class MatMul4Bit(torch.autograd.Function):
    @staticmethod
    def forward(ctx, A, B, ...):
        # 权重反量化
        dequant_weight = F.dequantize_4bit(B, quant_state, quant_type="nf4").to(A.dtype)
        # 浮点线性层 矩阵乘法
        output = torch.nn.functional.linear(A, dequant_weight, bias)

量化参数 Params4bit 依赖4比特量化算法函数quantize_4bit。

class Params4bit(torch.nn.Parameter):
    def _quantize(self, device):
        w = self.data.contiguous().to(device) # 获取连续权重
        # 执行4比特量化算法
        w_4bit, quant_state = bnb.functional.quantize_4bit(
            w,
            blocksize=self.blocksize, # 分块 量化 块大小
            compress_statistics=self.compress_statistics, # 双重量化
            quant_type=self.quant_type, # nf4 量化类型
        )
        self.data = w_4bit # 0~15 量化后的权重
        self.quant_state = quant_state # 记录量化参数的 量化状态
        return self
    def to(self, *args, **kwargs):
            return self._quantize(device)

4比特量化算法函数 quantize_4bit 实现如下。

def quantize_4bit(...):
    n = A.numel()
    # 4bit量化
    quantize_blockwise_fp32_nf4(code, A, absmax_out, quant_out, blocksize, n)
    # 绝对最大值的双重量化
    if compress_statistics:
        offset = absmax.mean() #  求均值
        absmax -= offset # 减去均值对称分布
        qabsmax, state2 = quantize_blockwise_nf8(absmax, blocksize=256)

上面对于权重的量化为4bit量化，对于分块参数absmax进行8bit量化，两者的量化均为非均匀量化，需要通过量化表进行量化。

# 查表量化
def quantize_by_code(x)
    # 绝对最大值
    a_max = absmax(x) # 1.76
    # 归一化
    x_n = x/a_max # [0.1818, -1, 0.0142, -0.6932]
    # 根据量化表查找索引，可以使用二分查找实现
    # 也可 采用 计算 浮点数和量化表绝对差值最小值的索引
    # 根据 NF4 进行舍入
    x_n_round = round_nf4(x_n, quant_code) # [0.1609, -1.0000, 0.0796, -0.6962] 
    # 输出位于NF4中的索引
    x_n_nf4 = index_nf4(x_n_round)  # [9, 0, 9, 1] 
    return x_n_nf4, a_max
# 查表反量化
def dequantize_by_code(quant_x, a_max):
    # 根据索引取数
    x_n_round = de_index_nf4(x_n_nf4, quant_code) 
    # 乘以最大值 a_max
    dx = x_n_round * a_max
    return dx

量化与反量化的实现 涉及到查找量化表和计算量化级别。在量化过程中，首先计算权重的绝对最大值，并将其用于归一化权重。然后，使用量化表查找最接近的量化级别，并将其存储为NF4格式的索引。在反量化过程中，使用索引从量化表中检索量化级别，并将其乘以绝对最大值以恢复原始的浮点数权重。

利用torch 代码设计上面的量化与反量化过程如下。

import torch

BNB_MAP = [-1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453, -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0, 0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224, 0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0]
MAPPING = torch.tensor(BNB_MAP, dtype=torch.float32).view(1, -1)

def py_quantize_nf4(A, blocksize=64):
    shape = A.shape
    absmax = A.view(-1, blocksize).abs().max(dim=1, keepdim=True).values
    a = A.view(-1, blocksize) / absmax.float()
    diff = torch.abs(a.unsqueeze(-1) - MAPPING)
    out = torch.argmin(diff, dim=-1)
    out = out.reshape(-1, 2)
    out = (out[:, 0] * 16 + out[:, 1]).to(torch.uint8)
    return out, absmax, shape

def py_dequantize_nf4(A, absmax, shape, dtype, blocksize=64):
    A = A.view(-1)
    A = torch.stack([A // 16, A % 16], dim=1).to(torch.int32)
    out = MAPPING.reshape(-1)[A] # 矩阵矩阵索引
    absmax = absmax.to(dtype=torch.float32)
    out = out.view(-1, blocksize) * absmax.reshape(-1, 1) #float类型
    out = out.reshape(*shape).to(dtype)
    return out