(含代码)利用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 模子每一层的最佳参数设置。 您可以根据特征图的具体特征开发自己的调解管道。 与基于客户需求的现有方法相比，此功能使 TensorRT 量化能够获得卓越的图像质量，生存原始图像的丰富细节。
根据 Q-Diffusion 的研究结果，激活分布在不同的时间步长内可能会有很大差异，并且图像的形状和整体风格主要在去噪过程的初始阶段确定。 因此，使用传统的最大校准会导致初始步骤中出现较大的量化偏差。

相反，我们有选择地使用所选步骤范围中的最小量化缩放因子，因为我们发现激活中的异常值对终极图像质量并不那么告急。 这种量身定制的方法，我们将其命名为“百分比定量”，重点关注步长范围的告急百分位。 它使 TensorRT 能够生成与原始 FP16 精度生成的图像几乎雷同的图像。

使用 TensorRT 8 位量化加快扩散模子

/NVIDIA/TensorRT GitHub 存储库现在托管端到端、SDXL、8 位推理管道，提供即用型办理方案以在 NVIDIA GPU 上实现优化的推理速率。
运行单个命令即可使用 Percentile Quant 生成图像，并使用 demoDiffusion 测量延迟。 在本节中，我们使用 INT8 作为示例，但 FP8 的工作流程根本雷同。

1. python demo_txt2img_xl.py "enchanted winter forest with soft diffuse light on a snow-filled day" --version xl-1.0 --onnx-dir onnx-sdxl --engine-dir engine-sdxl --int8 --quantization-level 3

复制代码

以下是该命令所涉及的主要步骤的概述：

• 校准
  
• 导出 ONNX
  
• 构建 TensorRT 引擎
  

校准

校准是量化过程中计算目标精度范围的步骤。 目前，TensorRT 中的量化功能封装在 nvidia-ammo 中，该依靠项已包罗在 TensorRT 8 位量化示例中。

1. # Load the SDXL-1.0 base model from HuggingFace
2. import torch
3. from diffusers import DiffusionPipeline
4. base = DiffusionPipeline.from_pretrained(
5.     "stabilityai/stable-diffusion-xl-base-1.0",
6.     torch_dtype=torch.float16,
7.     variant="fp16",
8.     use_safetensors=True
9. )
10. base.to("cuda")
12. # Load calibration prompts:
13. from utils import load_calib_prompts
14. cali_prompts = load_calib_prompts(batch_size=2,prompts="./calib_prompts.txt")
16. # Create the int8 quantization recipe
17. from utils import get_percentilequant_config
18. quant_config = get_percentilequant_config(base.unet, quant_level=3.0, percentile=1.0, alpha=0.8)
20. # Apply the quantization recipe and run calibration  
21. import ammo.torch.quantization as atq
22. quantized_model = atq.quantize(base.unet, quant_config, forward_loop)
24. # Save the quantized model
25. import ammo.torch.opt as ato
26. ato.save(quantized_model, 'base.unet.int8.pt')

复制代码

导出 ONNX

获得量化模子检查点后，您可以导出 ONNX 模子。

1. # Prepare the onnx export  
2. from utils import filter_func, quantize_lvl
3. base.unet = ato.restore(base.unet, 'base.unet.int8.pt')
4. quantize_lvl(base.unet, quant_level=3.0)
5. atq.disable_quantizer(base.unet, filter_func) # `filter_func` is used to exclude layers you don't quantize
6.   
7. # Export the ONNX model
8. from onnx_utils import ammo_export_sd
9. base.unet.to(torch.float32).to("cpu")
10. ammo_export_sd(base, 'onnx_dir', 'stabilityai/stable-diffusion-xl-base-1.0')

复制代码

构建 TensorRT 引擎

使用 INT8 UNet ONNX 模子，您可以构建 TensorRT 引擎。

1. trtexec --onnx=./onnx_dir/unet.onnx --shapes=sample:2x4x128x128,timestep:1,encoder_hidden_states:2x77x2048,text_embeds:2x1280,time_ids:2x6 --fp16 --int8 --builderOptimizationLevel=4 --saveEngine=unetxl.trt.plan

复制代码

总结

在生成式人工智能期间，拥有优先思量易用性的推理办理方案至关告急。 借助 NVIDIA TensorRT，您可以通过其专有的 8 位量化技术无缝实现高达 2 倍的推理速率加快，同时确保图像质量不受影响，从而实现卓越的用户体验。
TensorRT 对平衡速率和质量的承诺凸显了其作为加快 AI 应用步调的领先选择的地位，使您能够轻松交付尖端办理方案。
我将在 NVIDIA GTC 大会期间为大家带来免费中文在线解读：
NVIDIA CUDA 最新特性以及生成式 AI 相干内容，包括 Stable Diffusion 模子部署实践，以及介绍用于视觉内容生成的 Edify 模子，点击链接相识详情并注册参会：
https://www.nvidia.cn/gtc-global/session-catalog/?search=WP62435%20WP62832%20WP62400&ncid=ref-dev-945313#/

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。