重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和新库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

Stable Diffusion XL Int8 量化

此示例展示了如何使用 ModelOpt 来校准和量化 SDXL 的 UNet 部分。UNet 部分通常消耗超过 95% 的端到端 Stable Diffusion 延迟。

我们还提供了关于部署和运行带有 ModelOpt 量化 int8 UNet 的端到端 SDXL 流水线以生成图像并在目标 GPU 上测量延迟的说明。

要开始使用，您需要 NeMo 格式的预训练 SDXL 检查点。示例训练配置在此处可用：Stable Diffusion Examples NVIDIA/NeMo。

校准

第一步是使用默认配置运行量化脚本。该脚本会将量化的 unet 导出到 onnx 文件。以下是 SDXL 量化脚本的默认配置：NVIDIA/NeMo。

quantize
  exp_name: nemo
  n_steps: 20          # number of inference steps
  format: 'int8'       # only int8 quantization is supported now
  percentile: 1.0      # Control quantization scaling factors (amax) collecting range, meaning that we will collect the minimum amax in the range of `(n_steps * percentile)` steps. Recommendation: 1.0
  batch_size: 1        # batch size calling sdxl inference pipeline during calibration
  calib_size: 32       # For SDXL, we recommend 32, 64 or 128
  quant_level: 2.5     #Which layers to be quantized, 1: `CNNs`, 2: `CNN + FFN`, 2.5: `CNN + FFN + QKV`, 3: `CNN + Linear`. Recommendation: 2, 2.5 and 3, depending on the requirements for image quality & speedup.
  alpha: 0.8           # A parameter in SmoothQuant, used for linear layers only. Recommendation: 0.8 for SDXL

重要参数

• percentile：控制收集量化比例因子 (amax) 的范围。这意味着我们将在 (n_steps * percentile) 步长的范围内收集最小 amax。我们建议使用 1.0。

• alpha：SmoothQuant 中的一个参数，仅用于线性层。对于 SDXL，我们建议使用 0.8，对于 SD 1.5，我们建议使用 1.0。

• quant-level：要量化的层级，1：CNN，2：CNN + FFN，2.5：CNN + FFN + QKV，3：CNN + Linear。建议：2、2.5 和 3，具体取决于对图像质量和加速的要求。

• calib-size：对于 SDXL，我们建议使用 32、64 或 128，对于 SD 1.5，设置为 512 或 1024。

为量化的 ONNX UNet 构建 TRT 引擎

trtexec --onnx=./nemo_onnx/unet.onnx --shapes=x:8x4x128x128,timesteps:8,context:8x80x2048,y:8x2816 --fp16 --int8 --builderOptimizationLevel=4 --saveEngine=nemo_unet_xl.plan

重要参数

构建 TRT 引擎时，您必须按如下方式提供输入形状

• x：输入图像潜在形状 (B * C * H * W)

• context：输入文本条件 (B * S * hidden_dimension)

• y：附加嵌入 (B * adm_in_channels)

使用 NeMo 构建端到端 Stable Diffusion XL 流水线

我们提供了一个脚本来构建带有 NeMo 后端的端到端 TRT 推理流水线：SDXL Export Script NVIDIA/NeMo。

infer:
    out_path: sdxl_export
    width: 1024
    height: 1024
    batch_size: 2

trt:
  static_batch: False
  min_batch_size: 1
  max_batch_size: 8

重要参数

• out_path：用于保存 onnx 文件和 TRT 引擎文件的目录。

• width 和 height：推理输出的图像分辨率。

• batch_size：仅用于虚拟输入生成和 onnx 健全性检查。

• {min,max}_batch_size：TRT 引擎沿其动态轴的输入批大小。

运行端到端 Stable Diffusion XL TRT 流水线

推理脚本可以在这里找到：SDXL TRT Inference Script NVIDIA/NeMo。

unet_xl: sdxl_export/plan/unet_xl.plan
vae: sdxl_export/plan/vae.plan
clip1: sdxl_export/plan/clip1.plan
clip2: sdxl_export/plan/clip2.plan

out_path: trt_output

请指定 unet_xl 作为量化的 Unet 引擎来运行量化解决方案。系统将默认加载原始引擎文件。

推理加速

TRT 版本 9.3.0 GPU：H100

TRT int8 vs Framework fp16

流水线	批大小	延迟 (毫秒)	流水线	批大小	延迟	加速
Framework fp16 基准	1	3056.01	ModelOpt TRT Int8	1	1406.68	2.172498365
Framework fp16 基准	2	4832.24	ModelOpt TRT Int8	2	2403.29	2.01067703
Framework fp16 基准	4	8433.71	ModelOpt TRT Int8	4	4252.6	1.983189108

TRT int8 vs TRT fp16

流水线	批大小	延迟 (毫秒)	精度	批大小	延迟	加速
fp16 基准	1	1723.97	ModelOpt Int8	1	1406.68	1.225559473
fp16 基准	2	3004.47	ModelOpt Int8	2	2403.29	1.250148754
fp16 基准	4	5657.19	ModelOpt Int8	4	4252.6	1.330289705

FP16 推理 vs Int8 推理

提示：一只背着背包骑自行车的柴犬的照片。它戴着太阳镜和沙滩帽。（FP16 上方 vs Int8 下方）

提示：一只可爱的柯基犬住在一个用寿司做的房子里。（FP16 上方 vs Int8 下方）

提示：一位非常快乐的模糊熊猫的对比鲜明的人像，他打扮成厨师，在一个高端厨房里揉面团。他身后的墙上挂着一幅花卉画。（FP16 上方 vs Int8 下方）