(含代码)利用TensorRT的8位PTQ将Stable Diffusion速度提高 2 倍

利用TensorRT的8位PTQ将Stable Diffusion速度提高 2 倍

在生成人工智能的动态领域中，扩散模型脱颖而出，成为生成带有文本提示的高质量图像的最强大的架构。 像稳定扩散这样的模型已经彻底改变了创意应用。

然而，由于需要迭代去噪步骤，扩散模型的推理过程可能需要大量计算。 这对于努力实现最佳端到端推理速度的公司和开发人员提出了重大挑战。

从 NVIDIA TensorRT 9.2.0 开始，我们开发了一流的量化工具包，具有改进的 8 位（FP8 或 INT8）训练后量化 (PTQ: Post-Training Quantization)，可显着加快 NVIDIA 硬件上的扩散部署，同时保持图像质量 。 TensorRT 的 8 位量化功能已成为许多生成型 AI 公司的首选解决方案，特别是创意视频编辑应用程序的领先提供商。

在这篇文章中，我们讨论 TensorRT 与 Stable Diffusion XL 的性能。 我们介绍了使 TensorRT 成为低延迟稳定扩散推理的首选的技术优势。 最后，我们演示如何使用 TensorRT 通过几行更改来加速模型。

性能指标

与在 FP16 中运行的本机 PyTorch 的 torch.compile 相比，用于扩散模型的 NVIDIA TensorRT INT8 和 FP8 量化方案在 NVIDIA RTX 6000 Ada GPU 上实现了 1.72 倍和 1.95 倍的加速。 FP8 相对于 INT8 的额外加速主要归因于多头注意力 (MHA) 层的量化。 使用 TensorRT 8 位量化可以增强生成式 AI 应用程序的响应能力并降低推理成本。

除了加速推理之外，TensorRT 8 位量化还擅长保持图像质量。 通过专有的量化技术，它生成与原始 FP16 图像非常相似的图像。 我们将在本文后面介绍这些技术。

TensorRT 解决方案：克服推理速度挑战

尽管 PTQ 被认为是减少内存占用并加速许多 AI 任务推理的首选压缩方法，但它在扩散模型上并不能开箱即用。 扩散模型具有独特的多时间步去噪过程，并且噪声估计网络在每个时间步的输出分布可能会有很大变化。 这使得简单的 PTQ 校准方法不适用。

在现有技术中，SmoothQuant 作为一种流行的 PTQ 方法脱颖而出，可为 LLM 实现 8 位权重、8 位激活 (W8A8) 量化。 其主要创新在于解决激活异常值的方法，通过数学上等效的变换将量化挑战从激活转移到权重。

尽管它很有效，但用户在 SmoothQuant 中手动定义参数时经常遇到困难。 实证研究还表明，SmoothQuant 难以适应不同的图像特征，限制了其在现实场景中的灵活性和性能。 此外，其他现有的扩散模型量化技术仅针对单个版本的扩散模型量身定制，而用户正在寻找一种可以加速各种版本模型的通用方法。

为了应对这些挑战，NVIDIA TensorRT 开发了复杂的细粒度调整管道，以确定 SmoothQuant 模型每一层的最佳参数设置。 您可以根据特征图的具体特征开发自己的调整管道。 与基于客户需求的现有方法相比，此功能使 Tens