弥合模型优化与生产部署之间的差距
本教程将带你走完一个典型的端到端量化工作流，从 PyTorch 模型开始，到导出/编译 NVIDIA TensorRT 引擎以实现真实推理加速。阅读完整文章：

模型量化：利用 NVIDIA TensorRT 将 FP8 检查点转化为高性能推理引擎

来源：https://developer.nvidia.com/blog/model-quantization-turn-fp8-checkpoints-into-high-performance-inference-engines-with-nvidia-tensorrt/

将量化检查点转换为 NVIDIA TensorRT 引擎，可以弥合模型优化与生产部署之间的差距，从而在规模上实现更快的推理、更高的吞吐量以及更高效的 GPU 利用率。

在上一篇文章（https://developer.nvidia.com/blog/model-quantization-post-training-quantization-using-nvidia-model-optimizer/）中，我们使用 NVIDIA TensorRT Model Optimizer 生成了一份高质量的 FP8 量化对比语言-图像预训练（CLIP）检查点。本文将从上次结束的地方继续，讲解如何将该检查点导出为 ONNX 格式，并将其编译为可用于生产推理的 NVIDIA TensorRT 引擎。我们还将结果 FP8 TensorRT 引擎与 FP16 基线进行性能分析，以衡量量化模型带来的真实加速效果。

图 1 展示了典型端到端量化工作流的五个阶段。这是部署量化 CLIP 模型的标准流程。量化的大语言模型（LLM）通过 TensorRT-LLM（https://docs.nvidia.com/tensorrt-llm/index.html）走不同路径，相关教程见此处（https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_ptq#exporting-checkpoints）。

工作流示意图：从 PyTorch 模型到 TensorRT 运行时推理的五个阶段：模型首先被量化为 ModelOpt 检查点，导出为带 QDQ 节点的 ONNX，然后构建成 TensorRT 引擎，最后部署用于 TensorRT 运行时推理。
图 1. 使用 ModelOpt 和 TensorRT 的端到端量化与部署工作流

将模型导出为 ONNX 格式

https://developer.nvidia.com/blog/model-quantization-turn-fp8-checkpoints-into-high-performance-inference-engines-with-nvidia-tensorrt/#export_model_to_onnx_format

第一步是将 ModelOpt 检查点导出为 ONNX。以下伪代码展示了如何使用 Modelopt 内置辅助函数，对 FP8 量化的 CLIP 检查点执行此操作（导出目标为 ONNX opset 20+，其中完全支持 FP8 QuantizeLinear/DequantizeLinear）。它会将每个权重侧的量化-反量化（Q-DQ）对折叠为仅含 FP8 存储的 DQ 链，从而显著缩小 ONNX 文件。原则上，原生 torch.onnx.export 也可以工作，但需要我们编写自定义转换脚本。

import torch
from transformers import CLIPModel, CLIPTokenizer
from transformers.models.clip.modeling_clip import CLIPAttention
import modelopt.torch.opt as mto
import modelopt.torch.quantization as mtq
from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata
from modelopt.torch.quantization.plugins.diffusion.diffusers import _QuantAttention

# 轻量包装器，将单个前向暴露给 ONNX 导出器
class TextEncoder(torch.nn.Module):
    def __init__(self, m):
        super().__init__(); self.m = m
    def forward(self, x):
        return self.m.get_text_features(x)

class ImageEncoder(torch.nn.Module):
    def __init__(self, m):
        super().__init__(); self.m = m
    def forward(self, x):
        return self.m.get_image_features(x)

def prepare_for_fp8_onnx_export(model):
    # 1) 打开 FP8 注意力融合（默认关闭，重载后会丢失）。
    # 2) 清除 CLIP 中的 float `scale` —— 导出器会因它报错。
    for _, mod in model.named_modules():
        if isinstance(mod, _QuantAttention):
            mod._disable_fp8_mha = False
        if isinstance(mod, CLIPAttention) and getattr(mod, "scale", None) is not None:
            mod.scale = None

def export(wrapper, dummy, axis_name, out_name):
    """ModelOpt 的导出器将权重上的 Q+DQ 折叠为 FP8 存储的仅 DQ 链，
       并将 TRT 自定义算子重写为原生 ONNX QDQ —— 输出可直接用于 TRT。"""
    onnx_bytes, _ = get_onnx_bytes_and_metadata(
        model=wrapper,
        dummy_input=(dummy,),
        model_name=out_name,
        dynamic_axes={axis_name: {0: "batch"}},
        onnx_opset=20,
        weights_dtype="fp16",
    )
    OnnxBytes.from_bytes(onnx_bytes).write_to_disk("./onnx_output", clean_dir=False)

# 从 ModelOpt 检查点恢复 FP8 量化后的 CLIPModel
mto.enable_huggingface_checkpointing()
mtq.QuantModuleRegistry.register({CLIPAttention: "CLIPAttention"})(_QuantAttention)
model = (
    CLIPModel.from_pretrained(modelopt_ckpt, attn_implementation="sdpa", torch_dtype=torch.float16)
    .eval().cuda()
)
prepare_for_fp8_onnx_export(model)

# 导出文本编码器到 ONNX
tok = CLIPTokenizer.from_pretrained(model_ckpt)
dummy_text = tok(["a photo of a cat"], return_tensors="pt", padding="max_length", max_length=77)["input_ids"].cuda()
export(TextEncoder(model), dummy_text, "text_input", "text_clip_fp8")

# 导出图像编码器到 ONNX
dummy_image = torch.randn(16, 3, 224, 224, dtype=torch.float16).cuda()
export(ImageEncoder(model), dummy_image, "image_input", "image_clip_fp8")

表 1. CLIP ONNX 模型大小：FP8 与 FP16 对比

表 1 比较了 FP8 ModelOpt 检查点导出的 ONNX 文件大小与原始 FP16 HuggingFace 检查点导出的 ONNX 文件大小。FP8 检查点导出的 ONNX 文件明显更小，文本编码器缩小约 34%，图像编码器缩小约 50%。请注意，缩小 ONNX 文件只是便利性，并非必要。TensorRT 会在引擎构建时将权重侧的 Q 节点折叠到 FP8 权重中。ModelOpt ONNX 导出器提前在 ONNX 侧进行折叠，以减小磁盘文件大小。

我们可以使用 NVIDIA Nsight Deep Learning Designer（https://developer.nvidia.com/nsight-dl-designer）检查导出的 ONNX 文件，这是一个高效的 ONNX 模型编辑、性能分析和 TensorRT 引擎构建工具。图 2 显示了在 Nsight Deep Learning Designer 中可视化的导出 ONNX 图的一部分。可以看到图中现在包含 QuantizeLinear/DequantizeLinear（Q/DQ）节点，标记了 FP8 边界。

通过 Nsight Deep Learning Designer 可视化的 FP8 ONNX 图中注意力 MatMul 附近的部分。两条并行的量化/反量化链（Mul, QuantizeLinear, DequantizeLinear, Cast）输入到 MatMul，随后在 FP32 中执行 Softmax，然后是一对 QuantizeLinear 和 DequantizeLinear 将结果带回 FP8。
图 2. FP8 ONNX 图中注意力 MatMul 周围的 Q/DQ 节点

在引擎构建期间，TensorRT 会将这些节点与相邻层融合，以优化推理性能。这种融合消除了不必要的量化-反量化转换，从而能够在计算中使用优化的 FP8 内核。

使用 TensorRT 对 ONNX 模型进行性能分析

https://developer.nvidia.com/blog/model-quantization-turn-fp8-checkpoints-into-high-performance-inference-engines-with-nvidia-tensorrt/#profile_onnx_model_with_tensorrt

导出 FP8 ONNX 模型后，下一步是将其传递给 TensorRT 并测量运行速度。开始之前，请确保按照本教程（https://docs.nvidia.com/deeplearning/tensorrt/latest/installing-tensorrt/installing.html#installing-tensorrt）正确下载并安装 TensorRT。准备好后，我们将使用 trtexec（TensorRT 命令行包装器，https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec）对 ONNX 模型进行基准测试，命令如下：

# 设置 TensorRT 环境
export PATH=:$PATH
export LD_LIBRARY_PATH=:$LD_LIBRARY_PATH

# 使用 trtexec 对 ONNX 模型进行基准测试
trtexec --onnx=text_clip_fp8.onnx \
        --shapes=text_input:128x77 \
        --stronglyTyped \
        --saveEngine=text_clip_fp8.plan

trtexec --onnx=image_clip_fp8.onnx \
        --shapes=image_input:128x3x224x224 \
        --stronglyTyped \
        --saveEngine=image_clip_fp8.plan

• --onnx 指定输入的 ONNX 模型，TensorRT 将从中构建引擎。
• --shapes 固定输入形状，以便 TensorRT 针对该确切大小构建优化引擎。
• --stronglyTyped 强制 TensorRT 尊重 ModelOpt 嵌入到 ONNX 图中的精度注解，确保我们的 FP8 权重和激活实际在 FP8 中执行。
• --saveEngine 将构建的 TensorRT 引擎写入磁盘，以便后续复用，可用于独立 TensorRT 推理运行时（https://docs.nvidia.com/deeplearning/tensorrt/latest/getting-started/quick-start-guide.html#using-the-tensorrt-runtime-api）或通过 NVIDIA Triton 推理服务器（https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/index.html）进行服务（参见此示例，https://github.com/NVIDIA/TensorRT/tree/main/quickstart/deploy_to_triton#step-2-set-up-triton-inference-server）。

一个注意事项：ModelOpt 的导出器将注意力缩放包装在 FP32 往返中，--stronglyTyped 会拒绝该操作（可能会看到 Float 与 Half 类型不匹配的错误）。在 trtexec 基准测试之前，我们将这些缩放常数和 Cast 操作转换回 FP16，以获得干净的强类型引擎。

# 将 FP8 ONNX 中所有 FP32 初始化器和 Cast(to=FP32) 操作重新类型为 FP16。
import numpy as np
import onnx
from onnx import TensorProto, numpy_helper, shape_inference

model = onnx.load("clip_fp8.onnx")

for init in model.graph.initializer:
    if init.data_type == TensorProto.FLOAT:
        arr = numpy_helper.to_array(init).astype(np.float16)
        init.CopyFrom(numpy_helper.from_array(arr, name=init.name))

for node in model.graph.node:
    if node.op_type == "Cast":
        to_attr = next(a for a in node.attribute if a.name == "to")
        if to_attr.i == TensorProto.FLOAT:
            to_attr.i = TensorProto.FLOAT16

model = shape_inference.infer_shapes(model, data_prop=True, check_type=False)
onnx.save(model, "clip_fp8_strongtyped.onnx")

或者，我们也可以使用 Nsight Deep Learning Designer 通过官方用户指南（https://docs.nvidia.com/nsight-dl-designer/UserGuide/index.html）中的性能分析部分对 ONNX 模型进行 TensorRT 性能分析。

我们在 NVIDIA RTX 6000 Ada GPU 上使用 TensorRT 10.16，通过 trtexec 命令运行基准测试，静态批大小为 128。报告的延迟为默认测量窗口内所有推理迭代的中位数。

注意：FP8 仅在 Ada 及更高架构（计算能力 8.9 或以上）上支持矩阵乘法（GEMM）。有关哪些数据类型在哪些 GPU 上受支持的详细分类，请参见 TensorRT 支持矩阵（https://docs.nvidia.com/deeplearning/tensorrt/support-matrix/index.html）。

柱状图比较 CLIP FP8 与 FP16 在 TensorRT 上的表现。左图：TensorRT 引擎大小，图像编码器从 588 MB 降至 294 MB，减少 50%；文本编码器从 238 MB 降至 156 MB，减少 34%。右图：中位推理延迟，图像编码器从 204 ms 降至 127 ms，加速 1.61 倍；文本编码器从 19 ms 降至 11.4 ms，加速 1.67 倍。
图 3. CLIP FP8 与 FP16：TensorRT 引擎大小和推理延迟

图 3 展示了 FP8 量化在 TensorRT 引擎大小和推理延迟两方面相对于 FP16 的优势。左图显示，图像编码器从 588 MB 缩小到 306 MB（减少 48%），文本编码器从 238 MB 缩小到 156 MB（减少 34%），合计磁盘占用几乎减半。同样的节省也体现在推理时的 GPU VRAM 使用上，因为更小的引擎需要更少的内存来加载和运行。右图，延迟方面同样令人信服。图像编码器从 166.2 ms 降至 119.8 ms，文本编码器从 13.2 ms 降至 9.1 ms，图像端加速 1.39 倍，文本端加速 1.45 倍。

使用 Nsight Deep Learning Designer 对 FP16 和 FP8 CLIP 图像编码器进行的性能分析比较，显示 FP8 将 GEMM 延迟大致减半，消除了单独的融合条，并将大部分网络从 FP16 执行转变为 FP8 执行。
图 4. 使用 Nsight Deep Learning Designer 对 CLIP 图像编码器进行逐层性能分析：FP16（左）与 FP8（右）

FP8 加速究竟来自哪里？除了 trtexec 报告的原始数字外，Nsight Deep Learning Designer 提供了更丰富的可视化分解，为我们提供清晰答案。图 4 将 FP16 和 FP8 图像编码器的性能分析并排放置，三个差异立即显现。

1. GEMM 条从大约 1.8 ms 下降到 0.84 ms，主流 matmul 层实现了超过 2 倍的加速，这归功于 NVIDIA RTX 6000 Ada GPU 的 FP8 Tensor Core 内核。
2. FP16 性能分析中可见的“融合”层类别在 FP8 性能分析中消失了，因为 TensorRT 现在将整个注意力模块通过专门的 FP8 MHA 内核路由，产生了更精简的执行路径。
3. 精度饼图从大部分橙色（FP16）转变为大部分紫色（FP8）。

这些信号证实，我们的量化权重和激活正在 FP8 Tensor Core 上运行，这正是 FP8 收益的来源——在每个 matmul 密集型步骤中实现更高的计算吞吐量和更低的内存带宽使用。

TensorRT 中的量化工作原理

https://developer.nvidia.com/blog/model-quantization-turn-fp8-checkpoints-into-high-performance-inference-engines-with-nvidia-tensorrt/#how_quantization_works_in_tensorrt

导入 ONNX 模型时，TensorRT 会查找 QuantizeLinear / DequantizeLinear（Q/DQ）节点，这些节点标记了图中张量在全精度和低精度数据类型（如 FP8）之间转换的位置。在内部，TensorRT 要求每个可量化层的每个输入上有一对 Q/DQ 层。在引擎构建时，优化器将这些 Q/DQ 节点融合到相邻层中，并用直接操作低精度张量的专用内核替换原始层。这消除了量化-反量化的往返转换，使引擎能够以更高的计算吞吐量和更低的内存带宽执行。

图 5 展示了 FP8 GEMM 的这种转换。在导出的 ONNX 中，激活 (x_f) 和权重张量都包装在一对 QuantizeLinear/DequantizeLinear 中，经过 TensorRT 优化器融合后，剩下的是一个单一的 FP8 GEMM 内核，它直接接收 FP8 量化的激活和预先存储的 FP8 权重张量。

示意图展示了 TensorRT 如何围绕 GEMM 融合 FP8 QuantizeLinear 和 DequantizeLinear 节点，将导出 ONNX 中的原始 Q/DQ 节点折叠为单个带有预存储 FP8 权重的 FP8 GEMM 内核。
图 5. TensorRT 中的 FP8 GEMM Q/DQ 融合

有关 TensorRT 量化机制的深入探讨，请参阅文档（https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/work-quantized-types.html#working-with-quantized-types）。

开始使用

https://developer.nvidia.com/blog/model-quantization-turn-fp8-checkpoints-into-high-performance-inference-engines-with-nvidia-tensorrt/#get_started

在本文中，我们介绍了部署量化模型的完整 ModelOpt → ONNX → TensorRT 工作流。我们导出了带 Q/DQ 节点的 CLIP 检查点到 ONNX，构建了 TensorRT 引擎，并使用 trtexec 和 Nsight Deep Learning Designer 对 FP16 基线进行了基准测试。结果表明，在 RTX 6000 Ada GPU 上，FP8 量化相比原始 FP16 模型在速度和内存占用方面均有显著提升。我们还简要概述了 TensorRT 如何在构建时通过将 Q/DQ 节点融合为专用的低精度内核来实现这些增益。

尝试 NVIDIA Model Optimizer（https://github.com/NVIDIA/Model-Optimizer）和 NVIDIA TensorRT（https://developer.nvidia.com/tensorrt），探索模型量化带来的效率提升。