★ FEATURED ARTICLE 大家都知道大模型很强但部署起来钱包真的扛不住。我们有个业务场景需要把一套效果不错的开源基座模型塞进线上容器里跑实时推理显存预算卡得死延迟要求又高。一开始想直接上更小的模型但效果回调太明显业务方不答应。最后被逼着把模型蒸馏和剪枝这套组合拳完整走了一遍踩了不少坑也攒了一堆可以复用的经验。这篇文章就把整个工程实现过程、每一步的取舍思路、以及那些文档里不会写的细节全部摊开来讲。如果你正面临模型太大、推理太慢、部署成本超预算的问题或者想把自己辛辛苦苦训好的大模型压缩成能上线的小模型这篇文章应该能帮你少走很多弯路。我默认你已经有PyTorch、模型训练和推理的基础文中的代码片段都是工程里实际跑过的可以直接落到你的项目里改着用。1. 先想清楚压缩路线蒸馏和剪枝到底该先上哪个很多人在做模型压缩时第一反应都是“既然模型太大那就剪掉一些不重要的参数”。这个方向本身没错但在动手剪枝之前必须先把蒸馏和剪枝各自的边界搞清楚否则很容易出现“剪完模型变小了但精度崩了然后怎么调都回不来”的尴尬局面。1.1 蒸馏解决的是“模型能力迁移”问题剪枝解决的是“结构冗余”问题知识蒸馏的核心思路是让一个小模型学生去学习一个大模型教师的输出分布而不是仅仅学习硬标签。举个例子教师模型对于一张猫的图片可能输出“猫 0.7狗 0.2狐狸 0.1”这样的软概率分布而硬标签只告诉我们“答案是猫”。软分布里包含了“猫和狗有点像”这类信息这对小模型来说价值巨大因为它等于把小模型原本学不到的“常识”都塞给了它。剪枝的思路完全不同。它不改变模型的学习目标而是直接在结构上做文章。比如一个层有 512 个通道经过重要性评估后发现其中 100 个通道的权重接近零、对最终输出几乎没有贡献那我们就把这 100 个通道砍掉。这样矩阵乘法直接变小显存占用和计算量立竿见影地下降。关键区别在于蒸馏是从头训练一个小模型是在“教”适合压缩比大的场景能从 7B 压到 2B 这种量级但需要重新训练成本高。剪枝是在已有模型上直接做减法是在“切”适合在保留原模型结构体系的前提下做 20%-40% 的瘦身不需要从零开始成本低但压缩比有限。大多数项目真正需要的是先剪掉多余的结构让模型变小再通过蒸馏把剪枝损失的精度补回来。顺序反过来的话相当于先蒸馏一个完整的小模型再对这个小模型做剪枝两次精度损失叠加效果很难看。1.2 判断你属于哪种情况先做一次成本与收益的粗略测算动手之前别急着看论文先做一道小学数学题。拿我们的场景举例原始模型 7B显存占用大概 14GBFP16线上容器给的显存上限是 6GB单次推理延迟不能超过 500ms。用 FP16 精度估算6GB 显存大约只能装下 3B 参数量的模型加上激活值和中间缓冲实际可用参数量还得再打个折。所以压缩目标是把模型从 7B 压到 3B 以下同时保证下游任务精度不掉超过 3 个百分点。这个目标单纯靠剪枝基本不可能实现。7B 剪掉 50% 的通道结构上已经称得上“千疮百孔”注意力头和 FFN 比例完全被破坏下游任务精度必然崩。单纯靠蒸馏的话相当于从 7B 模型去蒸馏出一个 3B 学生训练成本会很高但理论上是可行的。所以我们的路线定为先做结构化剪枝把 7B 压到 4B 左右损失的部分精度通过蒸馏来恢复再从 4B 蒸馏微调到 3B。这样一个组合动作下来最终模型 3B 左右实际效果只比原始 7B 掉了 2 个百分点以内部署也稳稳落进了预算内。你可以按自己的场景套用这个判断逻辑压缩比小于 30% 优先剪枝压缩比大于 50% 优先蒸馏30%-50% 这个区间就剪枝加蒸馏组合。2. 蒸馏的工程实现从教师模型选型到训练策略的完整落地如果你决定走蒸馏这条路恭喜你接下来一段时间你都会和数据、loss、回调和各种玄学参数打交道。蒸馏最核心的问题有三个教师模型怎么选、软标签怎么用、蒸馏损失怎么组合。每一个都有讲究。2.1 教师模型选型的“门当户对”原则先说一个容易犯的错误为了追求极致效果选了比实际需求大得多的教师模型。某个测试场景里我们一开始用的是几百B级别的模型当教师蒸馏 3B 学生结果人的直觉是教师越强学生越强但实际的偏差在于两个模型的能力差距过大学生模型根本学不到教师那么多知识反而被教师的一些极端输出分布带偏蒸馏出来的模型在某些任务上的表现还不如直接用硬标签训练出来的模型。后来换成一套与目标学生更匹配的教师效果反而明显变好了。选教师模型时别只看它的指标有多强更要看它的输出分布是否友好、学生模型是否有能力拟合这个分布。圈里通常用的选择标准是学生在教师指导下能在验证集上稳定收敛如果蒸馏 loss 怎么都降不下去大概率是教师选偏了。教师模型和学生模型的 tokenizer 最好保持一致词汇表如果不同意味着学生要把教师的输出先映射到自己的词汇空间这个映射过程本身就会损失信息至少在蒸馏初始阶段会让训练非常不稳。2.2 蒸馏损失函数设计不止是 KL 散度现在谈最核心的部分。我们蒸馏用的损失函数可以拆解为三部分每一部分都有存在的理由。第一部分是硬标签损失就是标准的交叉熵。这一部分保证了模型依然能学到任务本身的真实结构不会因为教师的输出有偏而被带跑。第二部分是软标签损失使用的是教师模型输出概率分布与学生模型输出概率分布的 KL 散度。适当加热温度 T也就是把概率分布变得更平滑。温度 T 太高分布接近于均匀分布学生什么都学不到温度 T 太低软标签退化成硬标签蒸馏就没意义了。我们实践中文本模型 T 取 2-4图像模型 T 取 3-6 是相对稳定的区间因为图像数据类别间的相似度信息更丰富需要更高的温度才能暴露出来。第三部分是中间特征对齐损失。这个部分更多是经验之谈。只约束输出层的话学生模型虽然在最终预测上能对齐教师但中间层的表征结构可能完全不匹配导致学生在推理时泛化能力差。可选的做法是把教师和学生的中间层表征做一个 L2 距离约束让它们在某些层上的特征分布接近。需要注意的是如果教师和学生模型的隐层维度不同需要加一个小的映射层把维度对齐这个映射层在训练结束后要丢掉。下面是一段实际跑过的蒸馏训练代码骨架重点是自蒸馏时的 stop-gradient 处理、损失权重、以及软标签的生成方式。import torch import torch.nn as nn import torch.nn.functional as F def distill_loss(student_logits, teacher_logits, labels, T3.0, hard_weight0.3, soft_weight0.7, alpha0.5): # 硬标签损失 hard_loss F.cross_entropy(student_logits, labels) # 软标签损失temperature 加热 student_soft F.log_softmax(student_logits / T, dim-1) teacher_soft F.softmax(teacher_logits / T, dim-1) soft_loss F.kl_div(student_soft, teacher_soft, reductionbatchmean) * (T * T) # 中间特征对齐在训练循环里单独算这里只整合输出层损失 total_loss hard_weight * hard_loss soft_weight * soft_loss # alpha 是用于后续和特征对齐损失做加权的保留项 return total_loss, hard_loss, soft_loss有个细节必须提醒教师模型的输出应该提前一次性算好存成缓存不要每轮训练都重新过一遍教师模型。7B 教师模型跑一次前向非常费时如果放进训练循环里训练速度会慢一个数量级。我们当时是先离线把全量训练数据用教师模型跑一遍存成 logits 的 npy 文件训练的时候直接读取。2.3 蒸馏训练策略数据配比、学习率与步数的实操经验蒸馏不只是把损失函数换一下就完事数据配比和训练节奏同样重要。我们踩过的一个坑是直接拿全部原始训练数据做蒸馏结果学生模型在公共基准上表现不错但在业务特有的长尾数据上反而退化。原因是长尾数据在原始数据集里占比太低了学生模型根本没学到。后面调整了数据配比蒸馏数据中基础数据占六成业务长尾数据占四成效果立刻变好业务侧的指标明显回升。另外一个经验是蒸馏时学习率一定要比正常训练更低一般低 1-2 个量级。因为学生模型的目标不是从零学习而是已经被初始化成教师模型的一部分了在预训练阶段甚至直接复用教师的前若干层权重。高学习率会把已经学到的知识冲掉。蒸馏训练步数方面我们的经验是大概是正常训练步数的 0.5-1 倍就足够。如果训练步数太长学生模型会过拟合到教师模型的输出模式反而丢失自己的泛化能力。这个判断标准就是每隔一段时间在验证集上看效果效果开始波动甚至下降时果断提前停。3. 剪枝的工程实现结构化剪枝怎么切才能保住精度上限蒸馏是把模型“越教越好”剪枝则更接近手术刀活讲究怎么切得精准。剪枝分为非结构化剪枝和结构化剪枝业界成熟方案大多是结构化剪枝因为非结构化剪枝后的稀疏矩阵在通用硬件上加速效果极差。3.1 非结构化剪枝与结构化剪枝为什么我放弃了前者非结构化剪枝做的事情是把权重矩阵中绝对值小于某个阈值的元素直接置为零得到一个稀疏矩阵。这个思路在理论压缩率上很好看但在实际推理时由于 GPU 对稀疏矩阵的加速支持并不通用稀疏度不够高时甚至可能比稠密矩阵更慢。前几年我们试过把之类的方法套到 Transformer 上效果只能说惨淡。结构化剪枝则不同它移除的是整行、整列或者整个通道。在 Transformer 架构里结构化剪枝一般作用在两个地方注意力层剪掉部分注意力头FFN 层剪掉部分中间维度。这样剪完后矩阵形状是规则的不需要特殊算子支持任何推理框架都能直接跑。结构化剪枝的代价是需要重新训练恢复精度非结构化剪枝则几乎不需要训练。这是一个精度和效率的权衡为了部署时的性能我们必须选择结构化剪枝。3.2 剪枝目标选择哪些层可剪、哪些层不能碰不是所有层都适合剪。我们总结出的经验是注意力头的冗余度通常很高一个大模型里存在大量功能重复的注意力头一些注意力头对最终输出的贡献趋近于零这些头可以安全地剪掉。但 FFN 层剪起来要慎重尤其是靠近输出层的 FFN 层它们保存着大量任务相关的关键知识过度剪枝会直接导致模型表达能力断崖式下跌。这里要引入一个概念叫“层敏感性”。即使同一个类型的层在不同位置上被剪掉的容忍度也完全不同。靠近输入层的 embedding 层几乎不能剪因为它们负责把 token 映射成向量所有后续层都依赖这些向量的质量。中间层的冗余度较高可剪的空间大。靠近输出层的层负责把隐藏状态映射到词表概率分布剪掉一些维度可能直接影响最终预测。最好是先逐层尝试性地给每一层一个剪枝幅度预算然后观察每一层对剪枝的敏感程度敏感度低的层多剪一点敏感度高的层少剪甚至不剪。3.3 剪枝实操流程重要性评估、剪枝动作和微调恢复剪枝流程可以分为三步评估、执行、恢复。每一步都有独立的坑。评估阶段的核心是计算每个通道的重要性。常用的方法包括基于权重绝对值的 L1 范数、基于激活值的统计信息、以及基于损失函数梯度的敏感度分析。我们在实际项目中用的是激活值统计加损失函数梯度加权的方式比单纯看权重大小更准确因为权重绝对值大不代表它对最终任务重要关键要看它在推理时对中间特征的影响。执行阶段是真正动手剪。以 PyTorch 为例对一个 Linear 层做通道剪枝本质就是把权重矩阵的某些行或列删掉同时把相邻层的维度对齐。这个过程需要非常细心因为一个层剪掉维度后它的前一层和后一层的权重矩阵也要跟着修改。我们封装了一个简单的剪枝工具类方便复用。import torch import torch.nn as nn def prune_linear_layer(layer: nn.Linear, keep_indices: torch.Tensor, prune_bias: bool False): # keep_indices 是需要保留的通道索引 new_layer nn.Linear(layer.in_features, len(keep_indices), biaslayer.bias is not None) with torch.no_grad(): new_layer.weight.data layer.weight.data[keep_indices, :] if layer.bias is not None and prune_bias: new_layer.bias.data layer.bias.data[keep_indices] return new_layer这里有两个容易出错的细节。第一个是相邻层的对齐问题。如果剪掉了 A 层的输出通道那么接在 A 层后面的 B 层的输入维度必须同步调整对应地要把 B 层权重矩阵的相应列删掉。很多人在这一步搞错导致模型结构错乱运行时报维度不匹配。第二个是 B 层对应的 embedding 矩阵同步。如果剪到 embedding 层token embedding 和 lm_head 的维度是共享的剪一个必须两个都剪。剪枝执行完后不能直接拿去用必须做微调恢复这是整个剪枝流程里最重要的一步。剪枝之后模型精度一定会掉掉得多少取决于剪枝幅度。我们的经验是剪掉 20% 参数时微调只需要用原始训练数据的很小一部分大概 5%-10% 的量配合很低的学习率就可以恢复到接近原始水平。剪掉 40% 以上时就需要完整的蒸馏流程来补救了单纯微调已经不够。4. 评估体系剪完不能只看指标还要看中间层状态和推理画像模型压缩做完之后最怕的是“看起来指标都正常一上线就出错”。很多团队在压缩完模型后只跑一遍标准评测集看精度掉了多少就草草收工这样很容易漏掉隐藏的严重问题。我们后来把评估拆成了四个层面来执行。第一个层面是标准指标评估也就是验证集上的准确率、F1、BLEU 这类常规指标。这个层面只能反映压缩后的模型是否“大体可用”参考价值有限因为验证集往往跟线上真实分布存在偏差。第二个层面是中间层表征诊断。这一步是很多团队会漏掉的。我们会在压缩后的模型上跑一批数据把每一层的输出激活值的分布统计出来与压缩前的模型做对比。如果某一层的激活值分布发生了剧烈变化说明剪枝或蒸馏对那一层的破坏比想象中严重需要重点恢复。举个例子FFN 层的中间激活值往往存在大量幅值异常大的离群点剪枝后这些离群点如果被切掉一部分后续 LayerNorm 的缩放比例就会整体偏移表现就是模型在个别样本上输出极端异常。第三个层面是推理性能画像。这是工程视角独有的评估维度。我们需要在目标硬件上实际测延迟、吞吐量和显存占用而不是只看参数量。模型参数量减半显存不一定减半因为推理时中间激活值的显存占用占比很高。我们当时在某个推理框架里测了一个剪枝模型参数量减少了 35%但实际显存只下降了 18%原因就是中间激活值太大成为显存的新瓶颈。第四个层面是长尾数据回放测试。业务数据的分布往往是长尾的核心高频场景可能只占小部分大量的是低频但同样关键的小场景。我们会在压缩后的模型上跑一遍线上收集的真实长尾样例尤其那些容易混淆的、边界模糊的样本对比压缩前模型在这些样例上的表现差异。压缩模型常常在“简单样本”上表现和原来差不多但一到这些边界样本上就原形毕露。把这四层评估链路跑完一遍后才能比较确定地判断这次模型压缩是否真正合格。5. 组合拳实战一次完整的“剪枝 蒸馏”压缩项目复盘理论说了不少还是用一整段项目过程把这些东西串起来你更能感受到每一步之间的依赖关系。这是一个虚构的模拟项目X场景是对一套 7B 模型做压缩部署目标精度损耗控制在 3 个百分点以内显存降到一半以下。5.1 第一步诊断原始模型的冗余度与敏感层分布在动手剪之前我们先用一小批数据在原始 7B 模型上做了层敏感性分析。做法是逐层做小规模剪枝测试每次只剪某一层的 10% 通道观察验证集指标的变化。如果剪某一层时指标暴跌说明这一层敏感度极高不能动如果剪完后指标几乎不变说明冗余度很高可以在后续分配更多的剪枝额度。实验结果是注意力层的敏感度普遍较低尤其是中层的注意力头很多头剪掉后指标完全不变。但 FFN 层靠后的部分敏感度极高一旦剪多生成质量明显下降。于是剪枝预算分配策略定为注意力层承担大约六成的剪枝量FFN 层只承担四成且主要从靠近输入侧的中间层 FFN 下手。这个诊断过程虽然耗时但非常值得。没有这一步盲剪就像蒙眼走钢丝剪完哪一层出问题都找不到原因。5.2 第二步分阶段执行结构化剪枝并同步蒸馏恢复因为压缩比目标是超过 50%一轮剪完再恢复基本不现实。我们的做法是分三个阶段走。第一阶段把 7B 模型剪到 5.5B 左右注意力头剪掉约四分之一FFN 中间维度剪掉约 15%。这一轮剪完后精度掉落比较大在后续用小批量数据微调恢复了大约一半的损失。第二阶段在第一阶段恢复后的模型基础上以最初的 7B 模型作为教师进行一次完整的蒸馏训练。这次蒸馏的核心目标是补回第一阶段剪枝损失掉的精度学生模型的结构保持 5.5B 不变。训练数据用的是业务侧和公共数据的混合温度 T 设为 3.0训练约两万步。第三阶段把 5.5B 模型再剪到 3B 左右这次剪枝幅度大注意力头和 FFN 层都压得比较狠。这一轮剪完之后不再直接微调而是立刻接一轮以原始 7B 为教师的第二次蒸馏把精度尽量拉回来。实际效果是我们的精度在两个阶段的递减中最终和原始 7B 模型的差距控制在约 2.3 个百分点达到了业务能够接受的范围。5.3 第三步部署前推理优化与显存画像模型压缩完了还得能上线才算真正结束。在部署前我们做了这样几件事把模型从 FP16 换成 INT8 量化这一步再压缩约一半显存但需要特别关注敏感层不量化。用逐层做 INT8 量化评测每一层量化后的精度损失找出对量化最敏感的那几层让它们在推理时保持 FP16其余层用 INT8。开启了 KV Cache 的量化这是长文本场景下显存杀手。将 KV Cache 从 FP16 变成 INT8显存占用又降了一截且对精度影响很小。推理时对模型做层融合也就是把 LayerNorm 和 Attention 的计算合并、把激活函数和 FFN 的线性层融合减少内核启动次数延迟下降明显。最终上线时的模型参数量约 3.2B显存占用在 INT8 混合精度下约 5.8GB满足了 6GB 的预算线单次推理延迟从原始模型的接近 1000ms 降到约 450ms业务吞吐量提升了一倍以上。6. 部署阶段必须关注的几个隐性坑压缩和蒸馏做完整个项目其实才走完一半。部署阶段还有几个很容易被忽视的问题我单独拿出来说。6.1 推理框架对剪枝后模型结构的兼容性剪枝后的模型结构已经不是原始的标准结构了注意力头的数量、FFN 的中间维度可能都不是常见配置。很多优化过的推理框架在底层做算子优化和内存布局时是写死假设了常见结构的遇到自定义结构可能退回到低效的通用算子甚至直接报错。我们碰到过一次奇怪的现象剪枝后的模型在 PyTorch 原生推理时延迟正常但换到一个优化推理框架后延迟反而比原生推理还高。查了很久最后发现是这个框架对“不规整”的注意力头数不支持优化走了通用路径。解决方式是在推理框架的配置中显式关闭某些自动优化或者手动指定算子实现方式。这个属于经验积累提前知道能省很多排查时间。6.2 动态形状对推理性能的影响剪枝后的模型如果存在动态形状比如不同 batch 的序列长度变化很大对某些推理框架的显存分配策略非常不友好。我们把输入做了 padding 到固定长度并开启推理框架的显存预分配功能整体吞吐量提升了不少。这里补充一下padding 后用注意力掩码把填充部分遮掉保证计算结果不受影响即可。6.3 蒸馏和剪枝之后的模型校准集收集量化尤其是 INT8 量化需要一个校准集来统计数据分布范围。很多团队直接拿训练集的一部分做校准但我们发现训练集和线上真实分布差距大时量化后精度会掉得比预期多。更推荐的做法是拿线上真实流量中随机采样的数据做校准集采集量大概几百到几千条就够但代表性必须强。校准集的分布越接近真实线上分布量化后的量化误差就越小。整个项目做下来我最想分享的一条心得是模型压缩不是一次性的调参实验而是一套需要体系化评估、分阶段执行、部署联动验证的工程流程。每一步的单点效果都要放在完整链路里判断否则很容易陷入“压缩一时爽上线火葬场”的窘境。如果你正在规划类似的压缩项目建议先花一周时间把诊断和评估体系搭好再动刀。
阅读完成 · 觉得有帮助？
