Skip to content

压缩算法

大模型推理面临严峻的内存挑战。以LLaMA-3-70B为例,FP16权重需要约140GB内存,加上KV Cache后很容易超出单卡容量。量化技术通过降低数值精度来压缩模型和激活,是解决这一问题的关键技术。

7.4.1 KV Cache 量化

问题背景

KV Cache 是推理阶段的主要内存消耗来源:

\(\text{Memory}_{\text{KV}} = 2 \times n_{\text{layers}} \times n_{\text{heads}} \times d_{\text{head}} \times n_{\text{tokens}} \times \text{bytes}_{\text{dtype}}\)

对于长序列, KV Cache 可能超过模型权重本身:

  • LLaMA-3-70B,序列长度 32K,FP16:约84GB KV Cache vs 140GB权重

7.4.1.1 INT8量化

对称量化

对称INT8量化将FP16值映射到[-128, 127]范围:

\(x_{\text{int8}} = \text{round}\left(\frac{x_{\text{fp16}}}{s}\right)\)

\(x_{\text{dequant}} = x_{\text{int8}} \times s\)

其中缩放因子 \(s\) 计算为:

\(s = \frac{\max(|x|)}{127}\)

非对称量化

非对称量化使用零点(zero-point):

\(x_{\text{int8}} = \text{round}\left(\frac{x_{\text{fp16}} - z}{s}\right)\)

\(x_{\text{dequant}} = x_{\text{int8}} \times s + z\)

其中: \(s = \frac{\max(x) - \min(x)}{255}\)\(z = \text{round}\left(\frac{-\min(x)}{s}\right)\)

逐通道量化

KV Cache 通常按通道(channel-wise)量化以获得更好的精度:

python
def quantize_kv_cache_int8(K_cache, V_cache, axis=-1):
       """
       对KV Cache进行INT8量化
       K_cache, V_cache: [batch, heads, seq_len, head_dim]
       """
       #   计算每个通道的缩放因子
       K_max = K_cache.abs().amax(dim=axis, keepdim=True)
       V_max = V_cache.abs().amax(dim=axis, keepdim=True)

       K_scale = K_max / 127.0
       V_scale = V_max / 127.0

       #   量化
       K_int8 = torch.round(K_cache / K_scale).clamp(-128, 127).to(torch.int8)
       V_int8 = torch.round(V_cache / V_scale).clamp(-128, 127).to(torch.int8)


       return K_int8, V_int8, K_scale, V_scale


  def dequantize_kv_cache_int8(K_int8, V_int8, K_scale, V_scale):
       """ 反量化KV Cache"""
       K_fp16 = K_int8.to(torch.float16) * K_scale
       V_fp16 = V_int8.to(torch.float16) * V_scale
       return K_fp16, V_fp16

精度分析

模型任务FP16 PPLINT8 PPL相对增长
LLaMA-2-7BWikiText5.125.181.2%
LLaMA-2-70BWikiText3.323.350.9%
LLaMA-3-8BHumanEval28.0%27.5%-1.8%

结论:INT8 KV Cache量化通常带来<2%的精度损失,可接受。

7.4.1.2 INT4量化

分组量化

INT4的范围仅为[-8, 7],需要更精细的量化策略。分组量化(group-wise quantization)

将通道分成小组:

python
def quantize_kv_cache_int4(K_cache, group_size=128):
          """
          INT4  分组量化
          """
          batch, heads, seq_len, head_dim = K_cache.shape
          num_groups = head_dim // group_size

          # reshape  为 [..., num_groups, group_size]
          K_reshaped = K_cache.reshape(batch, heads, seq_len, num_groups, group_s

          #   每组计算缩放因子
          K_max = K_reshaped.abs().amax(dim=-1, keepdim=True)
          K_scale = K_max / 7.0


          #   量化到INT4(实际存储为INT8,但只使用4bit)
          K_int4 = torch.round(K_reshaped / K_scale).clamp(-8, 7)


          #   打包存储(两个INT4打包到一个INT8)
          K_packed = pack_int4(K_int4)


          return K_packed, K_scale

精度保持技术

INT4量化需要额外的精度保持技术:

  1. 异常值处理:识别并特殊处理离群值
  2. 动态缩放:根据激活分布动态调整缩放因子
  3. 混合精度:关键层使用INT8,其他层使用INT4

精度分析

量化方案压缩比典型精度损失
INT82x<2%
INT44x3-8%
INT4 + 异常值处理3.5x2-4%

7.4.1.3 FP8 量化

E4M3 和 E5M2 格式

FP8有两种主要格式:

  • E4M3:1位符号,4位指数,3位尾数,范围±448
  • E5M2:1位符号,5位指数,2位尾数,范围±57344

E4M3: S EEEE MMM (动态范围小,精度高)

E5M2: S EEEEE MM (动态范围大,精度低)

动态缩放

FP8 量化通常使用动态缩放因子:

\(s = \frac{\text{FP8\_max}}{\max(|x|) + \epsilon}\)

\(x_{\text{fp8}} = x \times s\)

实现(H100)

python
def quantize_fp8_e4m3(tensor):
      """使用E4M3格式量化到FP8"""
      # H100支持原生FP8张量核心
      fp8_max = 448.0

      #  计算缩放因子
      amax = tensor.abs().max()
      scale = fp8_max / amax

      #  量化(使用H100的FP8指令)
      tensor_fp8 = torch._scaled_quantize(tensor, scale, "E4M3")

      return tensor_fp8, scale

FP8 vs INT8

特性INT8FP8
动态范围固定可调
硬件支持广泛Hopper+
精度中高
速度更快(张量核心)

7.4.1.4 精度保持技术

离群值感知量化

KV Cache 中存在少量离群值(outliers),严重影响量化精度:

python
def outlier_aware_quantize(tensor, outlier_threshold=6.0):
    """
    离群值感知量化
    """
    #   识别离群值
    mean = tensor.mean()
    std = tensor.std()
    outliers = (tensor - mean).abs() > outlier_threshold * std

    #   离群值使用更高精度
    normal_mask = ~outliers

    #   量化非离群值
    normal_quantized = quantize(tensor[normal_mask])

    #   离群值保持FP16
    outliers_fp16 = tensor[outliers]

    return normal_quantized, outliers_fp16, normal_mask

混合精度策略

不同层使用不同精度:

  • 浅层(输入层):FP16(保留更多信息)
  • 中层:INT8
  • 深层:INT4(对精度影响较小)

7.4.2 模型量化

7.4.2.1 GPTQ

问题背景

训练后量化(Post-Training Quantization, PTQ)需要在的情况下将量化到低位宽。GPTQ是一种基于近似二阶信息的逐层量化方法。

GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers

核心思想

GPTQ(Generative Post-Training Quantization) 基于OBS(Optimal Brain Surgeon)框架,通过海森矩阵的逆来指导量化:

\(\delta w = -\frac{w_q - w}{[H^{-1}]_{\{ii\}}} \cdot H^{-1}_{\{:,i\}}\)

其中:

  • \(w\) 是原始权重
  • \(w_q\) 是量化后的权重
  • \(H\) 是海森矩阵逐层量化流程
python
def gptq_quantize_layer(layer_weight, bits=4, group_size=128):
    """
    GPTQ  逐层量化
    """
    W = layer_weight.data.clone()
    rows, cols = W.shape

    # 计算海森矩阵的逆(使用Cholesky分解)
    H = compute_hessian(W)
    H_inv = cholesky_inverse(H)

    # 逐列量化
    for i in range(cols):
          # 量化当前列
          w_col = W[:, i]
          scale = w_col.abs().max() / (2**(bits-1) - 1)
          w_q = torch.round(w_col / scale).clamp(-2**(bits-1), 2**(bits-1)-1)

          # 计算误差
          err = (w_col - w_q).unsqueeze(1)

          # 更新剩余列(补偿量化误差)
          if i < cols - 1:
               W[:, i+1:] -= err @ H_inv[i, i+1:].unsqueeze(0)

          W[:, i] = w_q

    return W

分组量化

为减少海森矩阵计算开销,GPTQ使用分组量化

python
def gptq_quantize_grouped(model, bits=4, group_size=128):
   """GPTQ  分组量化"""
   for name, layer in model.named_modules():
        if isinstance(layer, nn.Linear):
              weight = layer.weight.data

              #   按group_size分组
              for i in range(0, weight.shape[1], group_size):
                   end = min(i + group_size, weight.shape[1])
                   weight_group = weight[:, i:end]

                   #   每组独立量化
                   weight_q = gptq_quantize_layer(weight_group, bits, group_size)

                   weight[:, i:end] = weight_q

              layer.weight.data = weight

精度分析

GPTQ 在LLaMA模型上的表现:

模型FP16GPTQ-4bit精度损失
LLaMA-7B5.685.842.8%
LLaMA-13B5.095.253.1%
LLaMA-30B4.104.253.7%
LLaMA-65B3.533.684.2%

7.4.2.2 AWQ

核心思想

AWQ(Activation-aware Weight Quantization)发现:并非所有权重的bit都同等重要。保护对激活敏感的权重可以显著提升量化精度。

激活感知缩放

AWQ 通过分析激活分布来识别重要权重:

\(s = \left(\frac{|W|}{|X|^{\alpha}}\right)^{\beta}\)

其中 \(X\) 是激活值,\(\alpha, \beta\) 是超参数。保护重要权重

python
def awq_quantize_layer(weight, activation, bits=4, group_size=128):
      """
      AWQ 量化
      """
      #   计算每个通道的激活幅度
      act_scale = activation.abs().mean(dim=0)

      #   计算权重重要性
      weight_scale = weight.abs().max(dim=0)[0]

      #   计算保护缩放因子
      importance = act_scale ** 0.5 * weight_scale

      #   分组量化,使用重要性指导
      for i in range(0, weight.shape[1], group_size):
            end = min(i + group_size, weight.shape[1])
            group_importance = importance[i:end]

            #   重要组使用更高精度或特殊处理
            if group_importance.mean() > threshold:
                 #   保护重要权重
                 weight[:, i:end] = quantize_with_protection(
                      weight[:, i:end], bits, group_importance
                 )
            else:
                 weight[:, i:end] = normal_quantize(weight[:, i:end], bits)

      return weight

与GPTQ对比

特性GPTQAWQ
核心思想海森矩阵补偿激活感知保护
校准数据需要需要
量化时间较长较短
典型精度更高
实现复杂度中等较低

7.4.2.3 SmoothQuant

核心思想

SmoothQuant 解决激活量化的难题:激活比权重更难量化,因为激活有显著的离群值

核心洞察:将量化难度从激活迁移到权重。

数学推导

对于线性层 \(Y = XW\),引入平滑因子 \(s\)

\(Y = (X \cdot \text{diag}(s)^{-1}) \cdot (\text{diag}(s) \cdot W) = \tilde{X} \cdot \tilde{W}\)

选择 \(s\) 使得: \(s_j = \frac{\max(|X_j|)^{\alpha}}{\max(|W_j|)^{1-\alpha}}\)其中 \(\alpha\) 是迁移强度(通常0.5)。实现

python
def smoothquant_fuse_ln_linear(ln_layer, linear_layer, alpha=0.5):
       """
       融合LayerNorm和Linear层,应用SmoothQuant
       """
       #   获取激活统计
       act_scale = get_activation_scale(ln_layer)

       #   获取权重统计
       weight_scale = linear_layer.weight.abs().max(dim=0)[0]

       #   计算平滑因子
       smooth_scale = (act_scale ** alpha) / (weight_scale ** (1 - alpha))
       smooth_scale = smooth_scale.clamp(min=1e-5)

       # 应用平滑
       # 1. 缩放LayerNorm权重
       ln_layer.weight.data /= smooth_scale
       ln_layer.bias.data /= smooth_scale

       # 2. 缩放Linear权重
       linear_layer.weight.data *= smooth_scale.unsqueeze(0)
       if linear_layer.bias is not None:
             linear_layer.bias.data *= smooth_scale

       return ln_layer, linear_layer

精度分析

SmoothQuant 在W8A8(权重INT8,激活INT8)配置下的表现:

模型FP16SmoothQuant W8A8精度损失
OPT-175B8.348.461.4%
LLaMA-65B3.533.581.4%
LLaMA-2-70B3.323.381.8%

7.4.2.4 量化感知训练

问题背景

训练后量化虽然方便,但无法完全恢复量化带来的精度损失。量化感知训练(Quantization-Aware Training, QAT)在

直通估计器

QAT的核心挑战:量化函数不可导。使用直通估计器( Straight-Through Estimator,STE

python
前向: y = round(x / s) * s
反向: dy/dx = 1 (忽略round的梯度)

伪代码

python
class QuantizedLinear(nn.Module):
    def __init__(self, in_features, out_features, bits=8):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features))
        self.bits = bits
        self.scale = None

    def forward(self, x):
        #   训练时模拟量化
        if self.training:
            #  计算缩放因子
            w_scale = self.weight.abs().max() / (2**(self.bits-1) - 1)

            #  伪量化(前向量化,反向直通)
            weight_quantized = fake_quantize(self.weight, w_scale, self.bits)

            return F.linear(x, weight_quantized, self.bias)
        else:
            #  推理时使用真实量化
            return F.linear(x, self.quantized_weight, self.bias)


def fake_quantize(tensor, scale, bits):
    """伪量化:前向量化,反向直通"""
    quantized = torch.round(tensor / scale).clamp(-2**(bits-1), 2**(bits-1)-1)
    dequantized = quantized * scale

    # STE:    前向使用dequantized,反向使用原始梯度
    return tensor + (dequantized - tensor).detach()

QAT 流程

  1. 预训练:使用FP16训练模型到收敛
  2. 插入量化节点:在需要量化的层插入伪量化节点
  3. 微调:使用小学习率继续训练(通常1-10%的原始训练步数)
  4. 转换:将伪量化转换为真实量化

精度对比

方法配置LLaMA-7B PPL备注
FP16-5.68baseline
PTQINT85.751.2%损失
PTQINT46.209.2%损失
QATINT45.853.0%损失

用心记录,持续成长