TensorRT 10.x 在L40s GPU上运行Unet模型结果异常问题分析

问题背景

在使用TensorRT进行Unet模型(1.9GB/fp32)转换时，开发人员遇到了一个有趣的现象：在TensorRT 8.6版本下模型能够正确运行但性能下降15%，而在TensorRT 10.0.0/10.5/10.8版本下却出现了NaN结果。这一问题在NVIDIA L40s GPU上尤为明显。

问题现象深入分析

通过polygraphy工具进行调试时，发现了两个关键现象：

1. 当仅比较最终输出时，TensorRT和ONNX Runtime的结果存在显著差异
2. 当比较所有中间层输出时，两者的结果却又完全一致

这种矛盾现象暗示问题可能出在TensorRT的优化过程中，特别是与某些特定层的融合优化有关。

核心问题定位

经过逐层调试，发现问题集中在模型的attention机制实现部分。原始代码中的q/k/v张量处理在TensorRT 10.x版本中可能被过度优化或错误融合。具体表现为：

q, k, v = qkv.reshape(b, 3 * c, -1).chunk(3, dim=1)
h = self.attention(q, k, v)  # 原始代码会导致结果不一致

当修改为以下形式时，问题消失：

h = self.attention(v, v, v)  # 调试代码结果正确
h = q+k+v  # 替代方案也能正常工作

技术原理分析

TensorRT 10.x版本相比8.6版本引入了更多高级优化策略，包括：

1. 更激进的算子融合策略
2. 针对新型GPU架构(L40s)的特殊优化
3. 对attention机制的专门优化处理

在Unet模型的attention实现中，q/k张量的处理可能触发了TensorRT的某种优化路径，导致数值计算出现偏差。而通过强制保留这些张量(如将其加入输出)或改变计算方式，可以避免触发问题优化路径。

解决方案与建议

1. 显式保留关键张量：通过添加无实际影响的运算(如计算均值)来防止TensorRT过度优化特定张量

2. 精度控制：尝试使用混合精度或强制FP32模式，观察是否能规避问题

3. 版本兼容性检查：确认模型使用的ONNX opset版本与TensorRT 10.x的兼容性

4. 自定义插件替代：对于问题严重的部分，考虑使用TensorRT插件实现自定义算子

5. 详细日志分析：启用TensorRT的详细日志和性能分析工具，定位具体出错的优化阶段

经验总结

这一案例展示了深度学习模型在不同版本推理引擎中的兼容性挑战。对于复杂模型结构(特别是包含自定义attention机制的模型)，建议：

• 保持模型转换流程的版本一致性
• 建立完善的数值验证机制
• 对关键组件进行隔离测试
• 保留多种实现方案以应对兼容性问题

TensorRT 10.x在性能提升的同时，对模型结构的假设更为严格，开发者在升级时需要更加谨慎地验证模型行为。