ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

Model-Optimizer:从训练优化到推理加速的模型压缩全流程实践

Model-Optimizer:从训练优化到推理加速的模型压缩全流程实践 去年下半年我接手了一个文本分类服务的性能优化任务模型是 BERT-base线上 P99 延迟已经涨到了 380 毫秒显存稳稳吃掉一块 12GB 的 T4。老板给的指标很直接延迟砍一半精度不能掉超过 1 个点。当时我把散落在各个项目里的优化手段——换优化器、调学习率、上混合精度、做剪枝、量化、蒸馏、换推理引擎——全部收拢成了一条工具链给它起名Model-Optimizer。一开始它只是我笔记本里的一堆脚本后来慢慢变成了一个能让团队其他人也直接上手的工作流。这篇内容不打算讲什么高大上的理论就说说我在做训练优化、模型瘦身和推理加速这条路上实际踩过的坑、测过的数据、用过的配置以及 Model-Optimizer 最终沉淀下来的那套可复用流程。如果你手头也有一批模型要优化或者准备入职做算法工程化的活这里面的东西应该能直接抄作业。1. 我为什么做 Model-Optimizer三个真实痛点逼出来的工具链做模型优化之前我一直以为优化就是调参。后来发现完全不是。模型优化是一个从训练到部署的完整链路每一项改动都会牵动另一项的表现不把它当成一个系统来对待结果就是今天改了学习率明天显存不够后天上线延迟又回去了。1.1 训练侧loss 降不下去不等于模型结构有问题有段时间我在微调一个中文 NER 模型loss 一直卡在 0.35 下不去换了更深的层、加了 dropout收敛速度反而更慢。后来把优化器从 Adam 换成 AdamW又把学习率从固定的 2e-5 改成 warmup 余弦退火loss 在同样轮数内降到了 0.18。模型结构一个字没动纯靠训练侧的优化策略就拉开了差距。这个经历让我意识到很多团队把优化模型等价于改模型结构但训练策略侧的优化空间往往被严重低估。优化器的选择、学习率调度、梯度累积、混合精度、批量大小这些因素互相耦合单独调一个看不出效果组合起来变化非常明显。1.2 部署侧模型体积和延迟的账必须用端到端的口径算另一个痛点在部署。之前团队上线过一个蒸馏后的意图识别模型单看 PyTorch 的推理时间从 12 毫秒降到了 6 毫秒大家觉得已经达标。结果换到 ONNX Runtime 之后反而比原模型还慢 2 毫秒。排查了半天才发现问题出在动态 shape 没有固定、算子没有做层融合图优化根本没走完。模型优化必须用端到端的口径来算账从输入到输出、包含前后处理的整体耗时而不是只看某个单算子的时间。Model-Optimizer 里所有指标都以整条推理链路的延迟为准一开始就把这个口径固定下来后面才好做对比。1.3 实验管理侧没有基线记录的优化等于白做第三痛点是实验记录混乱。以前同事剪完枝说精度没怎么掉结果一问基线是什么、用了哪份校验集、跑了多少次平均全都说不清楚。没有可复现的基线任何优化结论都是靠感觉。Model-Optimizer 从设计上强制每一次优化都跑同样的评测脚本自动记录模型版本、优化方式、评估指标、耗时和显存占用。数据喂进去最后自动生成一张对比表。没有基线的优化实验系统直接拒绝执行。这三点是我做这套工具的出发点把训练策略、模型压缩、推理加速、实验管理统一到一个流程里。2. 训练阶段的优化器选型与参数联动不是玄学是收敛轨迹的代价权衡第一个模块是训练侧的优化策略配置。Model-Optimizer 里维护了一份我长期测试下来的优化器推荐矩阵不同任务类型配不同的优化器默认参数。2.1 主流优化器的适用边界从 SGD 到 AdamW 到 LAMB优化器适合场景收敛速度泛化性显存占用我的实测说明SGD MomentumCNN 分类、目标检测等 CV 任务慢好低ResNet-50 上大概多跑 1.5 倍 epoch 才能达到 Adam 的精度但最终泛化更好Adam通用任务快速收敛快中中文本分类直接上 Adam 很方便但带 weight decay 时表现不如 AdamWAdamWTransformer、预训练模型微调快好中BERT 微调我用它作为默认选项比 Adam 稳不少LAMB大 batch 分布式训练快好高batch size 开到 8K 以上才有明显收益小 batch 反而浪费刚开始做优化器选型时我犯过一个错只看收敛速度谁 loss 降得快选谁。后来做 CV 分类任务时发现 Adam 虽然前期猛后期精度被 SGD 反超。原因在于 Adam 的自适应学习率相当于给每个参数单独调步长收敛快但对梯度噪声的鲁棒性弱最后会在极小值附近震荡SGD 带 Momentum 虽然步子慢但整体轨迹更平稳能在更平缓的极小值处落脚泛化自然更好。因此 Model-Optimizer 里有一个收敛轨迹对比功能把优化器在验证集上的表现按 epoch 画成曲线而不是只看最终精度。如果曲线前期拉升快、后期抖动大我会优先考虑换成泛化性更好的优化器而不是死磕 epoch 数。2.2 学习率 warmup、批量大小与梯度裁剪的联动关系学习率是另一个很容易被单点调参误导的因素。固定学习率 2e-5 训练 BERT前面几百步容易把预训练权重冲坏后来加上 10% 步数的 warmup再配合余弦退火收敛速度和最终精度都明显更好。warmup 的逻辑很简单训练初期参数离最优解还很远太大步长容易让 loss 爆炸先用小学习率预热让优化器摸清梯度方向再逐步加大步长。批量大小和学习率也不是独立的。线性缩放规则说批量翻倍学习率也大致可以翻倍但要注意上限。BERT 类模型 max lr 我一般控制在 2e-5 到 5e-5 之间超过这个范围 loss 容易跳到无法恢复的位置。梯度裁剪值我固定在 1.0主要防止偶发的大梯度把训练干崩。这些参数看起来很常规可它们是一条链路上的几个阀门单独动哪一个都会影响其他几个的表现。Model-Optimizer 会把批量大小、warmup 比例、峰值学习率、衰减策略、梯度裁剪值打包成一个组合配置跑实验时像填表单一样统一提交而不是散落在不同脚本里各改各的。2.3 不改模型结构的前提下用混合精度、梯度累积和激活检查点省显存显存不够时很多人上来就改模型结构比如变小 hidden size。其实训练侧的显存优化手段还有一大把。我实测下来最有效的是三件套自动混合精度、梯度累积、激活检查点。自动混合精度AMP把部分算子从 FP32 降到 FP16显存直接省一半A100 和 V100 上都支持。要注意的是损失缩放和梯度裁剪的配合AMP 默认会动态调整 loss scale遇到 inf 时自动跳过这步更新训练日志里如果频繁出现 loss scale 下降就得检查学习率是不是太高了。梯度累积显存不够但又想用更大 batch 时把一个小 batch 的梯度攒够若干步再更新。Model-Optimizer 里我会这样配置真实 batch size 32显存只能跑 8那就做 4 步累积。注意 BatchNorm 在这种模式下统计量会偏CV 任务里要谨慎。激活检查点activation checkpointing前向不保留中间激活反向需要时重新计算。对长序列 Transformer 尤其有效BERT 可以做几层开启、几层不开启的组合而不是全开全关这样显存和计算之间能取一个平衡。我用一个 128 长度的中文情感分类任务batch size 32未开启任何优化时显存 6.2GB开 AMP 变 3.8GB再开激活检查点变成 2.1GB。模型结构完全没动训练照样正常收敛。3. 部署阶段的模型压缩三板斧剪枝、量化和蒸馏的实际收益训练搞定了接下来是部署。Model-Optimizer 的压缩模块提供三个独立但可以叠加的流程剪枝、量化、蒸馏。策略是先蒸馏再剪枝再量化顺序反了会互相干扰。3.1 结构化剪枝为什么我不用非结构化剪枝剪枝听起来简单就是把不重要的权重置零。但实际落地时有个大坑非结构化剪枝会把模型变成稀疏矩阵PyTorch 里推理速度不但没提升有时候反而更慢因为大多数硬件对稀疏计算支持有限。所以我只做结构化剪枝比如 Transformer 的 attention head 剪枝和 FFN 维度的剪枝。试验是在一个 6 层 BERT 上做的剪掉 2 个 attention head 之后F1 掉了 0.2 个点模型参数量减少约 10%。继续剪到 4 个 head 时 F1 掉了 1.4 个点这个损失已经超出容许范围。所以剪枝不是越多越好要画一条精度-剪枝比例的曲线找一个拐点。剪枝算法我用的是基于梯度的显著性评估每个 head 的梯度越大说明它对预测的贡献越关键。训练时记录每个 head 的平均梯度按显著性排序后裁剪。工程上我建议用 PyTorch 自带的torch.nn.utils.prune配合自己写的结构化剪枝逻辑不要直接拿整个大模型做实验先在一个小模型上验证流程。3.2 PTQ 与 QATINT8 量化里的精度回撤岔路量化是把 FP32 权重和激活变成 INT8推理时能显著提速。我踩过最大的坑是 PTQ 校准数据太少校准只喂了 100 条样本量完模型在某个分类类别上直接失灵。后来把校准集扩到 500 条并且覆盖各类别问题才缓解。PTQ 与 QAT 的区别在于PTQ 是训练后直接量化不改权重QAT 是在量化感知训练中让模型自己去适应量化误差。我在中文 NER 模型上的实测数据如下方案大小降幅F1 变化推理耗时FP32 原始模型-88.412 msPTQ INT8约 4 倍87.67 msQAT INT8约 4 倍88.17 msQAT 比 PTQ 的精度回撤少了 0.5 个点但训练成本高不少需要在量化模型上继续微调。如果任务的精度余量足够我建议先上 PTQ如果掉点不能接受再上 QAT。Model-Optimizer 默认输出两种量化模型让使用者自己选。3.3 蒸馏花教师 7% 的算力保住教师 95% 的效果知识蒸馏是一个性价比很高的方案。我拿 BERT-large 做教师蒸馏了一个 4 层的小 BERT 给学生见的做法是同时用软标签和硬标签计算损失软标签蒸馏让模型学到教师输出的分布硬标签监督保持真实任务的精度。最终学生的参数量只有教师的 27%推理延迟从 20ms 降到 6ms而学生模型的 F1 是教师模型的 95% 左右。这个损失买卖是比较划算的延迟降了 70%换来了 5% 的精度回撤。如果任务本身要求极高精度蒸馏方案要谨慎如果精度余量在 5 个点以上蒸馏往往是最优先考虑的压缩手段。蒸馏和之前讲的剪枝、量化是可以叠加的。我的固定顺序是先用教师蒸馏出小模型再对小模型做结构化剪枝最后做 INT8 量化。3.4 推理引擎加速从 PyTorch 到 ONNX Runtime 和 TensorRT模型压缩完还有一个常见加速手段是换推理引擎。同一份 BERT 模型我从 PyTorch 切到 ONNX Runtime默认设置下就快了 20% 左右再切到 TensorRT FP16能比 PyTorch 快 2 到 3 倍。但换推理引擎的坑很多。最开始我直接用公司的主干网络跑到 TensorRT结果 FP16 模式下精度掉了 2 个多点。后来定位到是某些算子对 FP16 支持不完整导致溢出解决方法是把特定算子显式保留为 FP32或者在 ONNX 里先做算子融合。ONNX Runtime 里有个工具叫onnxruntime.transformers.optimizer专门做 Transformer 的层融合和维度优化建议任何 Transformer 模型部署前都先跑一遍。另一个值得注意的点是固定序列长度。动态序列长度在 ONNX 和 TensorRT 里都会引入额外的 shape 推断开销如果业务场景里长度波动不大直接固定到一个合理的最大值延迟会更稳定。4. 避坑实录我在 Model-Optimizer 迭代过程中推翻过的方案整个工具链不是一次成型的。中间有好几次我自认为找到了最优方案没过多久就发现被某个隐蔽问题打脸。这部分单独拿出来说是为了帮你少走重复路。4.1 全量微调改成 LoRA 后我把过拟合误判成了优化失效有一段时间为了避免全量微调带来的显存压力我把一个文本分类模型切成 LoRA 训练。跑完验证集精度比全量微调版本低了 1 个点当时第一反应是 LoRA 结构表达力不够差点把它判定为不可用。后来排查发现真正的元凶是 LoRA 的 rank 和学习率设置不匹配。全量微调时峰值学习率是 3e-5LoRA 里因为可训练参数少学习率得调大一些才能充分更新改成 1e-4 之后精度追平了全量微调。这个现象说明改变训练策略时第一时间怀疑的是参数没对齐而不是方法本身不行。4.2 精度对比只盯 ACC导致一次优化方案的误判还有一次做类别不平衡任务的量化评估模型总体 ACC 下降了 0.3 个点大家都觉得可以接受。但拆到每个类别看样本数最少的那一类 F1 从 82 直接掉到 71。这类问题在整体指标上不明显却在真实线上场景里会被放大。从这以后Model-Optimizer 的评测报告强制按类别输出指标并且专门标注尾部类别的退化程度。任何优化如果导致尾部类别大幅回撤哪怕整体指标很漂亮也不允许直接上线。4.3 剪枝最优组合不等于逐层独立的贪心结果剪枝判断时最容易犯的错是一层一层独立找最优然后把每层最优组合起来。实操时会发现不同层的重要度是相互关联的前面层剪掉的东西可能后面层靠冗余还能兜底也可能前面层一旦剪多了后面层想补都补不回来。我在 6 层 BERT 上跑了全局剪枝搜索把各层剪枝比例当成一组超参数用少量步数做组合评估效果比逐层贪心好了不少。代价是实验次数变多所以 Model-Optimizer 里默认先做一层粗搜索找大致范围再做局部精调。5. Model-Optimizer 的落地工作流与留给团队的建议到这里Model-Optimizer 的功能和踩坑都讲完了。下面说下最终沉淀的工作流一个可以复现、可以交给他人的完整流程。5.1 一条命令跑通的优化流水线整个流程分六步每一步都有独立脚本但也有一个总入口传一个配置文件就能顺序执行python run_optimize.py --config configs/text_cls.yaml配置文件里主要包含模型路径、任务类型、优化器策略、剪枝比例、量化方案、评估集路径、延迟测试集路径。流程内部顺序是基线评估记录原始模型的精度、体积、延迟。训练策略优化按任务类型套用优化器推荐配置输出收敛曲线对比。模型压缩执行蒸馏或剪枝生成候选模型。精度验证按类别输出指标。量化PTQ 跑完检查掉点必要时启用 QAT。引擎转换与延迟测试导出 ONNX 或 TensorRT端到端计时。这套流程最核心的设计是每一步都会留一份中间产物后面发现问题随时能回溯到上一步而不是从头再来。5.2 在线上模型上的收益数字用这套流程优化一个中文意图识别服务初始模型是 BERT-baseFP32 推理平均延迟 22ms显存占用 1.4GB。走完全流程后蒸馏到 4 层模型然后 INT8 量化再切到 TensorRT FP16最终平均延迟 6ms显存占用 180MB。精度方面整体 ACC 从 91.2% 降到了 90.1%下降了 1.1 个点但尾部类别的 F1 退化控制在了 2 个点以内业务方可以接受。这个过程耗掉的开发时间是两个人两周其中一半时间花在排查各个工具之间的兼容性上。如果你打算在团队里推广这套流程建议留出至少一周的磨合周期。5.3 给普通团队的选型建议根据我的经验普通团队做模型优化不用一上来就全流程铺开。如果你的模型还在训练阶段优先改训练策略优化器、学习率调度、AMP这些是投入产出比最高的地方。如果模型已经部署先做推理引擎切换再做蒸馏压缩。如果模型对精度特别敏感QAT 比 PTQ 更值得投入。另外建议从一开始就建一个统一的评测基线所有优化都基于同一份数据、同一个脚本去评估。模型优化本质上是交易用一部分精度、一部分时间换显存、延迟和成本。没有基线交易就没有参照物所有数字都会变成安慰自己的报表。这套流程我现在每个新项目都在用基本养成了条件反射拿到新模型先跑一遍 Model-Optimizer 的 profile 脚本看清楚训练曲线和部署瓶颈再决定动哪一块。很多团队花大力气改网络结构之前其实先用这套流程压一遍往往收获比想象中大得多。
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进