Skip to content

推理优化技术

ZeRO推理优化是LLM部署的核心挑战,主要包括量化、剪枝和知识蒸馏三类技术。

2.6.1 量化(Quantization)

量化将模型权重和/或激活从高精度(FP32/FP16)转换为低精度(INT8/INT4/FP8),减少内存占用和计算量。

量化类型

latex
   ┌─────────────────────────────────────────────────────────────────┐
  │                           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)

python
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() * scale

INT4权重量化(W4A16)

INT4量化将两个INT4值打包到一个字节中

python
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

量化性能对比

量化方案模型大小推理速度精度损失硬件支持
FP16100%1.0×0%通用
W8A1650%1.2×<1%通用
W4A1625%1.5×2-3%通用
W4A825%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)

剪枝通过移除不重要的权重来减小模型大小。剪枝类型

latex
  ┌─────────────────────────────────────────────────────────────────┐
  │                        Pruning Types                            │
  ├─────────────────────────────────────────────────────────────────┤
  │                                                                               │
  │     1.   非结构化剪枝 (Unstructured Pruning)                                 │
  │          - 移除单个权重                                                     │
  │          - 高稀疏度,但难加速                                              │
  │          - 需要特殊硬件/库支持                                             │
  │                                                                               │
  │     2.   结构化剪枝 (Structured Pruning)                                       │
  │          - 移除整个神经元/通道                                             │
  │          - 易加速,但稀疏度受限                                             │
  │          - 通用硬件支持                                                     │
  │                                                                               │
  │     3.   半结构化剪枝 (Semi-structured)                                      │
  │          - 2:4稀疏模式                                                    │
  │          - Ampere GPU原生支持                                             │
  │          - 平衡稀疏度和加速                                               │
  │                                                                               │
  └─────────────────────────────────────────────────────────────────┘

2:4 结构化稀疏

Ampere架构GPU支持2:4稀疏模式(每4个权重保留2个):

latex
  原始权重:       稀疏化后:         压缩存储:
  [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%1.5-1.8×1-3%
75%25%2-3×5-10%
90%10%10×4-6×15-30%

2.6.3 知识蒸馏(Knowledge Distillation)

知识蒸馏通过训练小模型(学生)模仿大模型(教师)的行为来压缩模型。

蒸馏类型

latex
  ┌─────────────────────────────────────────────────────────────────┐
  │                     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 175BGPT-3 6.7B26×85-90%
LLaMA-2 70BLLaMA-2 7B10×80-85%
LLaMA-2 13BLLaMA-2 1.1B12×70-75%

2.6.4 推理优化技术选择指南

用心记录,持续成长