从OOM到Model-Optimizer：大模型部署量化剪枝蒸馏实战

上个月某个周五晚上，我盯着监控面板上刷出来的 CUDA OOM 报错，脑子只有一个念头：这个 Model-Optimizer 项目再不做完，下周一整周的线上扩容申请估计都要被砍掉。那个晚上我一边翻着各种量化脚本，一边在几个项目仓库里来回找之前写过的校准代码，最后实在忍不住，把散落在不同 commit 里的优化逻辑全部抽出来，重写成了一个统一的工具链——也就是后来我一直在用的 Model-Optimizer。

这个名字听起来挺大的，其实做的事情很聚焦：把训练好的模型从"学术指标很好"变成"生产环境能用"。包括参数量压缩、推理加速、精度对齐、后端导出，以及上线前的验证流程。如果你正在做 LLM 或 Transformer 类模型的部署，被显存爆掉、首 token 延迟超标、并发上不去这些问题困扰，那这篇从事故讲到实操的完整记录，应该能帮你少走很多弯路。

1. 一次线上OOM让我决定把优化流程重做一遍

1.1 事故现场还原

事情要从一次线上事故说起。我们当时把一个 13B 左右的对话模型切到生产，8 张 A100 撑推理，本来压测数据看着还行，结果业务真实流量一上来，不到二十分钟就触发了显存 OOM，容器直接重启，线上告警刷了一整屏。

当时查下来原因很朴素：输入长度一长，注意力机制的内存峰值远超压测预设；加上多路并发一起进来，显存碎片一涨，OOM 就在所难免。这种事情在推理优化领域几乎是"新人必经课"，但对当时的我来说，最头疼的不是 OOM 本身，而是我找不出一个快速的优化手段来缓解它——手头有一堆零散的量化脚本，但没人敢在生产环境直接跑。

1.2 传统"临时脚本式"优化流程的问题

我原来是这么干的：把 PyTorch 模型转成 ONNX，然后拿一个自己写过几次的校准脚本做 INT8 量化，再手动指定几个敏感层回退 FP16，跑一遍测试集看下指标没崩就部署。这套流程最要命的地方有三个。

第一，校准逻辑和代码强耦合。每次换模型、换数据分布，都要去改脚本里的阈值、层名、校准算法，出了 bug 得靠肉眼在代码里找。第二，不同项目之间完全不可复用。上个项目用的量化参数、剪枝比例、蒸馏配置，换个模型基本全部失效，等于每次从零开始。第三，精度评估不闭环。优化完了只看一个总准确率，线上延迟降了没、降了多少、有没有隐性劣化，全靠后来压测才发现。

1.3 Model-Optimizer 的目标定位

那次事故之后我复盘了一下，发现自己缺的不是某个量化算法，而是一套能把模型的"优化过程"当作工程问题来管理的框架。所以 Model-Optimizer 在设计之初就定了三个原则：

• 所有的优化策略都是可配置的插件，模型本身不感知优化过程。
• 中间表示统一走 ONNX，PyTorch、TensorFlow 甚至 Paddle 训练出来的模型都能进来。
• 优化结束后自动产出一份精度、延迟、显存对比报告，而不是让人去手动拼结果。

这套工具脱胎于实际问题，所以下面的每一条经验，都是那段时间真正踩过的坑。

2. 架构设计：优化策略怎么才能不变成另一堆胶水代码

2.1 中间表示层：为什么ONNX是绕不开的枢纽

Model-Optimizer 的第一层是模型接入层，我选择所有模型先进 ONNX，再走后续的量化、剪枝、蒸馏。原因很简单：ONNX 是一个稳定的中间图表示，几乎所有训练框架都有导出能力，几乎所有推理后端都有加载能力。一旦模型进了 ONNX 这个"通用语言"，优化策略就有统一的操作对象，不用针对每个训练框架写一套。

用生活中的类比来说，ONNX 就像是一个标准的集装箱规格：不管你的货物（模型）原本是装在 PyTorch 的卡车里还是 TensorFlow 的卡车上，只要统一换成 ONNX 集装箱，后面的吊机（量化工具）、卡车（推理引擎）、仓库（部署环境）就都能无缝对接。

在导出阶段有几点值得注意。比如torch.onnx.export的时候要把动态轴设对，否则序列长度一变，导出后的模型就废了。还有一个特别容易被忽略的参数是opset_version，如果设得太低，后面做算子级量化时很多 fuse 规则根本不会生效。我的经验是，Transformer 类模型建议至少指定 opset 15 以上，才能覆盖大多数量化场景所需的算子。

2.2 策略注册机制：把量化、剪枝、蒸馏做成可插拔能力

Model-Optimizer 的核心是一个策略注册表。每一个优化手段，比如PTQQuantizer、StructuredPruner、DistillationRunner，都实现同一个接口，然后注册到框架里。配置走 YAML 文件，不用改一行代码。

optimization: pipeline: - name: quantize type: ptq algorithm: entropy calib_samples: 512 operators: - matmul - layernorm fallback_nodes: - "/encoder/layer.0/attention/MatMul" - name: prune type: structured target_ratio: 0.2 importance: attention_head

这个设计的收益是：当我把这个工具给团队里其他人用时，他们不需要读懂量化算法的源码，只需要知道"我打算先量化再做轻量剪枝"，然后改几行 YAML 就行。而且因为每个策略都是独立的，某个策略出了诡异问题，可以直接在流水线里临时摘除，排查速度快很多。

2.3 一次完整的优化流水线长什么样

我跑得最多的一条流水线是：加载 ONNX → 数据校准 → PTQ 量化 → 结构化剪枝（小比例） → 精度评估 → 导出到 ONNX Runtime / TensorRT 格式 → 生成对比报告。

每一步之间传递的是同一个 ONNX 图对象，只是在图的属性里附加了一些元数据，比如量化参数、剪枝 mask、评估结果。这样的设计保证整个流水线可以随时从中间某一环接着跑，不会因为前面重跑一遍而丢失后面已经调好的参数。这个特性在做敏感层回退的时候特别有用：发现某一层量化后精度劣化，直接标记为 fallback，重新跑后面环节就行。

3. 量化落地的Real Talk：校准数据比算法选择更重要

3.1 PTQ为主、QAT兜底：我的选型逻辑

量化是 Model-Optimizer 里用得最多的策略，也是大家问得最多的部分。我的选型逻辑很简单：能 PTQ 就不上 QAT。PTQ（Post-Training Quantization）不需要重新训练模型，几分钟就能拿到一个 INT8 版本；QAT（Quantization-Aware Training）虽然精度更有保障，但要准备训练脚本、调学习率、重跑训练流程，周期至少以天为单位。

但在实际使用中，PTQ 对校准数据的要求相当高，没有好的校准集，再先进的算法也救不了精度。这条是我反复踩过的坑，值得单独展开讲。

3.2 校准集的构造：切片训练集是最省事但最容易翻车的做法

很多教程会告诉你"从训练集里随机抽 500 条当校准集"，这句话害了不少人。随机切片的训练集分布和线上真实输入经常差异很大，比如线上大多是长文本、多轮对话、带有格式噪音的数据，而训练集里可能以短文本居多。用这种校准集去统计激活值的动态范围，得出的量化 clip 点就会失真，表现为量化后的模型在线上偶发性输出劣化。

我在 Model-Optimizer 里定义了一套校准集的构造规则，核心就一句话：校准集要贴近线上推理时的输入分布，而不是贴近训练时的分布。具体做法是：

• 从线上日志里采一段时间的真实请求，清洗后作为校准数据；
• 如果拿不到线上数据，就按线上场景的输入长度分布、格式特征、领域分布去构造模拟数据；
• 校准集至少要涵盖所有可能出现的输入长度范围，尤其是长序列端。

3.3 算子级配置与敏感层回退

ONNX 量化的粒度分为 per-tensor 和 per-channel。我的一般策略是：MatMul、Gemm 这类权重占大头的算子用 per-channel，激活用 per-tensor；涉及 LayerNorm 和 Softmax 这类对数值范围敏感的算子先跳过量化，或者单独配置更高精度的 clip 方式。

实践中我发现一个规律：Transformer 的 Embedding 层、最后的 LM Head、残差连接后面的第一个 MatMul，往往是量化精度的重灾区。这些层要么输入动态范围特别宽，要么直接决定了输出分布。Model-Optimizer 的逻辑很简单：先全部量化，跑一遍评估；对精度掉得最多的前几个节点，自动标记为回退 FP16，然后重新评估。这样一个迭代式的"量化-评估-回退"循环，通常两三轮就能收敛到一个精度可接受且压缩率最大的配置。

3.4 精度验证：不能只看整体指标

量化后验证精度，最容易犯的错误是只看一个整体的 accuracy 或者 loss。我在事故中吃过亏：整体困惑度几乎没变，但具体的对话任务里，模型开始在某些专有名词上频繁出错。原因是这些 token 在整体指标里的权重太小，掩盖了量化引入的局部劣化。

所以 Model-Optimizer 的评估模块会同时输出整体指标和拆分指标，比如按输入长度分段、按 token 频率分桶、按语义任务分类的细粒度对比。把量化前后的偏差摊开来看，才能发现"延迟降了 30%，但长文本场景的生成质量掉了 2 个点"这类隐形风险。

4. 剪枝和蒸馏：先砍结构，再用小模型把精度找回来

4.1 结构化剪枝 vs 非结构化剪枝

说完量化，再讲剪枝。剪枝分两类：非结构化剪枝把权重张量里不重要的单个元素置零，模型体积能变小，但因为非零元素分布不规则，GPU 上的稀疏矩阵运算效率并不高，推理加速有限；结构化剪枝则是整行整列地移除权重，比如删掉某个注意力头或 FFN 的某个中间维度，结构规整之后，模型可以真正变小变快。

对于部署场景，我的建议很明确：优先结构化剪枝。虽然精度损失通常比非结构化大一些，但对显存和延迟的实际收益更直接，而且后续可以用蒸馏把精度拉回来，两条路配合着走。

4.2 注意力头和FFN维度的重要性评估

剪哪些结构不是靠拍脑袋，而是靠重要性评估。我用的一个简单有效的方法是：给每个注意力头算一个"输出方差贡献"和"梯度敏感度"的加权分数。输出方差贡献低的头，说明它对不同输入的响应区分度低，剪掉影响小；梯度敏感度低的头，说明它在下游任务中不被需要。

代码层面大概是这样的逻辑：

def estimate_importance(model, calib_loader): scores = {} for batch in calib_loader: hidden_states = hook_extract_attention_outputs(model, batch) for head_id, head_out in hidden_states.items(): var_score = torch.var(head_out, dim=0).mean() grad_score = compute_grad_norm(model, head_id, batch) scores[head_id] = var_score * grad_score return scores

算出分数后，按从低到高逐步剪除，每剪掉一批头就做一次评估，看指标下降曲线。通常我会控制在剪掉 10%~20% 的注意力头和 FFN 维度这个区间，再往上精度衰减会非常快。

4.3 蒸馏参数怎么调

剪完结构的模型，如果要恢复到接近原始精度，就要做蒸馏。蒸馏的核心是让小模型（学生）去模仿大模型（教师）的输出分布，训练目标一般是硬标签的交叉熵加上软标签的 KL 散度。

我们常用的配置是：蒸馏温度调到 4 左右，硬标签损失权重 0.2，软标签 KL 损失权重 0.8。这个比例不是固定不变的，如果发现学生的输出过于"平滑"、缺乏自信，就把硬标签权重适当调高；如果学生的泛化能力不够，就提高蒸馏温度，让教师输出的分布更"软"，携带更多类间关系信息。

学习率调度也有讲究。蒸馏阶段普遍再用 AdamW + warmup 加 cosine 衰减，峰值学习率比正常预训练低一个量级，比如 1e-5 起步。因为这时候不是在学新知识，而是在向教师模型的对齐，步子太大会直接冲过头。

4.4 剪枝蒸馏的一个翻车记录

有一次我自信满满地跑完"剪 30% + 蒸馏"之后，发现单测指标确实没掉，但生成文本开始出现重复片段。排查了很久才意识到，蒸馏的软标签会掩盖模型对低置信 token 的判断能力，等于让学生只学会了"大多数人怎么回答"，没学会"遇到模糊问题时该多谨慎"。

这个问题的解法是在蒸馏损失里额外加一个对低置信 token 的重加权项，让学生模型在那些教师也不确定的 token 上，不需要强行学一个概率分布，而是可以保留一定的"不确定性"。从那以后，我们所有蒸馏任务都会额外看一眼低置信区间的分布对齐情况，而不是只看平均 KL 散度。

5. 导出与后端对接：ONNX导出没你想象的那么"无损"

5.1 torch.onnx.export的动态轴设置

到这一步，优化后的 ONNX 图已经准备上线了。但很多人会在这里再翻一次车——导出的 ONNX 模型在本地跑没问题，一上线输入长度稍微一变就报 shape 不匹配。问题几乎都出在动态轴没设对。

torch.onnx.export里要明确指定哪些维度是动态的：

torch.onnx.export( model, dummy_input, "model_optimized.onnx", input_names=["input_ids", "attention_mask"], output_names=["logits"], dynamic_axes={ "input_ids": {0: "batch_size", 1: "seq_len"}, "attention_mask": {0: "batch_size", 1: "seq_len"}, "logits": {0: "batch_size", 1: "seq_len"}, }, opset_version=17, )

这里有个细节值得多说一句：如果模型内部有某些算子不支持动态 shape，即使你在 dynamic_axes 里声明了，导出时也可能悄悄被固定。所以我建议导出后做一个"形状探测器"，用多种不同长度的输入分别跑一遍导出的模型，确认所有动态维度都真实生效了再继续。

5.2 后端选择：ONNX Runtime还是TensorRT

Model-Optimizer 目前支持两个主流后端：ONNX Runtime 和 TensorRT。我的选型建议是：

从我的实测数据看，同一个优化后的模型，TensorRT 在延迟上通常能再压 15%~20%，但构建 engine 的时间明显更长，而且对模型中的某些算子的兼容性要求更苛刻。所以如果项目周期紧，先上 ONNX Runtime 往往是把模型跑起来的更优解。

5.3 端到端验证清单

导出到后端之后，我会按一份固定的清单做端到端验证，每一条都踩出过问题：

• 用至少 5 种不同长度的输入跑一遍，确认动态 shape 正常；
• 对比优化前后模型在相同输入上的输出差值，定位意外的数值漂移；
• 跑一遍真实线上回放数据，观察是否有超时、OOM、非法输出；
• 检查显存峰值，确认在多路并发下不会踩到显存上限；
• 记录 p50/p95/p99 延迟，确认优化不仅降低了平均延迟，还把长尾延迟也带下来了。

6. 上线前的数据回放和灰度回滚：最后一公里怎么守住

6.1 回放测试：用真实流量暴露优化痕迹

优化模型最终是要服务真实用户的，所以在上线前，我会把过去一周的线上请求数据和对应的历史输出做一个回放测试。回放的意思是把同样的输入喂给优化后的模型，比较它的输出与线上已服务模型输出的差异。

这个环节能发现两类问题。一类是"硬伤"：模型崩溃、输出格式错误、空响应。另一类是"软伤"：输出的语义发生了微妙变化，比如语气变冷、回答变短、拒绝类内容增多。硬伤靠监控能抓，软伤靠监控很难发现，必须有人逐条看回放结果。我的经验是，至少拉 200 条真实请求做人工抽检，并且集中在长输入、多轮会话、特殊符号这类容易暴露问题的样本上。

6.2 灰度策略

上线策略上，我一直坚持灰度。第一天先放 10% 流量，重点盯 p99 延迟和显存峰值；第二天如果稳定，放到 30%；再过一天没问题，才放开全量。灰度期间，新老两个版本并行跑，所有响应都做一次落库对比。

这样做的底气在于 Model-Optimizer 的配置即代码设计。每个优化版本都有独立的 YAML 配置和评测报告，回滚的时候只需要切换配置重新导出一份模型文件，不需要动任何代码逻辑。上线之后如果发现问题，我可以在一分钟内把流量切回老版本，最大限度降低对线上用户的影响。

6.3 配置即代码：让优化过程可复现

这个点我觉得值得单独说一句。在 Model-Optimizer 里，所有量化参数、剪枝比例、蒸馏配置、评估阈值，全部以 YAML 方式随代码库一起管理。这样做的直接好处是，任何人拉下来一个分支，只需要执行一条命令，就能复现某次优化的全过程，包括用哪份校准数据、哪个层回退 FP16、评估出来是多少分。

以前用零散脚本时经常出现一种乌龙：优化了一个效果很好的模型，但过了两周没人记得当时用的量化算法细节，想加一个 feature 却发现无从下手。配置即代码把这个问题根治了。

回到最初那次 OOM，现在再遇到类似问题，我整个排查和优化的流程大概是：先把模型导成 ONNX，用线上回放数据生成校准集，跑一遍 PTQ 量化，如果精度掉了再按敏感层回退和轻量剪枝的组合拳处理，最后按固定清单做端到端验证和灰度上线。整个过程从发现问题到拿到一个可上线的优化模型，基本能控制在一天内。跟之前那次手忙脚乱、连脚本都找不到的状况相比，好用了不知道多少倍。

Model-Optimizer 这个项目目前的定位就是一个解决"模型能用但部署不动"问题的工具套件。它不会替你选择训练什么模型，但能帮你把训好的模型压缩到可以上线、且上线后跑得稳的程度。优化这条路没有终点，模型、数据、业务需求都在变，但只要你把流程固化成可复现的工程能力，下一次遇到新问题时，你就不会从头开始踩坑了。