1. 先理清优化链路为什么是 ONNX 导出、TensorRT 优化、服务化这三步Transformer 模型的部署实战我最近做得比较频繁。很多团队的模型还在 PyTorch 训练阶段线上服务却已经要求更可控的延迟、更高的吞吐和更稳的显存占用于是 ONNX 导出、TensorRT 优化、服务化这一段链路就成了最常见也最省事的部署方案。这篇文章是我从实际踩坑角度出发把这三段里最容易卡住的地方拆开讲一遍适合那些已经跑通训练准备把模型从离线环境搬到 GPU 推理服务的同学参考。训练和推理的关注点其实差别很大。训练阶段你关心梯度能不能回传、显存够不够塞一个 batch追求的是灵活性推理阶段你关心的是单个请求延迟、吞吐、首 token 时间和显存峰值追求的是确定性和极致执行效率。PyTorch 的 eager 模式把每个算子拆成一个个独立的 CUDA kernel 去调度灵活归灵活但 kernel 之间有空隙而且很多连续操作可以融合却不会自动融合。把模型导出成 ONNX 再交给 TensorRT本质上就是把这层灵活让渡给一个为推理而生的运行时让它去做图优化、算子融合和精度调优。为什么中间要插一层 ONNX因为 TensorRT 的直接入口现在虽然也支持 PyTorch 导出的其他中间格式但 ONNX 在 Transformer 模型上的算子覆盖最全导出生态也最成熟它是训练框架和推理引擎之间一个相对中立的中间层。实际操作中ONNX 文件更像是模型的“图纸”它把计算图固定下来把动态 shape 暴露出来后续不管是用 ONNX Runtime 还是 TensorRT 读取都不再依赖原始训练框架。服务化这一步通常是被低估的。引擎构建完成只是拿到了一个能跑的二进制真正上线还要处理请求排队、动态 batch、并发实例、模型热更新这些问题。我见过不少团队在导出优化上花了两周最后却因为服务端配置不对导致 GPU 利用率只有百分之十几这显然是把部署实战理解窄了。所以这篇内容的路线很清楚先给出一个可以复现的 Transformer 编码器模型作为例子然后依次走 ONNX 导出、TensorRT 引擎构建、服务化挂载这三段每段我都会把重点参数和踩过的坑点写出来而不是只给结论。2. 环境版本矩阵与基线准备先量化你现在到底有多慢动手之前先定版本因为这套链路里 CUDA、PyTorch、ONNX Runtime、TensorRT 这四个组件的版本互相咬得很紧。版本不匹配最常见的现象是ONNX opset 版本支持不到位、TensorRT 解析器拒绝某个算子、导出时能跑但加载引擎时崩溃。我建议直接按下面这张表对齐这是我在实际项目中反复确认过的组合。组件推荐组合 A推荐组合 B备注操作系统Ubuntu 20.04Ubuntu 22.04最好是干净环境避免桌面套件干扰显存CUDA toolkit11.812.1TensorRT 安装包要匹配 CUDA 主版本Python3.8 / 3.93.10 / 3.11PyTorch 2.x 对 3.11 支持更好PyTorch1.132.1导出 API 大体一致2.x 的 dynamo 暂时不用ONNX1.131.15注意 opset 支持和 shape 推断能力TensorRT8.5.x10.x8.5 更稳定10.x 对新算子覆盖更全版本确定以后先做一个最朴素的基线压测。我去年代理的一个服务模型是一个 12 层 Transformer 编码器hidden size 768用来做文本分类。当时服务端直接调 PyTorch 的 GPU eager 推理单条 64 token 的输入在 A10 上平均要 9 到 12 毫秒看起来好像还行但一压并发就露馅8 路并发时 GPU 利用率忽高忽低p99 延迟能冲到 35 毫秒以上而且显存里有大量临时张量来不及释放。这个基线一定要用和线上一致的数据来测不要拿 GPU 的nvidia-smi里的显存占用来脑补。我当时写了一个很简单的脚本统计“请求进入模型到返回结果”的纯模型耗时不包含网络开销和预处理。压测时用 64 并发线程打 10 分钟记录平均延迟、p95、p99 和吞吐这部分数据后面做优化对比时非常关键。没有这一步后面 TensorRT 到底快了多少就全靠“感觉”而“感觉”在部署优化里是最不靠谱的东西。还需要准备一个相对干净的验证集后面做 FP16 或 INT8 精度对比时要用。验证集不需要很大几百条能覆盖不同序列长度的样本就够重点是比较导出前后的输出差异是否在可接受范围内。3. ONNX 导出动态 shape、opset 和算子兼容性的完整通关指南3.1 导出的最小可用模板我以一个自定义 Transformer 编码器为例。这个模型结构不复杂embedding 层加 12 层 TransformerEncoderLayer最后接一个分类头。代码结构刻意保持简洁方便你替换成自己的模型import torch from torch import nn class TextEncoder(nn.Module): def __init__(self, vocab_size30522, d_model768, nhead12, num_layers12, num_labels2): super().__init__() self.embed nn.Embedding(vocab_size, d_model) self.encoder nn.TransformerEncoder( nn.TransformerEncoderLayer( d_modeld_model, nheadnhead, batch_firstTrue, dropout0.0, ), num_layersnum_layers, ) self.classifier nn.Linear(d_model, num_labels) def forward(self, input_ids, attention_mask): x self.embed(input_ids) # attention_mask: True 表示需要被 mask 的位置这里直接用 bool 型 x self.encoder(x, src_key_padding_maskattention_mask) cls x.mean(dim1) return self.classifier(cls) model TextEncoder().eval()导出时最关键的是 dynamic_axes 配置。如果只固定 batch1 导出后续服务化会非常被动因为动态 batch 是提升吞吐的核心手段之一。还要把序列长度维度也设为动态因为线上请求不可能每一条都是固定长度例如刚上线时大家习惯性 pad 到训练时的最大长度结果短文本也占满 512 token 的计算量这在成本上是非常浪费的。dummy_ids torch.zeros((1, 64), dtypetorch.long) dummy_mask torch.zeros((1, 64), dtypetorch.bool) torch.onnx.export( model, (dummy_ids, dummy_mask), text_encoder.onnx, input_names[input_ids, attention_mask], output_names[logits], dynamic_axes{ input_ids: {0: batch, 1: seq}, attention_mask: {0: batch, 1: seq}, logits: {0: batch}, }, opset_version17, do_constant_foldingTrue, ) print(exported)opset_version 我直接用了 17这是 ONNX Runtime 和 TensorRT 现在都兼容得比较好的档位。太老的 opset 有些新算子不支持太高像 18 或 19 在部分 TensorRT 版本里解析会有问题。你如果用的是 TensorRT 8.x 而不是 10.x建议先查一下对应版本的算子支持表但通常 17 是个稳妥选择。3.2 导出后必须做的三项验证ONNX 导出成功不代表导出无误文件能生成只是第一步。我每次导出后固定做三件套检查第一是onnx.checker.check_model检查图结构是否合法这一步能发现明显的节点连接问题。import onnx onnx_model onnx.load(text_encoder.onnx) onnx.checker.check_model(onnx_model) print(onnx check passed)第二是用 ONNX Runtime 跑一遍推理拿输出和原始 PyTorch 模型对比。对比时要选好数据不要用全 0 输入因为全 0 会让很多数值差异被掩盖。我会生成随机 token 序列然后分别跑 PyTorch 和 ONNX Runtime比较 logits 的绝对差和相对差。通常允许的误差范围在 1e-3 以内FP32 导出一般能到 1e-5 以下如果超过 1e-2说明某个算子被错误替换了。第三是用 TensorRT 自带的 trtexec 做一次快速构建测试看整个图表能不能被解析。这一步有时候会直接暴露某个算子不支持等到服务化阶段才发现的话排查成本会高很多。3.3 三种最容易翻车的算子场景Transformer 导出最容易翻车的点不在标准 attention而在一些辅助操作上。最常见的是三种一是torch.where和 boolean mask 的组合。ONNX 里Where算子对三个输入的 shape 有严格 broadcast 要求如果你的 mask 是从attention_mask[:, :, None]之类的方式扩展出来的导出时 shape 推断偶尔会产生额外Unsqueeze节点倒不影响正确性但会在 TensorRT 图优化时形成一个小尾巴。建议在模型 forward 里手动把 shape 理清楚不要完全依赖隐式 broadcast。二是自定义的gelu实现。训练时代码里如果用了nn.GELU(approximatetanh)导出时会有多个数学算子拼成一个误差很小的公式TensorRT 对这类子图做融合时一般没问题但如果你的gelu是手写的0.5 * x * (1 torch.erf(x / sqrt(2)))则某些版本的 TensorRT 对 erf 的优化不如对 tanh 形式好。更常见的做法是直接使用 PyTorch 自带的 GELU 激活让它映射到标准 ONNX 子图。三是past_key_values相关的逻辑。如果是做生成模型比如 Transformer 解码器导出时要把use_cache相关的缓存逻辑固定下来否则 ONNX 图里会出现大量 if-else 分支。 TensorRT 处理动态 shape 时对这种控制流容忍度较低。我的建议是先导出一个不带缓存的最小推理图跑通流程再考虑带缓存的版本两边分开做不要一开始就追求一步到位。4. TensorRT 引擎构建精度、显存与动态 shape 的三方权衡4.1 trtexec 构建引擎的参数拆解拿到 ONNX 文件后最直接的引擎构建方式是用 trtexec。命令如下trtexec \ --onnxtext_encoder.onnx \ --saveEnginetext_encoder.engine \ --fp16 \ --minShapesinput_ids:1x16,attention_mask:1x16 \ --optShapesinput_ids:8x64,attention_mask:8x64 \ --maxShapesinput_ids:16x256,attention_mask:16x256 \ --workspace2048这里--fp16表示允许 TensorRT 把算子改写为 FP16 精度。--minShapes、--optShapes、--maxShapes三个参数是动态 shape profile缺一不可。如果你只导出了动态 shape却没有在构建引擎时指定这三个范围TensorRT 会默认用导出的 dummy shape 作为固定 shape之后一旦输入长度超出范围就会直接报错。--workspace控制构建期用来做 autotune 和层选择的最大显存单位是 MB。注意这不是运行时显存上限设得越大不代表推理时占用越大只是让 TensorRT 有更多空间去尝试不同 kernel 组合。我通常设 2048 或 4096如果服务器显存比较小可以把值调低但构建时间会变长且某些需要大 workspace 的融合策略可能被跳过。构建完引擎后我会先用 trtexec 自带的口径测一下吞吐和延迟比如加--duration60让它持续压测 60 秒。但这里有个容易误读的点trtexec 使用的是模拟输入不经过真实预处理和服务端框架所以它适合判断“引擎本身能跑多快”不代表“整个服务能跑到这个数”。4.2 精度模式选择FP32、FP16 还是 INT8我们实际对比过三种精度模式在相同动态 shape 下的表现数据大概是这样的同样的模型在 A10 上模式平均延迟相对 FP32 速度提升p99 抖动精度影响FP329.2 ms1.0x较小基准FP164.6 ms2.0x中等通常可忽略个别层有微小漂移INT83.1 ms3.0x较大需要校准集否则分类任务可能掉 1-2 个点FP16 几乎是白拿的性能。Transformer 结构里的 matmul 和 layernorm 对 FP16 精度足够宽容很多框架甚至默认就在用 FP16 做推理。真正需要小心的是LayerNorm它内部有均值、方差和小数加法如果整个层都被强制切成 FP16某些极端长尾数据上可能出现精度漂移。TensorRT 在做 FP16 图优化时会保留部分敏感算子为 FP32所以实测下来编码器模型几乎没有差别。INT8 则需谨慎。很多人一上来就冲 INT8觉得是极限优化但 INT8 需要对 activation 做量化校准校准集选不好会让线上效果显著下降。我自己的经验是先上 FP16确认吞吐和延迟已经达到目标再评估要不要继续压到 INT8。如果服务的瓶颈是在 GPU 利用率而不是延迟FP16 往往已经足够。4.3 动态 shape profile 的坑一次线上长序列事故有一个坑我想单独拿出来说。当时构建引擎时把optShapes设成了8x64也就是优化目标是为 batch 8、序列长度 64 调的优maxShapes设成了16x256。试运行都很正常结果上线后某个调用方传进来一条 260 token 的请求引擎直接拒绝报错信息是“input seq exceeds maximum profile shape”。这个报错是在线才暴露的因为测试集里最长只有 200 token而 max 卡在 256。这中间其实牵扯到 TensorRT 的一个约束动态 shape 不是无级变速引擎会按 profile 选择最优 kernel输入 shape 越贴近 optShapes执行效率越高一旦接近边界性能会明显下降甚至直接用不了。所以设置 maxShapes 时不要严格按照训练阶段看到的最大长度来要预留 20% 到 30% 的缓冲尤其是线上数据存在长尾的 NLP 服务。另外如果一个引擎要在同一个服务里服务不同 batch 和序列长度的组合通常需要配置多个 profile。比如一个 profile 给高吞吐的短请求另一个给低并发的长请求运行时根据实际 shape 切换 profile。配置多个 profile 会显著增加构建时间但能避免在极端 shape 下选到很差劲的 kernel。5. 服务化落地把 TensorRT 引擎挂上推理服务并让它高效跑起来5.1 自建轻量服务还是使用专用推理服务框架引擎构建好之后服务化有两种路线。一种是自己写一个模型加载器用 FastAPI 或 gRPC 暴露接口内部调用 TensorRT Python API 做 inference。这条路线适合原型验证、内部工具或者模型数量很少的场合。优点是代码透明依赖简单想加什么预逻辑都很自由缺点是并发调度、动态 batch、GPU 显存复用这些都得自己实现一旦请求量上来你会发现自己写的排队逻辑远不如现成框架健壮。另一种是使用专门做推理服务的框架比如 Triton Inference Server。它天然支持 TensorRT 引擎后端模型仓库管理、动态 batching、并发实例配置、指标上报都是现成的。我实际用下来生产环境强烈建议选这条路线不是因为框架复杂就一定好而是它把你线上推理最容易出问题的部分都收敛掉了。5.2 模型仓库目录与配置模型仓库的结构非常简单把引擎文件放到一个目录下再写一个config.pbtxt就够了。model_repository/ └── text_encoder/ ├── config.pbtxt └── 1/ └── text_encoder.engineconfig.pbtxt示例name: text_encoder platform: tensorrt_plan max_batch_size: 16 input [ { name: input_ids data_type: TYPE_INT64 dims: [ -1 ] }, { name: attention_mask data_type: TYPE_BOOL dims: [ -1 ] } ] output [ { name: logits data_type: TYPE_FP32 dims: [ -1, 2 ] } ] dynamic_batching { preferred_batch_size: [ 4, 8 ] max_queue_delay_microseconds: 100 } instance_group [ { count: 1 kind: KIND_GPU } ]这里的dims: [-1]表示序列维度是动态的batch 维度交给max_batch_size和动态 batching 管理。有一点容易混淆max_batch_size是服务请求级别的 batch不是 TensorRT 引擎内部的 batch。Triton 会收集多个请求组成一个推理 batch再喂给引擎。dynamic_batching里的preferred_batch_size我建议设置成引擎 optShapes 的 batch 值。比如引擎 optShapes 是 8x64那这里就填 8这样 Triton 会尽量凑满 8 个请求再发一次推理性能最稳定。max_queue_delay_microseconds是个权衡参数如果请求来得稀疏为了等满 8 个请求而让用户等待是不划算的所以要设一个最大延迟一般 50 到 200 微秒比较合理时间太长会让 p99 明显升高。5.3 动态 batching 的收益验证我们之前在同一个服务上做过开关对比。关闭动态 batching 时并发 16 路的吞吐大约在 280 次/秒开启后同样并发下吞吐涨到 520 次/秒左右而 p99 延迟只多了 1 到 2 毫秒。原因很简单GPU 的矩阵计算擅长处理大 batch单条请求的 matmul 其实非常浪费算力把多个独立请求拼成一个 batch 后矩阵乘法的形状从 1x64 变成 8x64计算效率没有线性下降但执行次数少了很多。不过动态 batching 不是万能药。它只对计算密集型模型有效如果你的模型里还有大量 Padding 和 Mask 操作不同长度的请求拼在一起会造成很多无效计算。所以服务端要做两个配合一是请求进来后按序列长度做分桶尽量让相似长度的请求进入同一个 batch二是把 padding 逻辑放在预处理里而不是等到了模型内部再做 mask。这一点我在实际项目里吃过亏。最初把变长请求直接丢给引擎一个 16 token 的请求和一个 256 token 的请求被拼进同一 batch短的那个浪费了大量计算在尾部 padding 上最终效果还不如固定长度 padding 后再跑。后来我改成按长度分桶调度整体吞吐提升了约 30%。6. 端到端压测、排障与几处容易翻车的细节6.1 第一次推理为什么那么慢冷启动与 warmupTensorRT 引擎加载后第一次推理往往非常慢有时一个 12 层的编码器要跑 200 到 400 毫秒而之后只有 4 到 5 毫秒。这不是引擎不稳定而是 CUDA context 初始化、显存分配、cuDNN 或 cuBLAS 的 lazy loading 都集中在第一次调用上。如果服务端一启动就接流量前几个请求的延迟一定会超时。解决办法是在服务启动流程里加一个预热阶段加载引擎后先跑 10 到 20 次 dummy 输入跑完再对外宣告服务可用。如果用的是 Triton可以直接在请求进来前发一个 health check 信号确保模型已经 ready。这个细节看起来小线上事故往往就是从这里开始的。我在第一次部署时没注意到这一点结果服务重启后健康检查虽然显示通过但第一批在线请求的 p99 全部超时排查了半天才发现是冷启动造成的。6.2 显存碎片与多实例配置TensorRT 引擎加载时经常一次性申请一大块显存如果服务里还有别的模型实例显存碎片会导致新模型分配不到足够空间。尤其是部署多个模型或者多版本模型做平滑切换时旧引擎释放的显存不一定能完全被新引擎复用。常用的处理办法是给 Triton 的instance_group里的count配 1同一时间只保留一个推理实例。滚动更新时先启动新版本模型等它加载完成后再释放旧版本这样显存峰值可控。如果你的 GPU 资源足够也可以把count设为 2得到更高的并发能力但要监控显存使用否则 OOM 会导致整卡服务崩溃。还有一点TensorRT 构建时的 workspace 参数和运行时显存预留有关。如果服务报告 OOM第一件事不是去加显存而是检查引擎是不是用了--workspace过大或者在 config 里给模型分配了不必要的gpu_memory_fraction。6.3 验证精度的完整闭环上线之前一定不要把“能跑”当作“跑对”。我见过不止一次ONNX 导出后模型精度在单个样本上对得上入库的大批量验证集却掉了一两个点最后定位到某个 Pad 相关的操作在导图时被错误常量折叠。我的验证流程是这样先选 300 到 500 条线上真实样本包含不同长度分布然后分别记录 PyTorch FP32、ONNX Runtime FP32、TensorRT FP16 三者的输出最后计算每个输出维度的平均绝对误差和最大绝对误差并跑一遍业务相关的指标。如果业务指标是分类准确率就直接比准确率如果是排序类指标就比 AUC 或 NDCG 的相对变化。只有当这些指标在误差范围内才允许切换流量的一个百分比做灰度。灰度这一步也别省。我会先把新服务切到 5% 流量跑一天观察线上指标和延迟再逐步放量到 30%、100%。不要因为离线数据好看就直接全量切在线请求的分布永远比离线验证集更野。6.4 上线后最容易被忽视的监控项服务化完成不代表事情结束。部署上线后至少要盯四个监控指标GPU 利用率、显存占用、请求 p99 延迟、推理 batch 平均大小。GPU 利用率如果长期低于 50%说明动态 batching 没生效或者模型本身计算量不够大。显存占用持续上涨且不回落大概率是显存碎片或某个实例泄漏。p99 延迟如果周期性出现尖峰要检查是不是服务端在做定期模型 reload 或者 CUDA context 释放。推理 batch 平均大小如果一直小于 preferred_batch_size很可能是请求涌入不均匀需要调max_queue_delay_microseconds或者在上游做批次缓存。我这边的经验是第一次上 TensorRT 服务通常要花两三天做压测和调参。不要指望引擎构建完一切就自动达到最优。尤其动态 batching 和 shape profile 这两块不同业务的数据分布不同参数必须跟着实际流量打磨。如果你也是从训练模型开始一步步做部署我建议第一次不要上 INT8先跑通 FP32 或 FP16 的完整链路把 ONNX 导出、TensorRT 构建、服务化挂载这几步的坑都摸一遍再回头考虑极致优化。毕竟部署实战的精髓不是把某一个环节做到最快而是让整条链路在长期运行中保持稳定、可控、可回滚。最后再分享一个我能给的最高优先级建议引擎构建、服务编排、灰度发布这些过程能自动化就自动化因为模型会更新框架会升级线上的坑会以你完全想不到的方式重演。