Model Optimizer 扩散模型量化部署实战：ONNX 导出与 TensorRT 引擎构建完整指南

本篇指南基于 Model Optimizer 仓库中的 examples/diffusers/quantization/ONNX-TRT-Deployment.md 及其配套脚本，系统讲解将扩散模型（SDXL、SD3、FLUX 等）经 Model Optimizer 量化（INT8/FP8/FP4）后导出 ONNX、再使用 trtexec 构建 TensorRT 引擎，并最终接入 demoDiffusion 或 DeviceModel 完成端到端文生图与延迟评测的完整工作流。读完本文，你将掌握：如何用一行脚本或精细化参数完成扩散模型主干网络的量化与 ONNX 导出；如何为不同模型正确配置 TensorRT 动态形状并生成 INT8/FP8 引擎；以及如何把量化后的主干无缝替换进推理管线，对比 FP16 基线与量化方案的生成质量与延迟。

本文描述的是可选的 ONNX 导出 + TensorRT 引擎工作流；若你只需要纯量化（PyTorch checkpoint / Hugging Face checkpoint），可参考 examples/diffusers/README.md 中的量化章节。

---

一、前置条件：环境与工作流总览

在开始之前，需要先完成 TensorRT 环境的搭建与 Model Optimizer 的安装（原文档要求参见 Pre-Requisites）：

• Docker 方式：建议直接使用 TensorRT 官方容器镜像（如 nvcr.io/nvidia/tensorrt），镜像内已包含 CUDA、TensorRT 等运行时依赖。
• 本地安装方式：使用 pip 安装带 ONNX 与 Hugging Face 依赖的 Model Optimizer 版本，并安装示例依赖：

pip install nvidia-modelopt[onnx,hf]
pip install -r requirements.txt

examples/diffusers/ 下各个子模块（fastgen、distillation、quantization 等）可能还有各自的 requirements.txt，需要按需单独安装。

整个工作流分为四个阶段，后续章节依次展开：

1. 量化与 ONNX 导出：使用 quantize.py（或一键脚本 build_sdxl_8bit_engine.sh）对扩散模型主干网络（backbone）执行 PTQ 校准，输出 PyTorch checkpoint 与 ONNX 模型；
2. TensorRT 引擎构建：使用 trtexec 读取 ONNX，按模型指定动态形状（--minShapes/--optShapes/--maxShapes）生成 .plan 引擎文件；
3. demoDiffusion 端到端验证：将量化引擎替换进 NVIDIA demoDiffusion 示例，对比 FP16 基线与 INT8 量化结果的图像质量与延迟；
4. DeviceModel 推理管线：不经过手动导出/构建，由 DeviceModel 接口自动完成"PyTorch checkpoint → ONNX → TensorRT engine"，以 torch 模型的方式直接跑通文生图管线。

---

二、第一步：量化并导出 ONNX

2.1 一键脚本：SDXL 主干 INT8/FP8 量化与导出

如果目标模型是 SDXL，可以直接运行仓库提供的 build_sdxl_8bit_engine.sh 一键完成"量化 + ONNX 导出"：

bash build_sdxl_8bit_engine.sh --format {FORMAT} # FORMAT 可选 int8 或 fp8

从脚本源码可以看到它内部封装的默认参数（build_sdxl_8bit_engine.sh）：

• 模型固定为 sdxl-1.0（Hugging Face ID 为 stabilityai/stable-diffusion-xl-base-1.0，见 models_utils.py）；
• FP8 路径：--batch-size 2 --calib-size 128 --n-steps 20 --collect-method default；
• INT8 路径：--batch-size 2 --calib-size 32 --collect-method "min-mean" --percentile 1.0 --alpha 0.8 --n-steps 20。

也就是说，脚本对 INT8 使用了"分位数校准（percentile=1.0）+ SmoothQuant（alpha=0.8）"的默认组合，对 FP8 使用默认校准器。导出成功后 ONNX 会输出到 <model>_<format>.onnx 目录（例如 sdxl-1.0_int8.onnx），同时保存 sdxl-1.0_int8.pt 量化 checkpoint。

2.2 自定义参数量化：INT8（FLUX-Dev / SD3-Medium / SDXL / SDXL-Turbo）

若需要在校准参数上做精细控制，或量化其他模型，请直接使用 quantize.py：

python quantize.py \
    --model {flux-dev|sdxl-1.0|sdxl-turbo|sd3-medium} \
    --format int8 --batch-size 2 \
    --calib-size 32 --alpha 0.8 --n-steps 20 \
    --model-dtype {Half/BFloat16} --trt-high-precision-dtype {Half|BFloat16} \
    --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --onnx-dir {ONNX_DIR}

2.3 自定义参数量化：FP8 / FP4（FLUX-Dev / SDXL / SDXL-Turbo / LTX-Video）

python quantize.py \
    --model {flux-dev|sdxl-1.0|sdxl-turbo|ltx-video-dev} --model-dtype {Half|BFloat16} --trt-high-precision-dtype {Half|BFloat16} \
    --format {fp8|fp4} --batch-size 2 --calib-size {128|256} \
    --n-steps 20 --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --collect-method default \
    --onnx-dir {ONNX_DIR}

其中 FP4 ONNX 导出目前仅支持 Flux 与 SDXL 两个模型家族（对应 ModelType.FLUX_DEV 与 SDXL 相关类型，见 models_utils.py）。

硬件要求：官方建议导出 ONNX 时使用 CPU + GPU 合计至少 48GB 内存 的设备；如果内存不足，请改用 CPU 进行 ONNX 导出（导出代码中会将 pipeline 切到 CPU、仅把 backbone 放到 CUDA 上执行，见 quantize.py）。

2.4 关键参数与底层实现

quantize.py 的参数全部在 create_argument_parser() 中定义（quantize.py），与量化直接相关的关键参数如下：

参数	默认值	说明
--model	flux-dev	目标模型，取值覆盖 sdxl-1.0、sdxl-turbo、sd3-medium、flux-dev、flux-schnell、ltx-video-dev、wan2.2-* 等（完整列表见 ModelType 枚举）
--format	int8	量化格式：int8 / fp8 / fp4
--quant-algo	max	量化算法：max（最大绝对值校准）、smoothquant（SmoothQuant）、svdquant（SVDQuant 低秩分解，适用于 NVFP4）
--percentile	1.0	分位数校准的百分位，仅对 INT8 生效（SmoothQuant 除外）；含义是只收集 (n_steps × percentile) 步范围内的 amax 统计。推荐值 1.0
--collect-method	default	校准统计量收集方法：global_min / min-max / min-mean / mean-max / default。FP8/FP4 目前仅支持 default（代码中对此有显式校验，见 quantize_config.py）
--alpha	1.0	SmoothQuant 的 alpha 参数，只作用于线性层。SDXL 官方推荐 0.8
--lowrank	32	SVDQuant 的低秩维度
--quantize-mha	False	是否将多头注意力（MHA）一并量化为 FP8；注意 INT8 格式下不支持该选项
--compress	False	是否压缩量化权重以降低内存占用（仅 FP8/FP4 支持）
--block-size	16	NVFP4 量化的块大小
--batch-size	2	校准 batch size
--calib-size	128	校准样本总数。SDXL INT8 推荐 32 或 64，SDXL FP8 推荐 128
--n-steps	30	校准去噪步数。SD/SDXL 推荐 20 或 30，SDXL-Turbo 推荐 4
--model-dtype	Half	加载 pipeline 的精度（Half / BFloat16 / Float），可用 --component-dtype vae:Half --component-dtype transformer:BFloat16 分组件指定
--trt-high-precision-dtype	Half	TensorRT 高精度层（如残差连接、部分归一化层）使用的精度，将被写入每个量化算子的 trt_high_precision_dtype 属性
--quantized-torch-ckpt-save-path	—	量化 PyTorch checkpoint 保存路径
--onnx-dir	—	ONNX 导出目录
--hf-ckpt-dir	—	Hugging Face 格式 checkpoint 导出目录（供 SGLang/vLLM/TRTLLM 等直接部署）
--restore-from	—	从已有 checkpoint 目录恢复量化模型（自动推断原量化格式，无需重新校准）
--prompts-file	—	使用自定义 prompt 文件替代默认校准数据集

量化配置的构建过程（对应 Quantizer.get_quant_config，quantize.py）值得展开说明：

• 根据 --format 选择基准预设：INT8 使用 INT8_DEFAULT_CONFIG（若 --quant-algo smoothquant 则改用 mtq.INT8_SMOOTHQUANT_CFG）；FP8 使用 FP8_DEFAULT_CONFIG；FP4 则针对 SDXL 系模型使用 NVFP4_FP8_CONV_CONFIG、对 Flux 系模型使用 NVFP4_FP8_MHA_CONFIG，其余使用 NVFP4_DEFAULT_CONFIG。这些预设通过 load_config("configs/ptq/presets/diffusers/...") 从 Model Optimizer 内置 recipe 加载（见 config.py）。
• INT8 且 --collect-method 非 default 时，会调用 reset_set_int8_config 为所有 nn.Conv2d 的输入量化器挂上 PercentileCalibrator，把 percentile、n_steps、collect_method 写进校准器参数（config.py）。
• FP4 场景会把 --block-size 覆盖进各量化算子的 block_sizes 配置。
• set_quant_config_attr 会把 trt_high_precision_dtype 与算法信息（SmoothQuant 的 alpha、SVDQuant 的 lowrank）注入量化配置。

校准流程由 Calibrator 驱动（calibration.py）：按 num_batches = ceil(calib_size / batch_size) 分批加载校准 prompt（默认来自 Gustavosta/Stable-Diffusion-Prompts 数据集，模型注册表见 models_utils.py），对每批 prompt 调用 pipe(prompt, num_inference_steps=n_steps) 跑完整去噪过程以收集激活统计量。LTX-Video、Wan 等视频模型则有专门的多阶段校准分支。

ONNX 导出最终落在 onnx_utils/export.py 的 modelopt_export_sd：

• 使用 PyTorch 原生 torch.onnx.export（dynamo=False），opset_version=20，并打开常量折叠；
• 每个模型家族都有预设的输入/输出名称与动态轴映射（如 SDXL 的 sample/timestep/encoder_hidden_states/text_embeds/time_ids → latent，FLUX 的 hidden_states/encoder_hidden_states/pooled_projections/timestep/img_ids/txt_ids[/guidance] → latent）；
• FP8 的 FLUX 模型会额外执行 flux_convert_rope_weight_type，把 Einsum 节点的 RoPE 权重转为 float32（见 export.py）；
• FP4 模型导出后会调用 NVFP4QuantExporter.process_model 完成 NVFP4 量化算子到 ONNX 的落图；
• ONNX 以外部数据（save_as_external_data=True）形式保存，所有张量合并到单个 _data 文件。

---

三、第二步：用 trtexec 构建量化主干的 TensorRT 引擎

[!IMPORTANT]

• 构建引擎前必须完成 TensorRT 环境搭建（见 Pre-Requisites）；
• INT8 要求 TensorRT 版本 ≥ 9.2.0；FP8 要求 TensorRT 版本 ≥ 10.2.0；
• 当前仓库示例暂不支持 SVDQuant 的部署。

3.1 SDXL：INT8/FP8 Backbone 引擎

# 适用于 SDXL
trtexec --builderOptimizationLevel=4 --stronglyTyped --onnx=./model.onnx \
    --minShapes=sample:2x4x128x128,timestep:1,encoder_hidden_states:2x77x2048,text_embeds:2x1280,time_ids:2x6 \
    --optShapes=sample:16x4x128x128,timestep:1,encoder_hidden_states:16x77x2048,text_embeds:16x1280,time_ids:16x6 \
    --maxShapes=sample:16x4x128x128,timestep:1,encoder_hidden_states:16x77x2048,text_embeds:16x1280,time_ids:16x6 \
    --saveEngine=model.plan

SDXL UNet 的输入包括：sample（潜空间特征图，通道数 4）、timestep（标量步数）、encoder_hidden_states（CLIP 文本编码，序列长 77）、text_embeds（池化文本嵌入，1280 维）与 time_ids（6 维时间 ID）。这些名称与 ONNX 导出时的 input_names 一一对应（见 export.py）。

3.2 SD3-Medium：INT8/FP8 Backbone 引擎

# 适用于 SD3-Medium
trtexec --builderOptimizationLevel=4 --stronglyTyped --onnx=./model.onnx \
    --minShapes=hidden_states:2x16x128x128,timestep:2,encoder_hidden_states:2x333x4096,pooled_projections:2x2048 \
    --optShapes=hidden_states:16x16x128x128,timestep:16,encoder_hidden_states:16x333x4096,pooled_projections:16x2048 \
    --maxShapes=hidden_states:16x16x128x128,timestep:16,encoder_hidden_states:16x333x4096,pooled_projections:16x2048 \
    --saveEngine=model.plan

SD3 的 DiT 主干输入为 hidden_states、timestep、encoder_hidden_states（joint attention，333 个文本 token）与 pooled_projections。

3.3 FLUX-Dev：FP8 Backbone 引擎

# 适用于 FLUX-Dev FP8
trtexec --onnx=./model.onnx --fp8 --bf16 --stronglyTyped \
    --minShapes=hidden_states:1x4096x64,img_ids:4096x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 \
    --optShapes=hidden_states:1x4096x64,img_ids:4096x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 \
    --maxShapes=hidden_states:1x4096x64,img_ids:4096x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 \
    --saveEngine=model.plan

FLUX 的 Transformer 主干输入较为特殊：hidden_states 是展平的图像 token（4096 × 64 通道）、img_ids/txt_ids 是图像与文本的位置 ID 张量、guidance 是引导标量（仅 flux-dev 具备，flux-schnell 没有该输入，对应导出逻辑见 export.py）。FLUX FP8 引擎还需要 --fp8 --bf16 标志并在 --stronglyTyped 模式下构建。

3.4 动态形状注意事项

maxShapes 代表给定张量允许的最大形状。 如果你想使用更大的 batch size 或调整其他维度，请相应修改这些值。optShapes 则是构建器执行优化时的基准形状，应设置为实际推理中最常见的形状以获得最优性能。

值得补充的是，这些形状定义与仓库内部的动态形状生成逻辑保持一致——例如 onnx_utils/export.py 中针对每个模型家族的 _gen_dummy_inp_and_dyn_shapes_* 函数会直接依据模型 config 生成 min/opt 形状（SDXL 的 min batch=2、opt batch=16；FLUX 默认 batch=1 等），_create_trt_dynamic_shapes 则将其整理为 minShapes/optShapes/maxShapes 三件套。你可以据此确认自己手写形状与官方导出的 ONNX 是否匹配。

---

四、第三步：用 demoDiffusion 跑端到端 Stable Diffusion 管线

如果你希望在目标 GPU 上，用 Model Optimizer 量化后的 UNet 完整跑通 SD/SDXL 文生图并测量延迟，可以按下面的步骤操作：

1. 克隆 demoDiffusion 示例：从 NVIDIA TensorRT 仓库的 demo/Diffusion 目录克隆一份（示例代码基于 TensorRT 10.2 分支编写）。注意首次运行时会花费较长时间构建 TRT 引擎。

2. 按 demoDiffusion 的 README 搭建管线，并先运行一个 FP16 基线 txt2img 示例：

# SDXL
python demo_txt2img_xl.py "enchanted winter forest, soft diffuse light on a snow-filled day, serene nature scene, the forest is illuminated by the snow" --negative-prompt "normal quality, low quality, worst quality, low res, blurry, nsfw, nude" --version xl-1.0 --scheduler Euler --denoising-steps 30 --seed 2946901
# 其他模型的示例命令请参考 demoDiffusion SD/SDXL 管线自带的 examples

3. 用上一步生成的 INT8 引擎替换 FP16 backbone 引擎，例如：

cp -r {YOUR_UNETXL}.plan ./engine/

重要提醒：引擎必须在同一张 GPU 上构建；且 INT8 引擎的文件名必须与 FP16 引擎的文件名一致，否则 demoDiffusion 管线无法识别替换。

4. 再次运行上述 txt2img 命令，即可对比 FP16 与 INT8 的生成图像与延迟。类似地，你也可以用 Model Optimizer 量化后的 backbone 配合 demoDiffusion 中其他扩散模型的示例，跑通各自的端到端管线。

---

五、第四步：用 DeviceModel 跑推理管线

5.1 DeviceModel 是什么

DeviceModel 是 Model Optimizer 提供的一层接口，用于像运行 torch 模型一样运行 TensorRT 引擎：它接收 torch 输入、返回 torch 输出；内部则完成"torch checkpoint → ONNX → TensorRT engine"的自动转换（对应仓库中的 modelopt/torch/_deploy/device_model.py 实现）。这使你可以把扩散管线中的 backbone 直接替换为 DeviceModel，按正常 pipeline 调用方式执行任意 prompt 的推理。

5.2 生成量化 torch checkpoint

先用 quantize.py 生成 FP8 量化 checkpoint（不导出 ONNX 也可）：

python quantize.py \
    --model {sdxl-1.0|sdxl-turbo|sd3-medium|flux-dev} \
    --format fp8 \
    --batch-size {1|2} \
    --calib-size 128 \
    --n-steps 20 \
    --quantized-torch-ckpt-save-path ./{MODEL}_fp8.pt \
    --collect-method default

5.3 使用 diffusion_trt.py 生成图像与评测延迟

然后用 diffusion_trt.py 加载量化 checkpoint 并跑文生图：

python diffusion_trt.py \
    --model {sdxl-1.0|sdxl-turbo|sd3-medium|flux-dev} \
    --prompt "A cat holding a sign that says hello world" \
    [--override-model-path /path/to/model] \
    [--restore-from ./{MODEL}_fp8.pt] \
    [--onnx-load-path {ONNX_DIR}] \
    [--trt-engine-load-path {ENGINE_DIR}] \
    [--dq-only] \
    [--torch] \
    [--save-image-as /path/to/image] \
    [--benchmark] \
    [--torch-compile] \
    [--skip-image]

脚本运行结束后会把生成的图像保存为 ./{MODEL}.png，并输出 TensorRT backbone 的推理延迟（device_model.get_latency()，见 diffusion_trt.py）。若想用 FP16/BF16 精度生成图像，去掉 --restore-from 参数直接运行即可。

各可选参数的作用（对应 diffusion_trt.py 的参数定义）：

参数	说明
--override-model-path	使用自定义模型路径，替代内置的模型 ID 映射
--restore-from	加载 Model Optimizer 量化 checkpoint 的路径
--onnx-load-path	加载已存在的 ONNX 模型（跳过 torch→ONNX 导出）
--trt-engine-load-path	加载已构建的 TensorRT engine（跳过 engine 编译，直接以 prepend_hash_to_bytes 方式装载，见 diffusion_trt.py）
--dq-only	将 ONNX 模型转换为 dq_only 模型后再构建引擎
--torch	使用纯 PyTorch 管线生成/评测（不构建 TRT），便于对照
--torch-compile	对 backbone 应用 torch.compile(mode="max-autotune")（需配合 --torch）
--torch-autocast	推理/评测时启用 torch.autocast
--benchmark	对 backbone 做独立延迟评测（--torch 模式下跑 warmup 10 次 + 评测 100 次并输出 avg/P50/P95/P99）
--save-image-as	自定义输出图像路径
--skip-image	跳过图像生成，只做引擎构建/评测
--num-inference-steps	去噪步数（默认 30，越小生成越快）

注意：使用 --trt-engine-load-path 加载已有引擎时，官方建议只加载由本管线生成的引擎（因为引擎内部与 ONNX 的输入/输出命名、SHA-256 状态哈希等绑定，见 diffusion_trt.py）。

从源码可以看到完整的底层调用链（diffusion_trt.py）：

1. 构造 deployment 配置 {"runtime": "TRT", "precision": "stronglyTyped", "onnx_opset": "17"}，经 RuntimeRegistry.get 取得 TensorRT runtime client；
2. 通过 get_onnx_bytes_and_metadata 得到 ONNX 字节流与输入/输出元数据（SDXL 模型会调用 remove_nesting 展平 added_cond_kwargs.text_embeds/time_ids 嵌套输入）；
3. 若未提供已有引擎，则 client.ir_to_compiled 将 ONNX 编译为引擎，并剥离前 32 字节的 SHA-256 哈希后写入 {model}.plan（哈希用于 runtime 维护状态，见 modelopt/torch/_deploy/_runtime/tensorrt/constants.py）；
4. 构造 DeviceModel(client, compiled_model, metadata, compilation_args, io_shapes, ignore_nesting)，将其挂载为 pipe.transformer（FLUX/SD3）或 pipe.unet（SDXL），随后直接调用 pipeline 完成图像生成。

---

六、量化效果对比示例

原文档在 Demo Images 一节给出了同一 SDXL 模型 FP16 与 INT8 主干生成的图像对比（示例图见仓库 assets 目录）：

SDXL FP16	SDXL INT8

两张图均以相同 prompt 与种子生成，用于直观评估 INT8 量化在主观画质上相对 FP16 基线的退化程度。在你自己的评测中，建议固定 prompt、seed、scheduler 与 denoising steps，仅切换 FP16/INT8 引擎，这样图像差异就只来源于量化本身。

---

七、注意事项与常见坑

汇总文档与源码中值得特别注意的几点：

1. TensorRT 版本下限：INT8 需 ≥ 9.2.0，FP8 需 ≥ 10.2.0；构建引擎的设备上必须预先完成 TensorRT 环境配置。
2. 引擎与 GPU 绑定：.plan 引擎必须在目标推理 GPU 上构建，换卡后需重新构建。
3. 文件名匹配：替换 demoDiffusion 的引擎时，量化引擎文件名必须与 FP16 引擎一致。
4. maxShapes 语义：它是张量允许的最大形状上限，扩大 batch 或分辨率前务必同步调整。
5. 校准与部署可分离：模型校准需要较多 GPU 算力，但不要求与部署目标 GPU 相同（见 examples/diffusers/README.md）。
6. 随机性：校准过程不固定随机种子，多次校准会得到不同的量化 amax 值，进而影响生成图像。若对画质敏感，可多跑几次校准并选用效果最佳的一次（详见 examples/diffusers/README.md）。
7. LoRA 融合：为获得最佳 INT8/FP8 性能，强烈建议在量化前先融合 LoRA 权重，否则可能破坏 TensorRT 对 INT8/FP8 QDQ 节点的 kernel 融合，造成性能损失。
8. FP4 导出范围：目前仅 Flux 与 SDXL 支持 FP4 ONNX 导出；SVDQuant 部署暂不支持。

至此，从"量化 + ONNX 导出"到"TensorRT 引擎构建"再到"demoDiffusion / DeviceModel 端到端推理"，一条完整的扩散模型量化部署链路已经打通。所有命令与脚本均可直接在仓库 examples/diffusers/quantization/ 目录下找到，按其参数说明即可复现。