把量化换到剪枝时踩过的坑

模型压缩 量化 剪枝 知识蒸馏 深度学习
作者: 醉月思 发布日期: 2022-01-10 分类: 技术实践

项目里那个 7B 参数的模型训练好了，部署到推理服务器上时，CPU 占用直接拉满，延迟干到了 3 秒多。

为什么要压缩

简单说就是硬件扛不住了。我遇到的真实场景是这样的：

  • 训练好的模型 7B 参数，FP16 精度下大约需要 14GB 显存
  • 推理服务器是一台 16GB 显存的 GPU 机器
  • 单次推理延迟 3.2 秒，QPS 不到 0.3
  • 客户要求延迟降到 500ms 以内，QPS 要到 2 以上

听起来很夸张，但这是真实的生产环境数据。大模型不是什么新玩意儿，但在实际部署时，硬件成本、推理延迟、并发能力这些硬指标摆在那里，不压缩就上不了线。

模型压缩的目标很简单：在不显著降低精度的前提下，让模型变得更小、更快。但这个"不显著降低"本身就充满了不确定性，不同任务对精度的敏感度完全不一样。

量化：从 FP16 到 INT8

量化是最直接的压缩方式。把 FP16 的参数量化到 INT8，理论上模型大小减半，推理速度提升 2-4 倍。

基础量化实践

用 PyTorch 做量化大概是这样的：

import torch
from torch.quantization import quantize_dynamic

model = torch.load('model_fp16.pt')
# 动态量化：线性层量化到 int8
quantized_model = quantize_dynamic(
    model,
    {torch.nn.Linear},
    dtype=torch.qint8
)
torch.save(quantized_model, 'model_int8.pt')

看起来很简单，但实际跑的时候问题一堆。第一个坑就是精度下降。量化后的模型在我的分类任务上准确率掉了 5 个百分点，从 92% 掉到 87%。这个下降还能接受，但在一些细粒度的分类任务上可能直接把模型废了。

静态量化更稳但更麻烦

动态量化是在推理时动态计算量化参数，简单但效果一般。静态量化需要先用校准数据集确定量化参数，效果好但步骤多：

import torch
from torch.quantization import prepare, convert

model = torch.load('model_fp16.pt')
model.eval()
# 融合操作
model.fuse_modules()
# 准备量化
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
prepare(model, inplace=True)

# 用校准数据集跑一遍
with torch.no_grad():
    for data in calibration_dataloader:
        model(data)

# 实际量化
convert(model, inplace=True)
torch.save(model, 'model_static_int8.pt')

静态量化后的准确率损失通常控制在 1-2 个百分点，但校准数据集的选择很关键。我第一次随便用了 100 条数据做校准，结果在验证集上的表现比随机还差，后来换了任务相关的 1000 条数据才稳定下来。

量化后的性能提升

在一台 RTX 3090 上的实测数据：

模型精度模型大小推理延迟QPS准确率
原始模型FP1614GB3.2s0.392.1%
动态量化INT87.5GB1.8s0.687.3%
静态量化INT87.5GB1.1s0.991.0%

静态量化确实有效，但离目标还有差距。延迟从 3.2 秒降到 1.1 秒，QPS 从 0.3 提升到 0.9，但距离 500ms 延迟和 QPS 2 的目标还差不少。

剪枝：砍掉不重要的连接

量化的本质是用更低的精度表示参数，剪枝则是直接砍掉不重要的连接。看起来更暴力，但有时效果不错。

非结构化剪枝

最简单的剪枝是非结构化剪枝，也就是直接把权重小的参数设为 0：

import torch.nn.utils.prune as prune

model = torch.load('model_fp16.pt')
# 对所有线性层进行剪枝
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Linear):
        prune.l1_unstructured(module, name='weight', amount=0.2)
        # 移除 mask，让剪枝永久生效
        prune.remove(module, 'weight')

torch.save(model, 'model_pruned.pt')

这里 amount=0.2 表示剪掉 20% 的参数。听起来合理，但实际跑的时候发现剪枝率不能随便设。30% 以上的剪枝率会导致模型性能断崖式下跌，我的任务上 20% 是个比较安全的上限。

非结构化剪枝的另一个问题是稀疏矩阵的推理加速不明显。理论上参数变少了，但现在的硬件对稀疏计算支持不够好，实际推理速度提升有限。

结构化剪枝更实用

结构化剪枝是砍掉整个神经元或者通道，而不是单个参数：

# 结构化剪枝示例
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Linear):
        prune.ln_structured(module, name='weight', amount=0.2, n=2, dim=0)
        prune.remove(module, 'weight')

结构化剪枝后的模型是真正的小模型，推理速度会有明显提升。但挑战在于剪枝后的模型结构不规整，需要重新训练或者微调来恢复精度。

我尝试过 30% 的结构化剪枝，然后重新训练了 3 个 epoch，准确率能恢复到原始水平的 95% 左右。模型参数从 7B 降到 4.9B，推理延迟从 3.2 秒降到 2.1 秒，提升明显但还不够。

知识蒸馏：用教师教学生

知识蒸馏的思路是用一个大的教师模型训练一个小的学生模型。教师的输出提供"软标签"，让学生在保持精度的同时变得更小。

基础蒸馏实现

import torch
import torch.nn as nn
import torch.nn.functional as F

# 教师模型
teacher = torch.load('teacher_fp16.pt').eval()
# 学生模型（结构更小）
student = torch.load('student_fp16.pt')

def distillation_loss(student_logits, teacher_logits, labels, temperature=3.0, alpha=0.5):
    # 软标签损失
    soft_loss = nn.KLDivLoss(reduction='batchmean')(
        F.log_softmax(student_logits / temperature, dim=1),
        F.softmax(teacher_logits / temperature, dim=1)
    ) * (temperature ** 2)
    # 硬标签损失
    hard_loss = nn.CrossEntropyLoss()(student_logits, labels)
    return alpha * soft_loss + (1 - alpha) * hard_loss

# 训练循环
for data, labels in train_loader:
    with torch.no_grad():
        teacher_logits = teacher(data)
    student_logits = student(data)
    loss = distillation_loss(student_logits, teacher_logits, labels)
    loss.backward()
    optimizer.step()

这里的关键是温度参数 temperature 和损失权重 alpha。温度越高，软标签的分布越平滑，学生学到的信息越多但噪声也更多。alpha 控制软标签和硬标签的平衡。

蒸馏的效果和坑

我用 7B 的教师模型蒸馏了一个 3B 的学生模型，训练了 10 个 epoch 后：

模型参数量推理延迟QPS准确率
教师 7B7B3.2s0.392.1%
学生 3B3B1.4s0.789.5%
学生 3B 蒸馏3B1.4s0.791.2%

蒸馏确实有效，学生模型的准确率提升了 1.7 个百分点。但问题在于学生模型的推理延迟还是不够，离 500ms 的目标差得远。

知识蒸馏的另一个坑是超参数敏感性。温度从 2.0 调到 4.0，准确率能差 3-4 个百分点。这部分需要针对具体任务调参，没有通用答案。

综合方案：量化和剪枝结合

单一方法通常不够，实际项目中往往需要组合多种技术。我最后采用的方案是：结构化剪枝 20% + 静态量化 INT8 + 蒸馏微调。

# 1. 结构化剪枝
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Linear):
        prune.ln_structured(module, name='weight', amount=0.2, n=2, dim=0)
        prune.remove(module, 'weight')

# 2. 重新训练恢复精度
model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
for epoch in range(5):
    for data, labels in train_loader:
        logits = model(data)
        loss = nn.CrossEntropyLoss()(logits, labels)
        loss.backward()
        optimizer.step()

# 3. 静态量化
model.fuse_modules()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
prepare(model, inplace=True)
with torch.no_grad():
    for data in calibration_dataloader:
        model(data)
convert(model, inplace=True)

这套组合拳下来，最终的模型表现：

方法参数量模型大小推理延迟QPS准确率
原始模型7B14GB3.2s0.392.1%
剪枝+量化+蒸馏5.6B5.8GB0.42s2.490.8%

组合方案与原始模型的延迟、QPS 和准确率差异很大，用柱状图对照比逐行读表更直观。

原始 7B 模型与剪枝+量化+蒸馏组合方案在延迟、QPS、准确率和模型大小上的对比

终于达标了：延迟 420ms、QPS 2.4，准确率仅降 1.3 个百分点，对多数生产任务足够可用。

踩过的坑总结

这些年折腾模型压缩，踩过的坑确实不少：

  1. 过度压缩导致模型崩坏。60% 的剪枝率听起来很诱人，但模型基本废了。压缩要循序渐进，先做小规模测试。

  2. 校准数据集选择错误。用 100 条无关数据做量化校准，模型性能暴跌。校准数据要和任务场景相关。

  3. 忽略硬件特性。非结构化剪枝理论上能减少参数量，但硬件加速支持不好，实际推理速度提升有限。

  4. 超参数敏感性。蒸馏的温度、剪枝率、量化参数这些都对最终效果影响巨大，必须针对具体任务调参。

  5. 评估不充分。只在验证集上看准确率，上线后发现实际场景性能不达标。压缩后的模型要在真实场景数据上充分评估。

  6. 部署复杂性。量化后的模型需要特定推理引擎支持，剪枝后的模型结构不规整，这些都会增加部署复杂度。

一些实用建议

结合经验，这里有一些实际项目中的建议：

  • 先量化，再剪枝。量化的改动较小，风险低，适合作为第一步。
  • 结构化剪枝优先。非结构化剪枝的理论收益在实际硬件上很难体现。
  • 蒸馏要有合适的学生模型。学生模型结构太小，教师再好也教不出来；太大则失去压缩意义。
  • 保留完整评估链。从训练到压缩再到部署，每个环节都要有完整的评估数据。
  • 考虑部署环境。推理引擎、硬件支持、并发模型这些都会影响最终方案选择。

模型压缩没有银弹，不同任务、不同硬件环境、不同精度要求都需要不同方案。但这套组合拳下来，大部分场景都能找到合理的平衡点。

后续的可能

现在的模型压缩技术还在不断发展，比如新的量化方法、更高效的剪枝策略、更好的蒸馏算法。这些技术能让更大的模型在更小的设备上运行，但离"无损压缩"还差得远。

在实际项目中，模型压缩往往和硬件升级、架构优化一起考虑。单纯靠压缩很难解决所有性能问题，但它确实是成本效益比较高的优化手段。

折腾模型压缩这些年，最大的感受是：理论上的最优往往不是实践中的可行，而工程实践中的可行方案往往不是理论上最优。这种张力和权衡，才是工程实践的真实面貌。

版权声明: 本文首发于 指尖魔法屋-把量化换到剪枝时踩过的坑（https://blog.thinkmoon.cn/post/176-model-compression-quantization-pruning/） 转载或引用必须申明原指尖魔法屋来源及源地址！

相关阅读

根据标签、分类和主题为你挑选的相近文章