共计 2820 个字符，预计需要花费 8 分钟才能阅读完成。
1. 背景：大模型推理的算力挑战
在部署 780m 参数量的 AI 模型时，初学者常遇到两个典型问题：
- 显存不足：模型加载后立即触发 OOM（Out Of Memory）错误，尤其在消费级显卡（如 RTX 3060 12GB）上更为明显
- 推理延迟高：单次请求处理时间超过 500ms，无法满足实时性要求（如对话系统需 <200ms）
以 Huggingface 的bert-base-uncased（110M 参数）为基准，780m 模型的显存占用约为其 7 倍，这意味着即使是 24GB 显存的 3090 显卡，也可能在未优化的情况下仅支持 batch_size= 1 的推理。
2. 技术选型：PyTorch 原生 vs TensorRT 优化
2.1 性能对比
| 指标 | PyTorch 原生 (FP32) | TensorRT (FP16) | 提升幅度 |
|---|---|---|---|
| 显存占用 (batch=1) | 6.8GB | 3.2GB | 53%↓ |
| 平均延迟 (ms) | 142 | 67 | 2.1x↑ |
| 最大 batch_size | 4 | 8 | 2x↑ |
测试环境：AWS g4dn.xlarge（T4 GPU/16GB 显存）
2.2 精度损失验证
使用 FP16 量化时，在 GLUE 基准测试上的准确率变化：
- MNLI-m：84.3% → 84.1% (Δ=0.2%)
- QQP：91.1% → 90.9% (Δ=0.2%)
精度损失在可接受范围内，适合大多数生产场景。
3. 核心实现步骤
3.1 Flask API 封装
from flask import Flask, request
import torch
app = Flask(__name__)
model = torch.load('780m_model.pt').cuda().half() # FP16 转换
@app.route('/predict', methods=['POST'])
def predict():
inputs = request.json['text']
tokens = tokenizer(inputs, return_tensors='pt').to('cuda')
with torch.no_grad():
outputs = model(**tokens)
return {'logits': outputs.logits.cpu().numpy().tolist()}
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000)3.2 TensorRT 转换（关键注释版）
import tensorrt as trt
# 初始化 builder
logger = trt.Logger(trt.Logger.INFO)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
# 加载 PyTorch 模型并转换为 ONNX
torch.onnx.export(
model,
dummy_input,
"model.onnx",
opset_version=13,
input_names=['input_ids', 'attention_mask'],
output_names=['logits']
)
# 解析 ONNX 模型
parser = trt.OnnxParser(network, logger)
with open("model.onnx", "rb") as f:
parser.parse(f.read())
# FP16 量化与层融合优化
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16) # 开启 FP16
config.set_flag(trt.BuilderFlag.STRICT_TYPES)
# 动态 shape 配置（应对变长输入）profile = builder.create_optimization_profile()
profile.set_shape(
"input_ids",
min=(1, 1),
opt=(1, 128),
max=(1, 512)
)
config.add_optimization_profile(profile)
# 生成引擎
engine = builder.build_engine(network, config)
with open("model.trt", "wb") as f:
f.write(engine.serialize())4. 性能优化实战
4.1 Batch Size 与显存关系
| batch_size | PyTorch 显存 | TensorRT 显存 |
|---|---|---|
| 1 | 6.8GB | 3.2GB |
| 2 | 8.1GB | 4.0GB |
| 4 | OOM | 6.5GB |
| 8 | – | 11.2GB |
建议：根据业务延迟要求选择最大 batch_size，通常 batch= 4 时性价比最高
4.2 CUDA Stream 并行
import torch.cuda.stream as stream
# 创建多个流
streams = [torch.cuda.Stream() for _ in range(4)]
# 并行处理
for i in range(4):
with torch.cuda.stream(streams[i]):
output = model(batch[i])
torch.cuda.synchronize() # 等待所有流完成实测可提升吞吐量 30%（RTX 3090 上从 45 QPS 升至 58 QPS）
5. 避坑指南
5.1 OOM 错误处理
- 现象 ：
CUDA out of memory报错 - 解决方案：
- 检查
torch.cuda.memory_summary()找到内存峰值 - 使用
with torch.inference_mode():减少计算图保存 - 添加
torch.cuda.empty_cache()清理碎片
5.2 动态 Shape 支持
- 问题：输入文本长度变化导致需要重新构建引擎
- 解决：在 TensorRT 中预定义多个 profile（如 32/64/128/256 长度）
5.3 量化后精度异常
- 排查步骤：
- 对比 FP32 和 FP16 的输出差异
torch.allclose(fp32_out, fp16_out, atol=1e-3) - 检查是否存在大数值计算（如 Softmax 前需做数值裁剪）
6. 部署建议
- 监控指标：
- 显存利用率（
nvidia-smi -l 1） - 分位数延迟（P50/P95/P99）
- 健康检查：
# 测试 API 响应 curl -X POST http://localhost:5000/predict -H "Content-Type: application/json" -d '{"text":"sample input"}' - 灰度发布：
- 先用 5% 流量测试新引擎
- 对比新旧版本的错误率和延迟
通过上述方法，我们在 T4 显卡上实现了 780m 模型的稳定部署，QPS（Queries Per Second）从最初的 12 提升到 35，显存成本降低 60%。建议初学者先从 TensorRT FP16 量化入手，逐步尝试更激进的 INT8 量化，最终实现在边缘设备的部署。
正文完
