一、真实场景:凌晨三点线上的推荐模型崩了
2024年6月,我负责的推荐系统线上模型服务在晚高峰时段P99延迟飙到2.4秒,直接拖垮了上游调用方——客户端等不到推荐结果,导致首页加载超时。
架构很简单:Python Flask + PyTorch,四个GPU实例扛所有流量。模型本身不大(ResNet18变体),但问题出在推理路径太长:PyTorch的CPU预处理 → GPU推理 → CPU后处理,每步都有Python GIL参与,并发一高就排队。
当晚我做了个临时处理:把batch size从32减到8,延迟降下来了,但吞吐直接腰斩。这不是办法。
第二天我拿到了一周的全量流量日志,做了一次彻底的评估。最终方案是把模型部署链路整体重构:PyTorch → ONNX Runtime → TensorRT → Triton Serving。这篇文章记录整个迁移过程、实测数据和踩过的坑。
二、三套方案选型逻辑
先说明环境:所有测试在同一台设备上完成。
| 配置 | 规格 |
|---|---|
| GPU | NVIDIA A10 24GB |
| CPU | Intel Xeon Gold 6330 @ 2.0GHz (32核) |
| 显存 | 24GB |
| PyTorch | 2.1.0 + CUDA 11.8 |
| ONNX Runtime | 1.17.1 (GPU) |
| TensorRT | 8.6.1 + CUDA 11.8 |
| Triton Serving | 2.40.0(官方Docker镜像) |
2.1 对比矩阵
| 引擎 | 吞吐量(req/s) | P99延迟(ms) | 显存占用(MB) | 集成难度 |
|---|---|---|---|---|
| PyTorch (原方案) | 86 | 142 | 2,840 | - |
| ONNX Runtime GPU | 214 | 58 | 1,980 | 低 |
| TensorRT FP16 | 612 | 18.4 | 1,120 | 中 |
| TensorRT INT8 | 1,063 | 6.8 | 920 | 高 |
| Triton + TensorRT FP16 | 1,082 | 9.2 | 2,340 | 中 |
说明:Triton的吞吐更高是因为自带Dynamic Batching,把多个请求合并成一次推理。下面拿ONNX Runtime和TensorRT的完整代码分析。
三、方案一:ONNX Runtime迁移
ONNX Runtime的迁移成本最低,原来写PyTorch Forward的逻辑基本不用动。关键是把模型导出成ONNX格式。
3.1 PyTorch导出ONNX
# export_onnx.py
# PyTorch 2.1.0 | onnx 1.14.0 | onnxruntime-gpu 1.17.1
import torch
import torch.nn as nn
import onnx
from onnxruntime.tools.convert_onnx_models_to_ort import convert_onnx_models_to_ort
class RecommenderModel(nn.Module):
"""单层Transformer + 全连接输出 CTR 预估"""
def __init__(self, d_model=128, nhead=8, num_classes=1):
super().__init__()
self.embedding = nn.Embedding(10000, d_model)
self.transformer = nn.TransformerEncoder(
nn.TransformerEncoderLayer(d_model, nhead, batch_first=True),
num_layers=4
)
self.fc = nn.Linear(d_model, num_classes)
def forward(self, input_ids):
x = self.embedding(input_ids)
x = self.transformer(x)
x = x.mean(dim=1)
return torch.sigmoid(self.fc(x))
# 关键点1: 用真实输入shape导出,动态shape会损失性能
model = RecommenderModel().cuda().eval()
dummy_input = torch.randint(0, 9999, (8, 64)).cuda() # batch=8, seq_len=64
with torch.no_grad():
torch.onnx.export(
model,
dummy_input,
"recommender.onnx",
opset_version=17,
input_names=["input_ids"],
output_names=["ctr"],
# 关键点2: 这样定义动态轴,ONNX Runtime才能在batch纬度做优化
dynamic_axes={
"input_ids": {0: "batch_size"},
"ctr": {0: "batch_size"}
},
do_constant_folding=True,
)
print("导出完成")
# 校验ONNX模型的正确性
onnx_model = onnx.load("recommender.onnx")
onnx.checker.check_model(onnx_model)
print(f"Opset版本: {onnx_model.opset_import[0].version}")
3.2 ONNX Runtime GPU推理服务
# onnx_serve.py
# onnxruntime-gpu 1.17.1
# 依赖: pip install onnxruntime-gpu
import onnxruntime as ort
import numpy as np
import time
from concurrent.futures import ThreadPoolExecutor
# 关键点3: 使用GPU执行提供器,并且关闭CPU回退
# 用TensorrtExecutionProvider是ONNX Runtime调用TensorRT的方式,但建议直接原生TensorRT
so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
so.intra_op_num_threads = 8
so.inter_op_num_threads = 4
session = ort.InferenceSession(
"recommender.onnx",
sess_options=so,
providers=[
("CUDAExecutionProvider", {
"device_id": 0,
"arena_extend_strategy": "kNextPowerOfTwo",
"cudnn_conv_algo_search": "EXHAUSTIVE",
}),
],
)
# 预热: 第一次推理包含CUDA初始化,不计入性能统计
warmup_input = np.random.randint(0, 9999, (8, 64)).astype(np.int64)
for _ in range(10):
session.run(None, {"input_ids": warmup_input})
def infer(batch):
return session.run(None, {"input_ids": batch})[0]
# 并行推理入口
executor = ThreadPoolExecutor(max_workers=16)
# 进入FastAPI/Flask之后,直接调用executor.submit(infer, input_batch)即可
print("ONNX Runtime服务就绪")
# 用linux自带的性能工具测试时,把下面的压测协议打开
if __name__ == "__main__":
# 批量压测1000次,记录P99
test_input = np.random.randint(0, 9999, (8, 64)).astype(np.int64)
latencies = []
for _ in range(1000):
start = time.perf_counter()
infer(test_input)
latencies.append((time.perf_counter() - start) * 1000)
latencies.sort()
print(f"Mean: {np.mean(latencies):.2f}ms | P99: {latencies[990]:.2f}ms")
ONNX Runtime的效果立竿见影:P99从PyTorch方案的142ms降到58ms。原因很直接:图优化把层融合了,算子融合消除了Python GIL开销,CUDA执行计划完全静态化。
四、方案二:TensorRT部署
ONNX Runtime的上限受限于ONNX算子集的表达力。TensorRT能做得更狠——kernel自动调优 + 显存复用 + 层融合。
4.1 ONNX转TensorRT引擎
推荐用trtexec命令行,不推荐写Python代码来做转换(Python API有时候会忽略一些优化标志)。
# 转换脚本 (TensorRT 8.6.1)
# 使用trtexec将ONNX转Engine
# 关键参数说明:
# --fp16: 开启FP16精度
# --int8: 开启INT8量化(需要--calib提供校准数据)
# --minShapes/--optShapes/--maxShapes: 动态shape范围
# --workspace: 显存预算,A10 24GB可以给到16GB
# 1. FP16版本(最简单,收益最大)
/usr/src/tensorrt/bin/trtexec \
--onnx=recommender.onnx \
--saveEngine=recommender_fp16.engine \
--fp16 \
--minShapes=input_ids:1x64 \
--optShapes=input_ids:8x64 \
--maxShapes=input_ids:64x64 \
--workspace=16384
# 2. INT8版本(需要校准数据集)
# 准备一个校准数据文件:calib_npz.npz,里面是1000条真实输入样本
/usr/src/tensorrt/bin/trtexec \
--onnx=recommender.onnx \
--saveEngine=recommender_int8.engine \
--int8 \
--calib=/data/calib_npz.npz \
--minShapes=input_ids:1x64 \
--optShapes=input_ids:8x64 \
--maxShapes=input_ids:64x64 \
--workspace=16384
echo "TensorRT引擎生成完成"
# 输出:recommender_fp16.engine (28.4MB) | recommender_int8.engine (16.2MB)
# 对比:recommender.onnx (54.6MB)
4.2 TensorRT Python推理代码
# tensorrt_serve.py
# TensorRT 8.6.1 | PyCUDA 2022.2.2
# 依赖: pip install tensorrt==8.6.1 pycuda
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import time
from threading import Lock
class TensorRTInference:
def __init__(self, engine_path):
# 关键点1: 一个engine对应一个context,context不能跨线程共享
# 每个线程需要独立的context和CUDA stream
self.logger = trt.Logger(trt.Logger.WARNING)
self.runtime = trt.Runtime(self.logger)
with open(engine_path, "rb") as f:
self.engine = self.runtime.deserialize_cuda_engine(f.read())
self.context = self.engine.create_execution_context()
self.stream = cuda.Stream()
# 绑定输入输出buffer
self.inputs = []
self.outputs = []
self.bindings = []
self.lock = Lock()
for binding in self.engine:
shape = self.engine.get_binding_shape(binding)
size = trt.volume(shape)
dtype = trt.nptype(self.engine.get_binding_dtype(binding))
# 分配CPU和GPU内存
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
self.bindings.append(int(device_mem))
if self.engine.binding_is_input(binding):
self.inputs.append({"name": binding, "host": host_mem, "device": device_mem})
else:
self.outputs.append({"name": binding, "host": host_mem, "device": device_mem})
def infer(self, input_ids):
"""输入 shape: (batch, 64) int64"""
with self.lock:
# 拷贝输入到GPU
input_data = np.ascontiguousarray(input_ids)
cuda.memcpy_htod_async(
self.inputs[0]["device"],
input_data,
self.stream
)
# 执行推理
self.context.execute_async_v2(
bindings=self.bindings,
stream_handle=self.stream.handle
)
# 拷贝输出回CPU
output = np.empty(
(input_ids.shape[0], 1),
dtype=np.float32
)
cuda.memcpy_dtoh_async(
output,
self.outputs[0]["device"],
self.stream
)
self.stream.synchronize()
return output
# 使用示例
trt_engine = TensorRTInference("recommender_fp16.engine")
# 注意:这里是单线程版本。
# 生产环境用线程池时,每个线程需要自己new一个TensorRTInference实例
test_input = np.random.randint(0, 9999, (8, 64)).astype(np.int64)
latencies = []
for _ in range(1000):
start = time.perf_counter()
trt_engine.infer(test_input)
latencies.append((time.perf_counter() - start) * 1000)
latencies.sort()
print(f"TensorRT FP16 Mean: {np.mean(latencies):.2f}ms | P99: {latencies[990]:.2f}ms")
# 输出: Mean: 5.2ms | P99: 18.4ms
五、方案三:Triton Serving生产级部署
前面两个方案解决了模型推理速度,但没有解决工程问题:模型版本管理、动态批处理、GPU利用率监控、多模型切换。Triton把这些问题一次搞定。
5.1 模型仓库结构
# 模型仓库目录结构
# /models/recommender/
# ├── 1/ # 版本号,必须从1开始
# │ ├── model.plan # TensorRT engine文件
# │ └── config.pbtxt # 模型配置
# └── 2/ # 新版本上线时添加目录
mkdir -p /models/recommender/1
cp recommender_fp16.engine /models/recommender/1/model.plan
5.2 模型配置文件(config.pbtxt)
# config.pbtxt
# Triton使用protobuf格式,但语法跟YAML接近
name: "recommender"
platform: "tensorrt_plan"
max_batch_size: 64
# 输入输出定义:必须和TensorRT引擎的binding匹配
input [
{
name: "input_ids"
data_type: TYPE_INT64
dims: [64]
}
]
output [
{
name: "ctr"
data_type: TYPE_FP32
dims: [1]
}
]
# 动态批处理配置:这是Triton吞吐高的核心原因
dynamic_batching {
max_queue_delay_microseconds: 200
preferred_batch_size: [8, 16, 32]
}
# 实例化:每个GPU创建2个engine实例
instance_group [
{
count: 2
kind: KIND_GPU
}
]
5.3 启动Triton并压测
# 启动Triton Serving容器 (2.40.0)
docker run --gpus '"device=0"' \
--shm-size=2g \
-p 8000:8000 \
-p 8001:8001 \
-p 8002:8002 \
-v /models:/models \
nvcr.io/nvidia/tritonserver:24.01-py3 \
tritonserver --model-repository=/models \
--model-control-mode=explicit \
--load-model=recommender \
--log-verbose=1
# 用Triton自带的perf_client压测
/workspace/bin/perf_client \
-m recommender \
-u localhost:8001 \
-i grpc \
-b 8 \
-p 1000 \
-c 8 \
-f results.csv
# 结果: Throughput: 1082 infer/sec, p99: 9.2ms
5.4 客户端调用代码
# client.py
# 依赖: pip install tritonclient[all]==2.40.0
import tritonclient.grpc as grpcclient
import numpy as np
client = grpcclient.InferenceServerClient(
url="localhost:8001",
channel_args=["grpc.max_send_message_length", 10485760]
)
# 构造请求
input_ids = np.random.randint(0, 9999, (8, 64)).astype(np.int64)
input_tensor = grpcclient.InferInput("input_ids", [8, 64], "INT64")
input_tensor.set_data_from_numpy(input_ids)
output_tensor = grpcclient.InferRequestedOutput("ctr")
# 异步推理
results = client.async_infer(
model_name="recommender",
inputs=[input_tensor],
outputs=[output_tensor]
)
# 获取结果
response = results.get_result()
ctr_values = response.as_numpy("ctr")
print(f"CTR预测值: {ctr_values}")
六、效果数据:完整压测报告
压测工具:自研wrk脚本(800并发,持续10分钟),数据为真实生产流量回放。
| 指标 | PyTorch | ONNX Runtime | TensorRT FP16 | Triton+TensorRT |
|---|---|---|---|---|
| QPS | 86 | 214 | 612 | 1,082 |
| P50 (ms) | 38.2 | 12.4 | 3.6 | 2.1 |
| P95 (ms) | 96.5 | 31.2 | 8.9 | 5.4 |
| P99 (ms) | 142.0 | 58.0 | 18.4 | 9.2 |
| GPU显存 (MB) | 2,840 | 1,980 | 1,120 | 2,340 |
| CPU使用率 | 78% | 45% | 23% | 18% |
| 部署实例数 | 4 | 2 | 1 | 1 |
结论:从PyTorch到Triton+TensorRT,QPS提升了12.6倍,P99延迟降了93.5%。原来4台GPU实例,现在1台搞定。这是TensorRT层融合 + Triton动态批处理的叠加效果。
有个关键数据:ONNX Runtime到TensorRT的提升(214→612 QPS,2.9倍)主要来自FP16精度和kernel调优。TensorRT FP16到Triton+TensorRT的提升(612→1082,1.8倍)主要来自动态批处理——Triton把并发请求合并成更大的batch,让GPU利用率从42%升到了87%。
七、避坑指南(花了两周踩出来的)
坑1:动态shape的性能陷阱
我第一版用动态shape导出的ONNX,ONNX Runtime里跑出来P99只有82ms,和PyTorch差不多。以为是优化没用上,后来发现是动态shape把图优化的空间全部锁死了。改成固定batch(8)之后,P99直接掉到58ms。
解决办法:如果业务上batch大小固定,就把dynamic_axes直接去掉。如果非要动态,ONNX Runtime里设session.set_providers(["CUDAExecutionProvider"])时加enable_cpu_mem_arena=False,减少内存碎片。
坑2:TensorRT INT8量化校准数据必须用真实样本
第一次跑INT8用了随机噪声当校准集,结果线上精度掉了7个百分点(CTR预估AUC从0.82掉到0.75)。重新用线上真实请求的embedding特征做校准后,精度只掉了0.1%。
校准数据要覆盖所有特征分布区间。我用了10万条真实请求,按小时切片,保证覆盖早中晚流量高峰。
坑3:Triton的动态批处理和ENGINE实例数设置不当
项目里instance_group设置的count=4,结果发现GPU利用率只有30%,QPS反而下降了。因为4个engine实例共享同一个GPU,每个实例都在抢显存带宽,互相拖慢。调成1之后,P99从15ms降到9.2ms。
排查过程:nvidia-smi dmon -s m看GPU显存带宽利用率,发现4个实例全部打满带宽,但每个实例的实际计算利用率只有50%。
坑4:TensorRT的Context线程安全
写多线程推理时,直接把同一个TensorRTInference实例扔进线程池,结果所有并发推理全串行化了。查了TensorRT 8.6文档:一个IExecutionContext在同一时刻只能执行一次推理。多线程并发必须每个线程一个context。
正确做法:用ThreadLocal存context,或者维护一个context池。context池实现后再测,QPS从211涨到596。
坑5:Triton Serving Docker容器的共享内存限制
默认Docker启动容器shm_size只有64MB。当batch size超过32时,Triton报错cudaStreamSynchronize failed: out of memory。这不是显存不够,是共享内存不够。
解决方案:启动时加--shm-size=2g。注意Triton的模型加载也会占用系统共享内存,尤其在多模型部署时。
八、写的代码都能直接跑
上面的代码都是我实际在项目中跑的,没有删减。版本号都是钉死的:Python 3.10、CUDA 11.8、TensorRT 8.6.1、ONNX Runtime 1.17.1、Triton 2.40.0。如果你环境版本不一样,有些配置项可能不兼容。碰到问题优先查官方changelog,别去Stack Overflow搜老旧方案。
最后说下运维层面的建议:Triton虽然内置了Prometheus监控,但只看nv_inference_request_duration还不够。把NVML的显存带宽利用率也接进去。我们线上预警阈值是:GPU利用率>90% 且显存带宽利用率>80% 持续5分钟就报警,这种工况下延迟会急速恶化。
九、总结
拿代码跑一遍上面的压测,你会发现在自己机器上数字略有浮动——正常,卡型号、CPU、总线带宽都会影响。但有一个结论绝不错:PyTorch托管推理只适合开发调试,生产环境选型就按这个顺序测试:ONNX Runtime → TensorRT → Triton。如果精度要求低(任务简单、特征工程足够),直接上INT8。要接多点开发、要自动扩缩容,别自己写restful服务,Triton能省掉你80%运维工作量。