《AI大模型应用开发实战从入门到精通共60篇》051、模型剪枝与蒸馏：让大模型变小变快的核心技术

051、模型剪枝与蒸馏：让大模型变小变快的核心技术

上周三凌晨两点，我盯着终端里那个报错发呆——一块A100 80G显存，跑一个7B的LLaMA推理，居然OOM了。检查了半天，发现是模型加载时把KV cache的max_seq_len设成了4096，加上batch size 4，显存直接炸穿。同事在旁边说：“要不换个更小的模型？”我摇头，业务场景要求必须保留这个特定微调后的能力。那晚我翻出了压箱底的模型压缩方案，最终把模型体积砍掉60%，推理速度提升3倍，精度只掉了不到1个百分点。

这不是魔法，是剪枝和蒸馏。

剪枝：砍掉那些“摸鱼”的神经元

先说说剪枝。很多人以为剪枝就是简单地把权重接近0的参数删掉，实际操作过就知道，直接这么干模型就废了。

我最早踩过一个坑——用L1范数对全连接层做非结构化剪枝，把权重绝对值小于0.01的全置零。结果模型输出全是乱码。后来才明白，剪枝不是“删参数”，而是“让参数变稀疏但保持功能”。

结构化剪枝才是工程上能用的方案。比如对Transformer的注意力头做剪枝。我习惯的做法是：先跑一批验证集数据，统计每个注意力头的平均注意力权重分布。那些对最终输出贡献极小的头（比如注意力权重几乎均匀分布，或者大部分时间都集中在[CLS] token上），直接砍掉。

代码里这样写：

# 这里踩过坑：千万别用随机batch统计，要用验证集全量数据defcompute_head_importance(model,dataloader):head_importance=torch.zeros(model.config.num_hidden_layers,model.config.num_attention_heads)model.eval()withtorch.no_grad():forbatchindataloader:outputs=model(**batch,output_attentions=True)# 注意：output_attentions=True会返回所有层的注意力权重# 别这样写：直接取mean，因为不同样本的注意力分布方差很大forlayer_idx,layer_attninenumerate(outputs.attentions):# layer_attn shape: [batch, heads, seq_len, seq_len]# 我们关心的是每个头对输出的影响，用attention weight的熵来衡量attn_entropy=-torch.sum(layer_attn*torch.log(layer_attn+1e-8),dim=-1)head_importance[layer_idx]+=attn_entropy.mean(dim=(0,2))returnhead_importance/len(dataloader)

统计完重要性后，我一般保留top-K的头，K根据压缩目标动态调整。比如目标压缩30%，那就砍掉重要性最低的30%的头。注意，砍头之后要重新调整模型配置，把num_attention_heads改小，同时确保hidden_size能被新的head数整除——这个细节我吃过亏，不改配置直接mask掉权重，推理时显存一点没省。

蒸馏：让大模型当老师

剪枝能砍掉冗余结构，但精度损失是硬伤。这时候蒸馏就派上用场了。

蒸馏的核心思想很简单：让大模型（Teacher）教小模型（Student）。但具体怎么教，门道很多。

我最早做蒸馏时，直接拿Teacher的logits做soft label，用KL散度训练Student。结果Student学了一堆噪声——因为Teacher在低概率区域也有输出，那些概率值虽然小，但累积起来会干扰Student的学习。

正确的做法是加温度系数。温度T越高，softmax输出的分布越平滑，Student能学到Teacher的“暗知识”。我一般T取2-4，具体看任务。

# 别这样写：直接用原始logits算KL散度# loss = F.kl_div(student_logits.log(), teacher_logits, reduction='batchmean')# 正确做法：加温度defdistillation_loss(student_logits,teacher_logits,temperature=3.0):# 这里踩过坑：softmax的dim要指定，默认是最后一维soft_student=F.log_softmax(student_logits/temperature,dim=-1)soft_teacher=F.softmax(teacher_logits/temperature,dim=-1)# KL散度乘以T^2是为了梯度尺度匹配loss=F.kl_div(soft_student,soft_teacher,reduction='batchmean')*(temperature**2)returnloss

除了logits层面的蒸馏，中间层特征也可以蒸馏。比如让Student的某层hidden state去拟合Teacher对应层的输出。但这里有个坑：Teacher和Student的hidden size可能不一样，需要加一个线性映射层对齐维度。这个映射层训练时要和Student一起更新，但推理时扔掉。

剪枝+蒸馏的组合拳

单独用剪枝或蒸馏，效果都有限。我试过只剪枝不蒸馏，压缩30%后精度掉了5个点；只蒸馏不剪枝，Student模型参数量减半但推理速度没提升多少（因为结构没变）。

真正的杀手锏是迭代式剪枝+蒸馏。流程是这样的：

1. 训练一个完整的Teacher模型（或者直接用现成的大模型）
2. 对Teacher做一次剪枝，得到压缩后的Student
3. 用Teacher的logits和中间层特征蒸馏Student
4. 对蒸馏后的Student再做一次剪枝
5. 重复步骤3-4，直到达到目标压缩率

我做过一个实验：对一个BERT-base模型（110M参数），经过3轮迭代剪枝+蒸馏，最终模型只有45M参数，在GLUE benchmark上平均精度只掉了1.2%。而直接剪枝到45M，精度掉了4.8%。

迭代的关键在于每轮剪枝的比例不要太大。我一般每轮剪10%-15%，然后蒸馏2-3个epoch。剪太多Student学不过来，精度会断崖式下跌。

工程落地的一些血泪教训

说几个实际部署时容易翻车的地方。

量化要放在剪枝和蒸馏之后。我试过先量化再剪枝，结果剪枝时因为量化后的权重分布变了，剪枝阈值完全失效。正确的顺序是：剪枝→蒸馏→量化。量化推荐用INT8，对精度影响小，推理速度提升明显。

剪枝后的模型要重新做batch normalization校准。这个很多人忽略。剪枝改变了网络结构，BN层的running mean和running variance需要重新统计。跑一遍验证集，更新BN参数，否则推理时输出会漂移。

蒸馏时Teacher和Student的输入要一致。听起来是废话，但我真见过有人用不同tokenizer处理数据，导致Teacher和Student看到的是不同的文本。蒸馏的前提是Teacher和Student在同一个语义空间里。

个人经验

做了两年模型压缩，最大的感悟是：不要追求理论上的最优压缩率，要追求工程上的可维护性。我曾经花两周时间把模型压缩到原来的20%，精度只掉了0.5%，但模型结构变得极其复杂，后续维护和迭代成本高得离谱。后来我改用结构化剪枝+蒸馏，压缩到40%，精度掉1%，但代码清晰，部署方便，团队里任何一个人都能接手。

另外，剪枝和蒸馏不是银弹。如果你的模型本身训练得就不够好（比如过拟合或者欠拟合），压缩后问题会放大。先确保Teacher模型足够强，再谈压缩。

最后，记得在部署前做一次完整的精度验证。我吃过一次亏：剪枝后的模型在测试集上精度达标，但上线后因为数据分布偏移，表现一塌糊涂。后来我在验证集里混入了20%的线上真实数据，才把问题暴露出来。

模型压缩的本质是权衡——用可控的精度损失换取推理效率。这个“可控”的边界在哪里，取决于你的业务场景。对于对话系统，1%的精度损失用户可能感知不到；但对于医疗诊断，0.1%的损失都不可接受。所以，别盲目追求压缩率，先搞清楚你的精度底线在哪里。