Triton推理服务器中ORT后端的GPU显存泄漏问题分析

问题概述

在Triton推理服务器使用ONNX Runtime(ORT)后端时，当以特定批处理大小序列发送推理请求时，会出现GPU显存持续泄漏的问题。这个问题在较新版本的Triton容器(nvcr.io/nvidia/tritonserver:24.12-py3和25.02-py3)中出现，但在较早版本(24.05-py3)中不存在。

问题表现

该问题表现为当以特定批处理大小序列(如[16,7,7,7...])发送请求时，每次推理后GPU显存都会增加约1024MB，最终导致显存耗尽和推理失败。而使用其他批处理大小序列(如[16,6,6,6...])则不会出现此问题。

技术背景

Triton推理服务器是一个高性能的机器学习模型服务系统，支持多种后端框架，包括ONNX Runtime。ORT后端允许用户在Triton中部署ONNX格式的模型。在模型推理过程中，Triton会管理GPU显存的分配和释放，确保高效利用硬件资源。

问题分析

触发条件

1. 批处理大小序列：当最大批处理大小为16时，后续请求使用7-15之间的批处理大小会触发泄漏
2. 模型结构：问题与模型中的特定操作(重采样器、CTC解码、数据类型转换)相关
3. Triton版本：仅出现在24.12和25.02版本中，24.05版本正常

可能原因

1. 内存池管理问题：ORT后端在特定批处理大小下可能没有正确释放临时分配的显存
2. 批处理机制缺陷：Triton的批处理机制在处理非最大批处理大小时可能存在资源回收问题
3. 版本兼容性问题：新版本Triton与ORT的交互方式可能引入了新的内存管理逻辑缺陷

解决方案与建议

临时解决方案

1. 使用Triton 24.05版本容器
2. 改用TensorRT后端(测试表明该后端无此问题)
3. 避免使用会触发泄漏的批处理大小序列

长期解决方案

1. 等待NVIDIA官方修复此问题
2. 考虑在应用层实现批处理大小控制逻辑，避免触发条件
3. 对模型进行优化，移除可能引发问题的操作(如重采样器或CTC解码)

技术验证方法

开发人员可以通过以下步骤验证问题：

1. 准备一个包含重采样器和CTC解码的ONNX模型
2. 配置Triton模型仓库，设置max_batch_size为16
3. 使用Python客户端发送特定批处理大小序列的请求
4. 监控nvidia-smi显示的显存使用情况

总结

Triton推理服务器在较新版本中出现的ORT后端显存泄漏问题，主要与特定批处理大小序列下的内存管理机制有关。虽然目前可以通过降级或更换后端来规避，但长期来看需要NVIDIA官方修复此问题。开发人员在生产环境中部署ONNX模型时，应充分进行压力测试和内存监控，确保系统稳定性。

对于性能敏感的应用场景，考虑使用TensorRT后端可能是一个更稳定的选择，但需要注意模型转换可能带来的精度损失问题。