LoRA的显存成本 QLoRA: Efficient Finetuning of Quantized LLMs

LoRA 冻结预训练权重，只训练 low-rank update，因此显著减少了可训练参数、gradient 和 optimizer state 的数量。但“参数高效”并不等于“显存占用已经足够低”：即使 $\underline{\underline{W_0}}$ 不参与更新，完整的基础模型仍然要放进显存，前向和反向计算也仍然需要经过它。

QLoRA 接着追问：

能否把冻结的基础模型压缩到 4-bit，同时仍然通过它反向传播，将 gradient 传给 LoRA adapter？

QLoRA 将 $\underline{\underline{W_0}}$ 以 4-bit 保存，使用时临时反量化到 BFloat16，再完成 $\underline{h}=\operatorname{dequant}(\underline{\underline{W_0^{\mathrm{NF4}}}})\underline{x}+\frac{\alpha}{r}\underline{\underline{B}}\underline{\underline{A}}\underline{x}$。基础权重保持冻结，真正被 optimizer 更新的仍然只有 $\underline{\underline{A}}$ 和 $\underline{\underline{B}}$。

QLoRA: Efficient Finetuning of Quantized LLMs

Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, Luke Zettlemoyer

NeurIPS · 2023

arXiv PDF Code 复制引用

论文将自己的三个主要技术贡献概括为：4-bit NormalFloat（NF4）、Double Quantization 和 Paged Optimizers。

贡献一：4-bit NormalFloat（NF4）#

量化要用有限的 bit pattern 表示连续数值。4-bit 只能提供 $2^4=16$ 个离散值，因此关键问题不是“能否只用 4 bit”，而是这 16 个值应该放在哪里。

预训练模型的权重通常近似服从以 $0$ 为中心的正态分布：大部分权重集中在零附近，绝对值较大的权重较少。如果仍然均匀放置 16 个量化值，分布稀疏的两端会占用过多编码，零附近却没有足够精细的表示。

NF4 根据标准正态分布的 quantile 选择 16 个候选值，使每个 quantization bin 在理论上包含相近数量的权重。于是，零附近的候选值更密集，尾部更稀疏，更符合权重本身的分布。

具体使用时，QLoRA 先把每个 block 中的权重除以该 block 的绝对值最大值，将它们缩放到 $[-1,1]$，再映射到最近的 NF4 候选值。存储的是每个权重的 4-bit 编号，以及恢复这个 block 所需的 quantization constant。

NF4 还需要精确表示 $0$。论文分别为负半轴选择 $8$ 个 quantile 值、为正半轴选择 $9$ 个 quantile 值，再删除两个集合中重复的一个 $0$，最终得到 $8+9-1=16$ 个候选值。

这里需要区分已有方法和论文贡献：block-wise quantization 并不是 QLoRA 首次提出的；QLoRA 的新贡献是为近似正态分布的预训练权重设计 NF4，并将它用于高精度的 4-bit fine-tuning。

贡献二：Double Quantization#

4-bit 编号并不足以恢复权重。每个 block 还需要一个 quantization constant，记录这个 block 在量化前被缩放了多少。block 越小，局部量化通常越准确，但需要保存的 constant 也越多。

QLoRA 的第一层量化每 $64$ 个权重使用一个 FP32 constant，因此 constant 的平均成本是 $32/64=0.5$ bit per parameter。Double Quantization 的做法，是把这些 constant 当作第二组待量化数据，再用 8-bit float 和 block size $256$ 对它们进行量化。

第二层量化后的平均成本变为 $8/64+32/(64\times256)=0.127$ bit per parameter，因此比直接保存 FP32 constant 少 $0.373$ bit per parameter。

对于一个 $30$B 参数模型，理论节省量约为 $30\times10^9\times0.373/8\approx1.40$ GB。它节省的不是 4-bit 权重本身，而是每个 block 附带的 quantization constant。

由于第一层 constant 都是正数，QLoRA 在第二次量化前先减去它们的平均值，使数据分布移动到零的两侧，从而更充分地使用对称量化的正负编码；反量化时再把平均值加回来。

贡献三：Paged Optimizers#

NF4 和 Double Quantization 减少的是训练期间长期存在的显存占用，但训练显存并不始终保持不变。处理较长序列或进行 gradient checkpointing 的重新计算时，显存可能短暂出现峰值。即使训练大部分时间只使用 $22$ GB，某一步突然需要 $25$ GB，在一张 $24$ GB GPU 上仍然会发生 out-of-memory。

Paged Optimizers 主要解决这种短时峰值，而不是进一步量化基础模型。它利用 NVIDIA Unified Memory，在 GPU 显存紧张时将部分 optimizer state page 移到 CPU RAM；需要执行参数更新时，再把相关 page 传回 GPU。

以 Adam 为例，optimizer 通常需要为可训练参数保存 gradient 的一阶 moving average $m_t$ 和平方 gradient 的 moving average $v_t$。QLoRA 只需要为 LoRA adapter 保存这些状态，不需要为冻结的 4-bit 基础权重保存；但当其他计算临时占用更多显存时，这些状态仍然可以作为可迁移的缓冲空间。

因此，三项贡献解决的是三个不同位置的显存问题：

1. NF4 用更适合权重分布的 4-bit 候选值压缩冻结的基础模型；
2. Double Quantization 继续压缩第一层量化产生的 constant；
3. Paged Optimizers 在显存峰值出现时，将 optimizer state 暂时迁移到 CPU。

QLoRA 没有改变 LoRA 对 weight update 的 low-rank 假设。它改变的是基础模型的保存与计算方式，使 LoRA 不只减少需要训练的参数，也能用于单张 GPU 无法以 16-bit 容纳的更大模型。