簡介

LORA (Low-Rank Adaptation) 是一種高效微調大型預訓練模型的方法。它通過凍結預訓練模型的權重，并在Transformer架構的每一層中引入可訓練的秩分解矩陣，顯著減少了可訓練參數的數量，從而確保了更加高效的適應過程。具體來說，它將一個大矩陣分解為兩個低秩矩陣的乘積，即 weight[ho] = w1[hr] @ w2[ro]，其中 r 是秩，是一個關鍵的超參數。通常，r 的值設置為4、8或12，以平衡表(biao)達(da)力(li)和計算效率。

QLoRA 是LoRA的量化版本，它結合了量化技術來進一步減少內存和計算成本。在QLoRA中，LoRA的可訓練低秩矩陣 w1 和 w2 保持不量化，以便進行反向傳播和優化。然而，原始模型的權重 W 被凍結并(bing)量(liang)化，以減少內存占用。

low-rank低秩矩陣分解(jie)，只對(dui)模型(xing)的(de)部分權重進行(xing)更新，而(er)不是(shi)全量微調(diao)，可大幅較(jiao)少(shao)參數更新量，減(jian)少(shao)計算量。

兩個(ge)(ge)低秩矩(ju)陣替代一(yi)個(ge)(ge)大(da)矩(ju)陣，將 權(quan)(quan)重(zhong) weight 分(fen)解(jie)為(wei) 兩個(ge)(ge) 小張量 weight[ho] = w1[hr]@w2[ro], 權(quan)(quan)重(zhong)更(geng)新(xin)時對(dui) 兩個(ge)(ge)低秩 矩(ju)陣w1/w2(適(shi)配(pei)器權(quan)(quan)重(zhong))進行(xing)更(geng)新(xin)。其中 r稱為(wei) rank，是一(yi)個(ge)(ge)重(zhong)要的超(chao)參數。r越(yue)大(da)，可訓(xun)練參數越(yue)多，表達力(li)越(yue)強(qiang)，不過一(yi)般也(ye)不需要太大(da)， r = 4,8,12 是常見的設置。

此外，在 Transformer block 中，對哪些參數矩陣(zhen)(zhen)的更新量進行(xing)模擬？一般是選擇自注意力(li)層的 Q, V 矩陣(zhen)(zhen)；或者范圍更大(da)一些：Q, K, V, O（輸出(chu)矩陣(zhen)(zhen)）。

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

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

LoRA的(de)參數 w1 w2 是不量化的(de)，因為它(ta)們需要反向(xiang)傳(chuan)播優(you)化。而(er)原始模型的(de)參數 W 是freeze的(de)，因此可以量化。

量化技術 在QLoRA中扮演(yan)了(le)(le)至關重要的(de)角色。特(te)別是，QLoRA采用(yong)了(le)(le)非均勻量化(hua)（如NF4）來量化(hua)權重，這有助于減少量化(hua)誤(wu)差并(bing)保留更(geng)多的(de)信息。非均勻量化(hua)使用(yong)量化(hua)表將浮點數映射到離散的(de)量化(hua)級別，從而實現了(le)(le)更(geng)高的(de)壓縮(suo)率和(he)更(geng)低(di)的(de)量化(hua)誤(wu)差。

在(zai)QLoRA中，模(mo)型的權(quan)重以NF4格(ge)(ge)(ge)式(shi)存(cun)儲(chu)，但在(zai)計(ji)算(suan)時使用BF16格(ge)(ge)(ge)式(shi)。這意味著在(zai)每次前向傳播(bo)之前，需要將(jiang)NF4格(ge)(ge)(ge)式(shi)的權(quan)重反(fan)(fan)量化(hua)為(wei)BF16格(ge)(ge)(ge)式(shi)進(jin)行計(ji)算(suan)。計(ji)算(suan)完成后，再將(jiang)結(jie)果量化(hua)為(wei)NF4格(ge)(ge)(ge)式(shi)進(jin)行存(cun)儲(chu)。這種(zhong)雙重反(fan)(fan)量化(hua)過程進(jin)一步節省(sheng)了量化(hua)常數的空間占用。

針對離群值(zhi)/異常值(zhi)：分(fen)塊量化，不同數據塊，有不同的量化系(xi)數 c，如 分(fen)N塊單獨量化，需要N個(ge)量化系(xi)數c2。

雙(shuang)重反(fan)量(liang)化(hua)：w存儲(chu)(chu)為 nf4類(lei)型(xing)，通過兩次反(fan)量(liang)化(hua) dequant(dequant(c1,c2), W_nf4)，將存儲(chu)(chu)數(shu)據(ju)類(lei)型(xing)轉換為計算數(shu)據(ju)類(lei)型(xing)， 進一步(bu)節(jie)省量(liang)化(hua)常數(shu)的(de)空間(jian)占用；

c2為(wei)分塊量(liang)化(hua)(hua)(hua)(hua)系(xi)數(shu)，每塊一個量(liang)化(hua)(hua)(hua)(hua)系(xi)數(shu)，為(wei)降(jiang)低量(liang)化(hua)(hua)(hua)(hua)系(xi)數(shu)存儲(chu)，利用(yong)c1將量(liang)化(hua)(hua)(hua)(hua)的量(liang)化(hua)(hua)(hua)(hua)系(xi)數(shu)c2反量(liang)化(hua)(hua)(hua)(hua)為(wei)浮(fu)點，再(zai)將量(liang)化(hua)(hua)(hua)(hua)的權(quan)重量(liang)化(hua)(hua)(hua)(hua)為(wei)浮(fu)點。

QLoRA 中(zhong)，模型的權(quan)重有(you)兩種(zhong)格式：用(yong) NF4 存儲；用(yong) BF16 計算。需(xu)要(yao)用(yong)相應權(quan)重計算前向傳播時，對 NF4 的權(quan)重反(fan)量化為(wei) BF16；計算結(jie)束后，再量化為(wei) NF4。

QLoRA 步驟：

• 初始量(liang)化：首先，大型(xing)語言模型(xing) (LLM) 權重 被量(liang)化為(wei) 4 位，顯著減少(shao)了內存占用。

  • 量化權重
  • 量化 量化系數

• LoRA微調(diao)：然后，執(zhi)行 LoRA 訓練。

  • 反量化 量化系數
  • 反量化 權重
  • 計算Y 梯度更新 前向

bitsandbytes 實現原理

BitsandBytes (bnb) 是一個用于實現QLoRA的庫。它提供了高效的4位和8位量化算法，以及相關的量化配置和模型替換功能。通過配置bnb的量化參數，可以輕松地將模型中的nn.Linear層替換為量化的bnb.nn.Linear4bit層。

在bnb的實現中，Linear4bit類繼承自nn.Linear，并覆蓋了其forward方法以使用4位矩陣乘法。MatMul4Bit是一個自定義的autograd函(han)數，用于執行4位矩(ju)陣(zhen)乘法(fa)。它(ta)首(shou)先(xian)反(fan)量化權重，然后(hou)使用浮(fu)點線(xian)性層進(jin)行矩(ju)陣(zhen)乘法(fa)。

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層(ceng)替換(huan)為`bnb.nn.Linear4bit量化線性(xing)層(ceng)。

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子類，它(ta)(ta)負(fu)責執(zhi)行(xing)4位量(liang)化算法并(bing)存儲量(liang)化狀態。在量(liang)化過程(cheng)中，它(ta)(ta)使(shi)用量(liang)化表(biao)將(jiang)浮點數映(ying)射到NF4格(ge)式的(de)量(liang)化級別(bie)。在反量(liang)化過程(cheng)中，它(ta)(ta)使(shi)用相同的(de)量(liang)化表(biao)將(jiang)量(liang)化級別(bie)映(ying)射回(hui)浮點數。

其中量化線性層依賴量化參數 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比(bi)特量化算(suan)法函數 quantize_4bit 實現(xian)如下。

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)

上面(mian)對(dui)于(yu)權重的量(liang)化(hua)(hua)為(wei)4bit量(liang)化(hua)(hua)，對(dui)于(yu)分塊(kuai)參(can)數absmax進行(xing)8bit量(liang)化(hua)(hua)，兩(liang)者的量(liang)化(hua)(hua)均為(wei)非均勻(yun)量(liang)化(hua)(hua)，需要通過量(liang)化(hua)(hua)表進行(xing)量(liang)化(hua)(hua)。

# 查表量化
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

量化與反量化的實現 涉及到查找量化(hua)(hua)(hua)(hua)(hua)表和計算(suan)量化(hua)(hua)(hua)(hua)(hua)級(ji)別。在量化(hua)(hua)(hua)(hua)(hua)過程中，首先計算(suan)權重(zhong)的絕對最大值，并(bing)(bing)將(jiang)其(qi)用于歸(gui)一化(hua)(hua)(hua)(hua)(hua)權重(zhong)。然后，使用量化(hua)(hua)(hua)(hua)(hua)表查找最接(jie)近(jin)的量化(hua)(hua)(hua)(hua)(hua)級(ji)別，并(bing)(bing)將(jiang)其(qi)存儲(chu)為(wei)NF4格(ge)式的索引。在反量化(hua)(hua)(hua)(hua)(hua)過程中，使用索引從量化(hua)(hua)(hua)(hua)(hua)表中檢索量化(hua)(hua)(hua)(hua)(hua)級(ji)別，并(bing)(bing)將(jiang)其(qi)乘以絕對最大值以恢(hui)復原始的浮點數權重(zhong)。

利用torch 代碼設(she)計上(shang)面的量(liang)化與反量(liang)化過程(cheng)如下。

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