TensorRT中SDXL模型INT8量化性能问题分析与解决方案

概述

在TensorRT 9.3版本中，用户报告了SDXL(Stable Diffusion XL)模型在INT8量化后推理速度反而比FP16版本慢的问题。本文将深入分析这一现象的原因，并介绍NVIDIA官方的解决方案。

问题现象

测试环境使用A800 GPU，对比FP16和INT8量化版本的SDXL模型推理性能：

FP16版本：

• UNet模块30次降噪耗时2616.81ms
• 整体管道耗时2851.01ms
• 吞吐量0.35 image/s

INT8版本(quantization-level 3)：

• UNet模块30次降噪耗时5550.13ms
• 整体管道耗时5785.81ms
• 吞吐量0.17 image/s

可以看到，INT8量化后的模型推理速度仅为FP16版本的一半左右，这与预期的量化加速效果相反。

问题原因分析

经过NVIDIA工程师调查，发现主要原因在于：

1. MHA融合限制：TensorRT当前不支持序列长度大于512时的INT8 MHA(Multi-Head Attention)融合，这是出于精度考虑的设计决策。

2. 量化级别选择：当使用quantization-level 3(全量化)时，由于上述限制，会导致性能下降而非提升。

3. ONNX模型大小异常：用户还观察到INT8量化后的ONNX模型体积异常增大，但最终的.plan文件大小正常。

解决方案

NVIDIA工程师提供了以下解决方案：

1. 调整量化级别：建议使用--quantization-level 2.5参数，这样MHA部分不会被量化，可以避免性能下降问题。

2. 代码修改：对于TensorRT 9.3版本，需要手动修改utilities.py文件，移除choices=range(1,4)的限制，才能使用2.5级别的量化。

3. 版本升级：NVIDIA表示将在TensorRT 10.0 GA版本中修复此问题，包括更新AMMO版本和最新的校准脚本。

性能对比

用户尝试了不同解决方案后的性能表现：

1. 升级PyTorch 2.0：带来轻微性能提升(UNet从2616ms降至2315ms)，但效果有限。

2. 使用quantization-level 2.5：在某些情况下性能反而更差(UNet耗时增至7193ms)，这表明问题可能更复杂。

3. TensorRT 10.0测试：有用户在A100上测试发现，即使使用TRT 10.0和AMMO量化，INT8版本仍略慢于FP16(UNet 377ms vs 373ms)。

最佳实践建议

基于当前情况，建议用户：

1. 对于SDXL等大模型，谨慎使用INT8量化，特别是在序列长度较大的场景下。

2. 关注TensorRT 10.0的正式发布，其中包含对量化性能的改进。

3. 可以尝试FP8量化(即将支持)，这可能更适合大模型场景。

4. 在实际应用中，建议进行充分的基准测试，量化不总是能带来性能提升，需要结合具体模型和硬件进行评估。

总结

TensorRT中SDXL模型的INT8量化性能问题揭示了深度学习模型量化在实际应用中的复杂性。NVIDIA正在积极改进量化支持，特别是针对大模型和长序列场景。用户应保持对最新版本的关注，并根据自身应用场景选择合适的量化策略。