276 lines
7.8 KiB
Markdown
276 lines
7.8 KiB
Markdown
# QLoRA（Quantized Low-Rank Adaptation）
|
||
|
||
## 概述
|
||
|
||
QLoRA（Quantized Low-Rank Adaptation）是由Tim Dettmers等人于2023年提出的高效微调方法，结合了量化技术与低秩适配（LoRA），使得在单张48GB GPU上微调65B参数模型成为可能。
|
||
|
||
**核心思想**：将预训练模型量化到4-bit NF4存储，但在计算时反量化为BF16进行训练，同时使用LoRA适配器学习任务相关知识。
|
||
|
||
## NF4数据类型（4-bit NormalFloat）
|
||
|
||
### 问题背景
|
||
|
||
标准INT4量化对于正态分布的权重不是最优的，因为正态分布的权重理论上应该均匀分布在量化网格上。
|
||
|
||
### NF4原理
|
||
|
||
NF4（4-bit NormalFloat）是一种针对正态分布数据优化的4位数值格式：
|
||
|
||
```
|
||
NF4量化网格: [-1, -0.696, -0.525, -0.351, -0.175, 0, 0.175, 0.351, 0.525, 0.696, 1]
|
||
(共11个值点，用于量化16个4-bit表示)
|
||
```
|
||
|
||
**设计原则**：
|
||
- 量化边界基于正态分布的分位数确定
|
||
- 保证每个量化值承载的信息量相等
|
||
- 减少小值区域的量化误差
|
||
|
||
### 与INT4对比
|
||
|
||
```
|
||
INT4: 均匀分布 [-7, -5, -3, -1, 1, 3, 5, 7]（对称8值）
|
||
NF4: 非均匀分布，基于正态分布的分位数
|
||
```
|
||
|
||
### PyTorch实现
|
||
|
||
```python
|
||
# NF4 量化网格（预计算）
|
||
NF4_GRID = [-1.0, -0.69619224, -0.52507305, -0.35093744,
|
||
-0.17546872, 0.0, 0.17546872, 0.35093744,
|
||
0.52507305, 0.69619224, 1.0]
|
||
|
||
def quantize_nf4(tensor):
|
||
# 计算每个值最近的NF4量化点
|
||
scale = tensor.abs().max()
|
||
tensor_normalized = tensor / scale
|
||
quantized = torch.zeros_like(tensor, dtype=torch.uint8)
|
||
for i, grid_val in enumerate(NF4_GRID):
|
||
quantized[tensor_normalized == grid_val] = i
|
||
return quantized, scale
|
||
|
||
def dequantize_nf4(quantized, scale):
|
||
# 反量化，使用查找表
|
||
values = torch.tensor(NF4_GRID, device=quantized.device)
|
||
tensor_normalized = values[quantized]
|
||
return tensor_normalized * scale
|
||
```
|
||
|
||
## 双重量化（Double Quantization）
|
||
|
||
### 问题
|
||
|
||
在QLoRA中，量化权重需要存储量化参数（scale、zero_point）。对于NF4，每个参数块（block）需要一个FP32 scale。
|
||
|
||
```
|
||
Block size = 64 参数 → 1个FP32 scale
|
||
模型参数量：65B → 约1B个Block → 1B × 4 bytes = 4GB（仅scale）
|
||
```
|
||
|
||
### 解决方案：双重量化
|
||
|
||
对量化参数本身再进行量化：
|
||
|
||
```python
|
||
# 第一层：对模型权重进行NF4量化
|
||
W_quant, W_scale = quantize_nf4(W) # W_scale: FP32, block_size=64
|
||
|
||
# 第二层：对scale进行INT8量化
|
||
W_scale_quant, W_scale_scale = quantize_int8(W_scale)
|
||
# W_scale_scale: FP32, block_size=256（更大的block，因为scale分布更集中）
|
||
```
|
||
|
||
### 显存节省
|
||
|
||
| 量化方式 | scale显存 |
|
||
|---------|----------|
|
||
| 无双重量化（FP32） | 4GB（65B参数模型） |
|
||
| 双重量化（INT8 scale） | ~0.5GB |
|
||
| 节省 | ~3.5GB |
|
||
|
||
## 梯度分页优化器（Paged Optimizer）
|
||
|
||
### 问题
|
||
|
||
在多GPU训练时，梯度累积需要存储大量优化器状态。当显存不足时，通常需要CPU卸载，但CPU-GPU数据传输成为瓶颈。
|
||
|
||
### 分页优化器原理
|
||
|
||
借鉴操作系统虚拟内存的页面调度思想：
|
||
|
||
```python
|
||
class PagedOptimizer:
|
||
"""梯度分页优化器"""
|
||
|
||
def __init__(self, model_params, optimizer_class, block_size=1024):
|
||
self.blocks = {} # device_id -> list of blocks
|
||
self.block_size = block_size
|
||
|
||
def step(self, gradients):
|
||
# 检查是否需要页面调度
|
||
for param in gradients:
|
||
if self._need_offload(param):
|
||
# 将不活跃的优化器状态换出到CPU
|
||
self._page_out(self._find_inactive_block())
|
||
|
||
# 正常优化器步骤
|
||
self.optimizer.step()
|
||
|
||
def _need_offload(self, param):
|
||
# 判断是否需要换页
|
||
return self._gpu_memory_pressure() > self._threshold
|
||
```
|
||
|
||
### 在QLoRA中的应用
|
||
|
||
```python
|
||
# QLoRA训练配置
|
||
config = {
|
||
"lora_r": 64,
|
||
"lora_alpha": 16,
|
||
"lora_dropout": 0.05,
|
||
"weight_loader": {
|
||
"nf4": True,
|
||
"double_quant": True,
|
||
"quantization_storage_type": "nf4"
|
||
},
|
||
"paged_optimizer": {
|
||
"enabled": True,
|
||
"block_size": 1024,
|
||
"cpu_offload": True
|
||
}
|
||
}
|
||
```
|
||
|
||
## LoRA与量化的结合
|
||
|
||
### 架构设计
|
||
|
||
```
|
||
输入X
|
||
↓
|
||
┌───────────────────────┐
|
||
│ 冻结的4-bit权重 W │ ← NF4量化存储，反量化后用于计算
|
||
└───────┬───────────────┘
|
||
↓
|
||
┌───────────────────────┐
|
||
│ 添加LoRA适配器 │ ← BF16训练
|
||
│ ΔW = A × B │
|
||
│ A: (in_features, r) │
|
||
│ B: (r, out_features)│
|
||
└───────────────────────┘
|
||
↓
|
||
输出Y = X @ (W + ΔW)
|
||
```
|
||
|
||
### 量化权重前向传播
|
||
|
||
```python
|
||
def quantized_linear(x, weight, lora_a, lora_b, scale):
|
||
"""
|
||
x: 输入 (batch, seq, in_features)
|
||
weight: NF4量化权重 (存储用)
|
||
lora_a, lora_b: LoRA参数 (BF16)
|
||
"""
|
||
# 反量化权重到BF16进行计算
|
||
w_bf16 = dequantize_nf4(weight) # (out_features, in_features)
|
||
|
||
# 计算量化部分贡献
|
||
y_quant = F.linear(x, w_bf16)
|
||
|
||
# 计算LoRA部分贡献（BF16）
|
||
y_lora = F.linear(F.linear(x, lora_a), lora_b)
|
||
|
||
return y_quant + y_lora
|
||
```
|
||
|
||
### 训练流程
|
||
|
||
```python
|
||
# QLoRA训练循环
|
||
for batch in dataloader:
|
||
optimizer.zero_grad()
|
||
|
||
# 1. 前向传播：反量化权重，计算损失
|
||
x = batch["input_ids"]
|
||
outputs = model(x, use_qlora=True)
|
||
loss = criterion(outputs, batch["labels"])
|
||
|
||
# 2. 反向传播：只更新LoRA参数和LayerNorm
|
||
loss.backward()
|
||
|
||
# 梯度裁剪
|
||
torch.nn.utils.clip_grad_norm_(model.lora_params(), max_norm=0.5)
|
||
|
||
# 3. 优化器步骤（只在GPU上更新LoRA参数）
|
||
optimizer.step()
|
||
|
||
# 4. 页面调度（如需要）
|
||
paged_optimizer.check_memory_pressure()
|
||
```
|
||
|
||
### 可训练参数占比
|
||
|
||
| 模型规模 | 总参数量 | LoRA参数量（r=64） | 训练比例 |
|
||
|---------|---------|-------------------|---------|
|
||
| 7B | 7,000M | ~50M（7B × 64 × 2） | 0.7% |
|
||
| 13B | 13,000M | ~50M | 0.4% |
|
||
| 65B | 65,000M | ~50M | 0.08% |
|
||
|
||
极低训练参数量使得单卡微调超大模型成为可能。
|
||
|
||
## 量化参数配置
|
||
|
||
### NF4参数块大小
|
||
|
||
```python
|
||
# 量化参数
|
||
BITS = 4
|
||
BLOCK_SIZE = 64 # 每64个参数共用一个scale
|
||
|
||
# 双重量化参数
|
||
DQ_BLOCK_SIZE = 256 # 每256个scale共用一个量化scale
|
||
```
|
||
|
||
### 量化存储类型
|
||
|
||
| 配置 | 存储精度 | 计算精度 | 适用场景 |
|
||
|------|---------|---------|---------|
|
||
| NF4 + BF16 | NF4 | BF16 | 标准QLoRA |
|
||
| INT8 + BF16 | INT8 | BF16 | 精度要求高 |
|
||
| NF4 + FP16 | NF4 | FP16 | 硬件限制 |
|
||
|
||
## 实验结果
|
||
|
||
| 模型 | 量化精度 | 任务表现 | 单GPU显存 |
|
||
|------|---------|---------|---------|
|
||
| LLaMA-65B | QLoRA (NF4) | 接近FP16基线 | 48GB |
|
||
| LLaMA-33B | QLoRA (NF4) | 接近FP16基线 | 24GB |
|
||
| LLaMA-7B | QLoRA (NF4) | 超过FP16 | 6GB |
|
||
|
||
## 实现框架
|
||
|
||
```python
|
||
# transformers + peft
|
||
from peft import LoraConfig, get_peft_model, prepare_model_for_kint4_training
|
||
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
"model_name",
|
||
load_in_4bit=True,
|
||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||
bnb_4bit_use_double_quant=True,
|
||
bnb_4bit_quant_type="nf4"
|
||
)
|
||
|
||
model = prepare_model_for_kint4_training(model)
|
||
config = LoraConfig(r=64, lora_alpha=16, target_modules=["q_proj", "v_proj"])
|
||
model = get_peft_model(model, config)
|
||
```
|
||
|
||
## 注意事项
|
||
|
||
1. **量化校准**：需要代表性数据集估计量化范围
|
||
2. **量化误差**：累积误差可能导致训练不稳定
|
||
3. **LoRA秩选择**：过大r可能导致过拟合，过小r表达能力不足
|
||
4. **混合精度**：部分层（如Embedding）保持FP16以保证精度 |