问题：推理速度跟不上业务，直接上GPU也不行？

去年我们上线一个BERT文本分类模型，PyTorch FP32推理，batch_size=1，单次耗时50ms。业务要求QPS=200，延迟P99小于20ms。加T4显卡后，原始模型只能跑到40QPS，延迟80ms——因为显存不够、batch不够大。

试过PyTorch JIT、TorchScript，加速有限。后来把模型转ONNX，配合TensorRT，再上Triton Inference Server做batching，最终延迟降到8ms（FP16），QPS冲到500+。

本文将完整记录这个优化过程，并给出可复现的代码和踩坑点。

环境与基准

基准模型：基于bert-base-uncased微调的情感分类模型，输入长度max_seq_len=128，输出2类。原始PyTorch模型FP32推理延迟：

• batch_size=1: 51.2ms
• batch_size=8: 162.3ms (平均每样本20.3ms)

方案一：ONNX Runtime 直接加速

ONNX Runtime支持GPU执行，不需要额外编译，兼容性好。转ONNX的代码：

import torch
from transformers import BertForSequenceClassification

model = BertForSequenceClassification.from_pretrained("./bert_sentiment")
model.eval()
dummy_input = torch.randint(0, 30522, (1, 128), dtype=torch.long)
dummy_attention_mask = torch.ones(1, 128, dtype=torch.long)

torch.onnx.export(
    model,
    (dummy_input, dummy_attention_mask),
    "bert_sentiment.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch_size"},
        "attention_mask": {0: "batch_size"},
        "logits": {0: "batch_size"}
    },
    opset_version=17,
    do_constant_folding=True
)
print("ONNX exported.")

关键参数：dynamic_axes允许batch维度动态；opset_version=17是ONNX Runtime 1.17支持的最新版本。用更高的opset有时能拿到更好的算子融合，但要注意兼容性（后面避坑会讲）。

# 使用ONNX Runtime GPU推理
import onnxruntime as ort
import numpy as np
import time

providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
session = ort.InferenceSession('bert_sentiment.onnx', providers=providers)
input_name = session.get_inputs()[0].name
mask_name = session.get_inputs()[1].name

dummy_input = np.random.randint(0, 30522, (1, 128)).astype(np.int64)
dummy_mask = np.ones((1, 128), dtype=np.int64)

# warmup
for _ in range(10):
    session.run(None, {input_name: dummy_input, mask_name: dummy_mask})

# benchmark
latencies = []
for _ in range(500):
    start = time.perf_counter()
    session.run(None, {input_name: dummy_input, mask_name: dummy_mask})
    latencies.append((time.perf_counter() - start) * 1000)
print(f"ONNX Runtime GPU batch=1: {np.median(latencies):.2f}ms")

结果：batch=1 median 23.8ms，比PyTorch快54%。但还没达到10ms目标。

方案二：TensorRT FP16 + 动态shape

TensorRT能进一步融合算子、使用FP16或INT8。但转换过程麻烦：先转ONNX，再用trtexec或Python API转TensorRT。下面是使用Python API转FP16 TensorRT engine：

import tensorrt as trt
import os

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

def build_engine(onnx_file_path, engine_file_path, fp16=True):
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(EXPLICIT_BATCH)
    parser = trt.OnnxParser(network, TRT_LOGGER)
    
    with open(onnx_file_path, 'rb') as f:
        if not parser.parse(f.read()):
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            raise RuntimeError("Failed parse ONNX")
    
    config = builder.create_builder_config()
    if fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
    
    # 动态shape profile
    profile = builder.create_optimization_profile()
    profile.set_shape("input_ids", (1, 128), (8, 128), (32, 128))
    profile.set_shape("attention_mask", (1, 128), (8, 128), (32, 128))
    config.add_optimization_profile(profile)
    
    serialized_engine = builder.build_serialized_network(network, config)
    with open(engine_file_path, 'wb') as f:
        f.write(serialized_engine)
    print(f"Engine built: {engine_file_path}")

build_engine("bert_sentiment.onnx", "bert_fp16.trt", fp16=True)

注意：TensorRT 8.x以后推荐用build_serialized_network而不是旧版build_engine。dynamic shape必须设profile，否则推理时shape必须固定。

# 加载TensorRT engine并推理
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

with open("bert_fp16.trt", "rb") as f:
    runtime = trt.Runtime(TRT_LOGGER)
    engine = runtime.deserialize_cuda_engine(f.read())
    
context = engine.create_execution_context()
context.set_binding_shape(0, (1, 128))  # input_ids
context.set_binding_shape(1, (1, 128))  # attention_mask

# 分配GPU内存
d_inputs = []
d_outputs = []
bindings = []
for i in range(engine.num_bindings):
    dims = context.get_binding_shape(i)
    size = trt.volume(dims) * 1  # 用float32输出
    dtype = trt.nptype(engine.get_binding_dtype(i))
    if engine.binding_is_input(i):
        d_input = cuda.mem_alloc(size)
        d_inputs.append(d_input)
        bindings.append(d_input)
    else:
        d_output = cuda.mem_alloc(size)
        d_outputs.append(d_output)
        bindings.append(d_output)

# 准备数据
h_input = np.random.randint(0, 30522, (1, 128)).astype(np.int32)  # TRT要求int32
h_output = np.empty(1*2, dtype=np.float32)

# warmup
for _ in range(10):
    cuda.memcpy_htod(d_inputs[0], h_input)
    context.execute_v2(bindings)
    cuda.memcpy_dtoh(h_output, d_outputs[0])

# benchmark
latencies = []
for _ in range(500):
    start = time.perf_counter()
    cuda.memcpy_htod(d_inputs[0], h_input)
    context.execute_v2(bindings)
    cuda.memcpy_dtoh(h_output, d_outputs[0])
    latencies.append((time.perf_counter() - start) * 1000)
print(f"TensorRT FP16 batch=1: {np.median(latencies):.2f}ms")

结果：batch=1 median 7.6ms，FP16推理精度损失几乎可忽略（准确率从92.3%降到92.0%）。

方案三：Triton Inference Server 服务化

有了TensorRT engine，但直接暴露给业务太原始。Triton Inference Server提供模型管理、动态batching、并发推理。部署配置如下：

# model_repository/bert_sentiment/config.pbtxt
name: "bert_sentiment"
platform: "tensorrt_plan"
max_batch_size: 64
input [
  {
    name: "input_ids"
    data_type: TYPE_INT32
    dims: [128]
  },
  {
    name: "attention_mask"
    data_type: TYPE_INT32
    dims: [128]
  }
]
output [
  {
    name: "logits"
    data_type: TYPE_FP32
    dims: [2]
  }
]
dynamic_batching {
  preferred_batch_size: [8, 16, 32]
  max_queue_delay_microseconds: 100
}
instance_group [
  {
    count: 2
    kind: KIND_GPU
  }
]

启动Triton：

docker run --gpus=1 --rm -p 8000:8000 -p 8001:8001 \
  -v /path/to/model_repository:/models \
  nvcr.io/nvidia/tritonserver:23.10-py3 \
  tritonserver --model-repository=/models

客户端请求（gRPC）：

import tritonclient.grpc as grpcclient
import numpy as np

client = grpcclient.InferenceServerClient(url="localhost:8001")
input_ids = np.random.randint(0, 30522, (1, 128)).astype(np.int32)
mask = np.ones((1, 128), dtype=np.int32)

inputs = [grpcclient.InferInput("input_ids", input_ids.shape, "INT32"),
          grpcclient.InferInput("attention_mask", mask.shape, "INT32")]
inputs[0].set_data_from_numpy(input_ids)
inputs[1].set_data_from_numpy(mask)

outputs = [grpcclient.InferRequestedOutput("logits")]

response = client.infer(model_name="bert_sentiment", inputs=inputs, outputs=outputs)
logits = response.as_numpy("logits")
print(logits)

使用Triton的dynamic_batching，当多个请求同时到达时自动合并成batch，大幅提高GPU利用率。压测结果：

• 并发线程4，每个线程发无限请求：QPS从单请求的130提升到460，延迟P99=12ms
• 如果禁用dynamic batching：QPS=150，P99=22ms

效果数据汇总

关键发现：ONNX Runtime是“零成本”的第一档加速，TensorRT是深度加速，Triton是生产级服务化。三者可叠加：ONNX→TRT→Triton。纯模型推理延迟，TensorRT比ONNX快3倍，比PyTorch快7倍。加上batching，吞吐提升10倍以上。

避坑指南（实际踩过的坑）

1. ONNX opset版本与算子兼容

用opset=17时，BERT的Gelu近似函数可能不支持。我们当时在转ONNX时报:"Unsupported: operator 'com.microsoft::Gelu'"。解决方案：设置opset=15或显式定义Gelu为近似。或者升级onnxruntime到1.17+并安装onnxruntime-gelu扩展。最终我们改成opset=14，一切正常。

测试：opset=14与opset=17在RTX 4090上性能无明显差异，但兼容性更好。

2. TensorRT动态shape的profile设置

如果client请求的batch size超出profile范围，会报错"Profile shape is out of range"。必须覆盖最小、最优、最大三个值。最优（opt）设成业务最常用的batch size，比如8或16。如果业务batch从1到64都可能，profile要设(1, 8, 64)，但性能在非最优shape时下降。我们线上固定batch=16，所以profile只设了一个点(16,16,16)，性能最佳。

3. Triton版本与TensorRT engine兼容

Triton 2.43打包的TensorRT是8.6.1，你必须用同一版本生成engine，否则加载报错。不同大版本之间完全不兼容，例如TRT 8.5生成的engine在8.6上会报"incompatible library version"。建议使用Triton官方镜像内的trtexec或python环境进行转换。

4. 显存泄漏

Triton默认每个模型实例独立加载engine，如果有多个instance_group，每个实例都会分配显存。我们一开始设了4个instance，结果16GB显存爆了。根据显存大小调成2个instance，并设置optimization { priority: "PREFER_THROUGHPUT" }，平衡显存与并发。

5. 动态batching的max_queue_delay

max_queue_delay_microseconds设太大会增加延迟，太小又无法充分batching。我们经验值：100-200微秒，延迟增加<1ms，但吞吐提升30%。需要根据实际压测调整。

6. PyTorch转ONNX时dynamic_axes不能用batch维度以外的动态？

我们的模型支持可变长度seq_len，但在转ONNX时dynamic_axes只设了batch维度。如果要动态seq_len，也必须加上。但TensorRT要求profile必须指定所有动态维度的范围。我们采取了截断+padding到固定长度128，省去了动态seq的麻烦。如果业务需要变长，建议用TensorRT的IElementWiseLayer手动实现，或者用ONNX Runtime的DynamicQuantizeLSTM等。

总结建议

• 快速起步：直接上ONNX Runtime GPU，零修改，提速2-3倍。
• 追求极致性能：TensorRT FP16，最好配合Triton的dynamic batching。
• 生产环境必须上Triton：模型版本管理、动态batching、并发控制、Prometheus监控。
• 不要迷信INT8：精度损失严重，且需要校准数据集，除非模型非常大或者对精度不敏感。

以上所有代码和配置均可在GitHub仓库找到：github.com/example/bert-deployment