推理优化技术
ZeRO推理优化是LLM部署的核心挑战,主要包括量化、剪枝和知识蒸馏三类技术。
2.6.1 量化(Quantization)
量化将模型权重和/或激活从高精度(FP32/FP16)转换为低精度(INT8/INT4/FP8),减少内存占用和计算量。
量化类型
┌─────────────────────────────────────────────────────────────────┐
│ Quantization Types │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 1. 权重量化 (Weight-only Quantization) │
│ - 只量化权重 │
│ - 激活保持FP16 │
│ - 简单,但加速有限 │
│ │
│ 2. 权重-激活量化 (Weight-Activation Quantization) │
│ - 同时量化权重和激活 │
│ - 需要量化感知训练(QAT) │
│ - 加速效果好,但精度损失大 │
│ │
│ 3. KV Cache量化 │
│ - 专门量化KV Cache │
│ - 减少推理内存占用 │
│ - 对长序列场景重要 │
│ │
└─────────────────────────────────────────────────────────────────┘线性量化公式
\(x_{quant} = \text{round}\left(\frac{x}{\text{scale}}\right) + \text{zero\_point}\)
\(x_{dequant} = (x_{quant} - \text{zero\_point}) \times \text{scale}\)
其中:
- scale: 缩放因子
- zero_point: 零点偏移
INT8权重量化(W8A16)
def quantize_weight_int8(weight):
"""
将FP16权重量化到INT8
weight: [out_features, in_features]
"""
# 计算缩放因子(per-channel)
w_max = weight.abs().max(dim=1, keepdim=True).values
scale = w_max / 127.0
# 量化
weight_int8 = torch.round(weight / scale).clamp(-128, 127).to(torch.int)
return weight_int8, scale
def dequantize_weight_int8(weight_int8, scale):
""" 反量化"""
return weight_int8.float() * scaleINT4权重量化(W4A16)
INT4量化将两个INT4值打包到一个字节中
def quantize_weight_int4(weight):
"""
将FP16权重量化到INT4(打包存储)
"""
# 计算缩放因子(per-group,如128个元素一组)
group_size = 128
num_groups = weight.shape[1] // group_size
weight_reshaped = weight.reshape(-1, num_groups, group_size)
w_max = weight_reshaped.abs().max(dim=-1, keepdim=True).values
scale = w_max / 7.0 # INT4 范围: -8 ~ 7
# 量化到INT4
weight_int4 = torch.round(weight_reshaped / scale).clamp(-8, 7).to(torch.int)
# 打包: 两个INT4 → 一个INT8
weight_packed = pack_int4(weight_int4)
return weight_packed, scale量化性能对比
| 量化方案 | 模型大小 | 推理速度 | 精度损失 | 硬件支持 |
|---|---|---|---|---|
| FP16 | 100% | 1.0× | 0% | 通用 |
| W8A16 | 50% | 1.2× | <1% | 通用 |
| W4A16 | 25% | 1.5× | 2-3% | 通用 |
| W4A8 | 25% | 2.0× | 3-5% | 有限 |
| FP8 (W8A8) | 50% | 2.0× | <1% | H100 |
| GPTQ (W4) | 25% | 1.5× | 3-4% | 通用 |
| AWQ (W4) | 25% | 1.5× | 1-2% | 通用 |
GPTQ 与AWQ
:
- 基于 OBS ( Optimal BrainSurgeon)方法
- 逐层量化,最小化输出误差
- 支持任意位宽量化
:
- 考虑激活值幅值进行量化
- 保护重要权重通道
- 精度通常优于GPTQ
2.6.2 剪枝(Pruning)
剪枝通过移除不重要的权重来减小模型大小。剪枝类型
┌─────────────────────────────────────────────────────────────────┐
│ Pruning Types │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 1. 非结构化剪枝 (Unstructured Pruning) │
│ - 移除单个权重 │
│ - 高稀疏度,但难加速 │
│ - 需要特殊硬件/库支持 │
│ │
│ 2. 结构化剪枝 (Structured Pruning) │
│ - 移除整个神经元/通道 │
│ - 易加速,但稀疏度受限 │
│ - 通用硬件支持 │
│ │
│ 3. 半结构化剪枝 (Semi-structured) │
│ - 2:4稀疏模式 │
│ - Ampere GPU原生支持 │
│ - 平衡稀疏度和加速 │
│ │
└─────────────────────────────────────────────────────────────────┘2:4 结构化稀疏
Ampere架构GPU支持2:4稀疏模式(每4个权重保留2个):
原始权重: 稀疏化后: 压缩存储:
[0.1, 0.9, [0.0, 0.9, [0.9, 0.8, 0.7, 0.5]
0.8, 0.2, 0.8, 0.0, 元数据: [1, 0, 0, 1, 0, 1, 1, 0]
0.3, 0.7, 0.0, 0.7,
0.5, 0.4] 0.5, 0.0]
50% 稀疏度,2×加速剪枝效果
| 稀疏度 | 模型大小 | 理论加速 | 实际加速 | 精度损失 |
|---|---|---|---|---|
| 50% (2:4) | 50% | 2× | 1.5-1.8× | 1-3% |
| 75% | 25% | 4× | 2-3× | 5-10% |
| 90% | 10% | 10× | 4-6× | 15-30% |
2.6.3 知识蒸馏(Knowledge Distillation)
知识蒸馏通过训练小模型(学生)模仿大模型(教师)的行为来压缩模型。
蒸馏类型
┌─────────────────────────────────────────────────────────────────┐
│ Knowledge Distillation Types │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 1. 输出蒸馏 (Logit Distillation) │
│ - 学生模仿教师的softmax输出 │
│ - 使用温度参数软化概率分布 │
│ - 最简单直接 │
│ │
│ 2. 中间层蒸馏 (Hidden State Distillation) │
│ - 学生模仿教师的中间层表示 │
│ - 需要设计映射层 │
│ - 效果更好 │
│ │
│ 3. 注意力蒸馏 (Attention Distillation) │
│ - 学生模仿教师的注意力矩阵 │
│ - 捕捉结构信息 │
│ - 适合Transformer │
│ │
└─────────────────────────────────────────────────────────────────┘蒸馏损失函数
软目标损失(KL散度):
\(\mathcal{L}_{\text{soft}} = T^2 \cdot \text{KL}\left(\text{softmax}\left(\frac{z_T}{T}\right) |\text{softmax}\left(\frac{z_S}{T}\right)\right)\)
其中 \(T\) 是温度参数,通常 \(T = 2\) 到 \(5\)。
硬目标损失(交叉熵):
\(\mathcal{L}_{\text{hard}} = \text{CE}(y, \text{softmax}(z_S))\)
\(CE(P,Q) = H(P) + D_{KL}(P \parallel Q)\)
总损失:
\(\mathcal{L} = \alpha \cdot \mathcal{L}_{\text{soft}} + (1-\alpha) \cdot \mathcal{L}_{\text{hard}}\)
蒸馏效果示例
| 教师模型 | 学生模型 | 参数量比 | 精度保持 |
|---|---|---|---|
| GPT-3 175B | GPT-3 6.7B | 26× | 85-90% |
| LLaMA-2 70B | LLaMA-2 7B | 10× | 80-85% |
| LLaMA-2 13B | LLaMA-2 1.1B | 12× | 70-75% |