使用量化类型#

量化简介#

TensorRT 支持使用低精度类型来表示量化的浮点值。量化方案是对称量化——量化值以有符号 INT8、FP8E4M3 (简称 FP8)、有符号 INT4 或 FP4E2M1 (简称 FP4) 表示,从量化值到非量化值的转换只是一个乘法。反向转换中,量化使用倒数比例,然后进行钳位和舍入(对于整数)或强制转换(对于 FP8 和 FP4)。

TensorRT 量化 INT8、FP8 和 FP4 的激活值和权重。INT4 支持仅权重量化。

量化工作流程#

创建量化网络有两种工作流程

训练后量化 (PTQ)

在训练网络后导出比例因子。TensorRT 提供了一个名为校准的 PTQ 工作流程。它测量每个激活张量内激活值的分布,当网络在代表性输入数据上执行时,然后使用该分布来估计每个张量的比例值。

量化感知训练 (QAT)

在训练期间使用伪量化计算比例因子,模拟量化和反量化过程。这允许训练过程补偿量化和反量化操作的影响。

TensorRT 的 量化工具包 是一个 PyTorch 库,可帮助生成 TensorRT 可以优化的 QAT 模型。该工具包的 PTQ 配方也可以在 PyTorch 中执行 PTQ 并导出到 ONNX。

显式与隐式量化#

注意

隐式量化已弃用。您应该使用 TensorRT 的 量化工具包 来创建具有显式量化的模型。

量化网络可以通过两种(互斥的)方式处理:使用隐式或显式量化。两种处理模式之间的主要区别在于您是否需要显式控制量化,还是让 TensorRT 构建器选择要量化的操作和张量(隐式)。以下部分提供更多详细信息。仅当为 INT8 量化时才支持隐式量化。它不能与强类型一起使用(因为类型不是自动调整的,并且将激活值转换为 INT8 以及从 INT8 转换回来的唯一方法是通过 Quantize (Q) 和 Dequantize (DQ) 运算符)。

当网络具有 QuantizeLayerDequantizeLayer 层时,TensorRT 使用显式量化模式。当网络中没有 QuantizeLayerDequantizeLayer 层,并且在构建器配置中启用了 INT8 时,TensorRT 使用隐式量化模式。隐式量化模式仅支持 INT8。

隐式量化网络中,每个激活张量量化候选者都具有一个关联的比例,该比例通过校准过程推导出来或由 API 函数 setDynamicRange 分配。如果 TensorRT 决定量化张量,它将使用此比例。

在处理隐式量化网络时,TensorRT 在应用图优化时将模型视为浮点模型,并机会性地使用 INT8 来优化层执行时间。如果某层在 INT8 中运行速度更快,并且在其数据输入和输出上分配了量化比例,则会将具有 INT8 精度的内核分配给该层。否则,将分配高精度浮点(FP32、FP16 或 BF16)内核。如果需要高精度浮点以保证精度,但会牺牲性能,则可以使用 API Layer::setOutputType Layer::setPrecision 来指定。

显式量化网络中,量化和反量化操作由图中的 IQuantizeLayer (C++, Python) 和 IDequantizeLayer (C++, Python) 节点显式表示 - 这些节点在下文中将称为 Q/DQ 节点。与隐式量化相比,显式形式精确地指定了执行转换为量化类型和从量化类型转换的位置,并且优化器将仅执行模型语义指示的转换为量化类型和从量化类型转换,即使

  • 添加额外的转换可能会提高层精度(例如,选择 FP16 内核实现而不是量化类型实现)。

  • 添加或删除转换会产生更快的引擎(例如,选择量化类型内核实现来执行指定为具有高精度的层,反之亦然)。

ONNX 使用显式量化表示:当 PyTorch 或 TensorFlow 中的模型导出到 ONNX 时,框架图中的每个伪量化操作都导出为 Q,然后是 DQ。由于 TensorRT 保留了这些层的语义,因此用户可以期望精度非常接近框架中看到的精度。虽然优化保留了量化和反量化运算符的算术语义,但它们可能会更改模型中浮点运算的顺序,因此结果不会是按位相同的。

TensorRT 的 PTQ 功能生成具有隐式量化的校准缓存。相比之下,在深度学习框架中执行 QAT 或 PTQ,然后导出到 ONNX 将导致显式量化模型。

隐式与显式量化#

隐式量化(已弃用)

显式量化

支持的量化数据类型

INT8

INT8、FP8、INT4、FP4

用户对精度的控制

全局构建器标志和每层精度 API。

直接编码在模型中。

API

  • 模型 + 比例(动态范围 API)

  • 模型 + 校准数据

带有 Q/DQ 层的模型。

量化比例

  • 权重
    • 由 TensorRT 设置(内部)

    • 每通道量化

    • INT8 范围 [-127, 127]

  • 激活值
    • 由校准设置或由用户指定

    • 每张量量化

    • INT8 范围 [-128, 127]

  • 权重和激活值
    • 使用 Q/DQ ONNX 运算符指定

    • INT8 范围 [-128, 127]

    • FP8 范围:[-448, 448]

    • INT4 范围:[-8, 7]

    • FP4 范围:[-6, 6]

  • 激活值使用每张量量化。

  • 权重使用每张量量化、每通道量化或块量化。

有关量化的更多背景知识,请参阅以下论文

量化方案#

给定比例 \(\text{s}\),我们可以将量化和反量化操作表示如下

\(x_{q}=quantize\left(x, s \right)=roundWithTiesToEven(clip(\frac{x}{s}, -128,127))\)

其中

  • \(\text{x}\) 是要量化的高精度浮点值。

  • \(x_{q}\) 是范围在 [-128,127] 内的量化 INT8 值。有关更多信息,请参阅显式与隐式量化部分。

  • \(\text{roundWithTiesToEven}\)此处 描述。

\(\text{x}=dequantize\left(x_{q}, s\right)=x_{q}\ast s\)

在显式量化中,您负责选择所有比例。在隐式量化模式下,您可以使用 TensorRT 的校准算法之一配置或确定激活比例。TensorRT 根据以下公式计算权重比例

\(\text{s}=\frac{max(abs(\mathrm{x}_{min}^{ch}),abs(\mathrm{x}_{max}^{ch}))}{127}\)

其中 \({x}_{min}^{ch}\)\({x}_{max}^{ch}\) 是权重张量的通道 \(\text{ch}\) 的浮点最小值和最大值。

使用 FP8 时,仅支持显式量化;因此,您负责量化比例的值。

\(x_{q}=quantize\left(x, s \right)=castToFp8(clip(\frac{x}{s}, -448,448))\)

其中

  • \(\text{x}\) 是要量化的高精度浮点值。

  • \(x_{q}\) 是范围在 [-448, 448] 内的量化 FP8E4M3 值。

  • \(\text{s}\) 是使用 16 位或 32 位浮点表示的量化比例。

  • \(\text{castToFp8}\) 舍入到 FP8E4M3 中可表示的最接近的值,关系舍入到偶数,如 此处 所述。

\(\text{x}=dequantize\left(x_{q}, s\right)=x_{q}\ast s\)

不允许在同一网络中使用 FP8 和 INT8。

使用 INT4 时,仅支持显式量化,因此您负责量化比例的值。

\(x_{q}=quantize\left(x, s \right)=roundWithTiesToEven(clip(\frac{x}{s}, -8,7))\)

其中

  • \(\text{x}\) 是要量化的高精度浮点值。

  • \(x_{q}\) 是范围在 [-8, 7] 内的量化 INT4 值。

  • \(\text{s}\) 是使用 16 位或 32 位浮点表示的量化比例。

  • \(\text{roundWithTiesToEven}\)此处 描述。

\(\text{x}=dequantize\left(x_{q}, s\right)=x_{q}\ast s\)

TensorRT 仅支持用于权重量化的 INT4(Q/DQ 层放置建议)。

使用 FP4 时,仅支持显式量化,因此您负责量化比例的值。

\(x_{q}=quantize\left(x, s \right)=castToFp4(clip(\frac{x}{s}, -6,6))\)

其中

  • \(\text{x}\) 是要量化的高精度浮点值。

  • \(x_{q}\) 是范围在 [-6, 6] 内的量化 FP4 值。

  • \(\text{s}\) 是使用 16 位或 32 位浮点表示的量化比例。

  • \(\text{castToFp4}\) 舍入到 FP4E2M1 中可表示的最接近的值,关系舍入到偶数,如 此处 所述。

\(\text{x}=dequantize\left(x_{q}, s\right)=x_{q}\ast s\)

当量化 FP4 激活值时,建议使用 动态量化

量化模式#

支持三种量化比例粒度

  1. 每张量量化:单个比例值(标量)用于缩放整个张量。

  2. 每通道量化:比例张量沿给定轴广播 - 对于卷积神经网络,这通常是通道轴。

  3. 块量化:张量沿单个维度划分为固定大小的 1 维块。为每个块定义一个比例因子。

量化比例必须包含所有正高精度浮点系数(FP32、FP16 或 BF16)。舍入方法是 舍入到最近的偶数 并钳位到有效范围,对于 INT8 为 [-128, 127],对于 FP8 为 [-448, 448],对于 INT4 为 [-8, 7]

使用显式量化时,激活值只能使用每张量量化进行量化。权重可以在任何量化模式下量化。

在隐式量化中,权重由 TensorRT 在引擎优化期间量化,并且仅使用每通道量化。TensorRT 量化卷积、反卷积、全连接层和 MatMul 的权重,其中第二个输入是常数,并且两个输入矩阵都是 2D。

当将每通道量化与卷积一起使用时,量化轴必须是输出通道轴。例如,当使用 KCRS 表示法描述 2D 卷积的权重时,K 是输出通道轴,权重量化可以描述为

For each k in K:
    For each c in C:
        For each r in R:
            For each s in S:
                output[k,c,r,s] := clamp(round(input[k,c,r,s] / scale[k]))

比例是系数向量,并且必须具有与量化轴相同的大小。

反量化的执行方式类似,只是逐点运算定义为

output[k,c,r,s] := input[k,c,r,s] * scale[k]

块量化

在块量化中,元素被分组为 1D 块,块中的所有元素共享一个公共比例因子。块量化支持最多 3 个维度的输入。

INT4 块量化支持仅权重量化 (WoQ)。

FP4 块量化支持权重和激活值。为了最大限度地减少量化误差,建议对激活值使用动态量化。

当使用块量化时,比例张量维度等于数据张量维度,除了执行分块的一个维度(分块轴)。例如,给定一个 2-D RS 权重输入,R(维度 0)作为分块轴,B 作为块大小,则分块轴中的比例根据块大小重复,可以这样描述

For each r in R:
    For each s in S:
        output[r,s] = clamp(round(input[r,s] / scale[r//B, s]))

比例是维度为 (R//B, S) 的 2D 系数数组。

反量化的执行方式类似,只是逐点运算定义为

output[r,s] = input[r,s] * scale[r//B, s]

设置动态范围#

动态范围 API 仅适用于 INT8 量化。

TensorRT 提供了直接设置动态范围(必须由量化张量表示)的 API,以支持在 TensorRT 外部计算这些值的隐式量化。

该 API 允许使用最小值和最大值设置张量的动态范围。由于 TensorRT 当前仅支持对称范围,因此比例是使用 max(abs(min_float), abs(max_float)) 计算的。请注意,当 abs(min_float) != abs(max_float) 时,TensorRT 使用的动态范围大于配置的动态范围,这可能会增加舍入误差。

您可以按如下方式设置张量的动态范围

1tensor->setDynamicRange(min_float, max_float);
1tensor.dynamic_range = (min_float, max_float)

sampleINT8API 说明了这些 API 在 C++ 中的用法。

使用校准的训练后量化#

注意

本节介绍已弃用的 API。建议使用显式量化。

校准仅适用于 INT8 量化。

在训练后量化中,TensorRT 计算网络中每个张量的比例值。此过程称为校准,需要您提供代表性输入数据,TensorRT 在这些数据上运行网络以收集每个激活张量的统计信息。

所需的输入数据量取决于应用程序,但实验表明,对于校准 ImageNet 分类网络,大约 500 张图像就足够了。

给定激活张量的统计信息,确定最佳比例值并非精确科学 - 它需要在量化表示中的两个误差源之间取得平衡:离散化误差(随着每个量化值表示的范围变得更大而增加)和截断误差(值被钳位到可表示范围的限制)。因此,TensorRT 提供了多个校准器,它们以不同的方式计算比例。较旧的校准器还执行 GPU 的层融合,以在校准之前优化掉不需要的张量。当使用 DLA 时,这可能会出现问题,因为融合模式可能不同,并且可以使用 kCALIBRATE_BEFORE_FUSION 量化标志覆盖。

校准批大小也会影响 IInt8EntropyCalibrator2IInt8EntropyCalibrator截断误差。例如,使用多个小批量的校准数据进行校准可能会导致直方图分辨率降低和比例值不佳。对于每个校准步骤,TensorRT 都会更新每个激活张量的直方图分布。假设它在激活张量中遇到一个值,该值大于当前直方图最大值。在这种情况下,直方图范围会增加 2 的幂次方,以适应新的最大值。除非直方图在最后一个校准步骤中重新分配,否则此方法效果良好,从而导致最终直方图具有一半的空 bin。这样的直方图可能会产生较差的校准比例。这也使得校准容易受到校准批次顺序的影响;不同的顺序可能会在不同的点增加直方图大小,从而产生略有不同的校准比例。为了避免此问题,请尽可能使用单个大批量进行校准,并确保校准批次是充分随机化的且具有相似的分布。

IInt8EntropyCalibrator2

熵校准选择张量的比例因子来优化量化张量的信息论内容,并且通常抑制分布中的异常值。这是当前推荐的熵校准器,并且是 DLA 所必需的。默认情况下,校准发生在层融合之前。校准批大小可能会影响最终结果。建议用于基于 CNN 的网络。

IInt8MinMaxCalibrator

此校准器使用激活分布的整个范围来确定比例因子。它更适用于 NLP 任务。默认情况下,校准发生在层融合之前。建议用于 NVIDIA BERT 等网络(Google 官方实现的优化版本)。

IInt8EntropyCalibrator

这是原始的熵校准器。它比 LegacyCalibrator 更简单,并且通常产生更好的结果。校准批大小可能会影响最终结果。默认情况下,校准发生在层融合之后。

IInt8LegacyCalibrator

此校准器用于与 TensorRT 2.0 EA 兼容。它需要用户参数化,并且是其他校准器产生不良结果时的后备选项。默认情况下,校准发生在层融合之后。您可以自定义此校准器以实现百分位数最大值。例如,对于 NVIDIA BERT 和 NeMo ASR 模型 QuartzNet,观察到 99.99% 百分位数最大值具有最佳精度。

构建 INT8 引擎时,构建器执行以下步骤

  1. 构建一个 32 位引擎,在校准集上运行它,并记录每个张量的激活值分布直方图。

  2. 从直方图构建一个校准表,为每个张量提供一个比例值。

  3. 从校准表和网络定义构建 INT8 引擎。

校准可能很慢;因此,步骤 2 的输出(校准表)可以缓存和重用。这在给定平台上多次构建同一网络时很有用,并且所有校准器都支持此功能。

在运行校准之前,TensorRT 查询校准器实现,以查看它是否可以访问缓存表。如果是,则直接进行步骤 3。缓存数据作为指针和长度传递。

校准缓存数据可以在不同设备之间移植,只要校准发生在层融合之前。具体来说,当使用 IInt8EntropyCalibrator2IInt8MinMaxCalibrator 校准器或设置了 QuantizationFlag::kCALIBRATE_BEFORE_FUSION 时,校准缓存是可移植的。例如,这可以通过在具有独立 GPU 的机器上构建校准表,然后在嵌入式平台上重用它来简化工作流程。不能保证跨平台或设备融合相同,因此在层融合后校准可能不会产生可移植的校准缓存。通常,校准缓存不能跨 TensorRT 版本移植。

TensorRT 使用对称量化,量化比例是使用权重张量中找到的最大绝对值计算的。对于卷积、反卷积和全连接权重,比例是每通道的。

注意

当构建器配置为使用 INT8 I/O 时,TensorRT 仍然期望校准数据为 FP32。您可以通过将 INT8 I/O 强制转换为 FP32 精度来创建 FP32 校准数据。此外,FP32 强制转换校准数据应在 [-128.0F, 127.0F] 范围内,并转换为 INT8 数据而不会有任何精度损失。

INT8 校准可以与动态范围 API 一起使用。手动设置动态范围会覆盖从 INT8 校准生成的动态范围。

注意

校准是确定性的 - 也就是说,如果您在同一设备上以相同的顺序为 TensorRT 提供相同的校准输入,则生成的比例在不同的运行中将是相同的。当使用相同的设备以相同的批大小生成校准缓存时,校准缓存中的数据将是按位相同的,前提是提供了相同的校准输入。当使用不同的设备、批大小或校准输入生成校准缓存时,不能保证校准缓存中的确切数据是按位相同的。

使用 C++ 进行 INT8 校准#

要为 TensorRT 提供校准数据,必须实现 IInt8Calibrator 接口。

构建器按如下方式调用校准器

  1. 首先,它查询接口以获取批大小,并调用 getBatchSize() 以确定要期望的输入批大小。

  2. 然后,它重复调用 getBatch() 以获取输入批次。批次必须正好是 getBatchSize() 的批大小。当没有更多批次时,getBatch() 必须返回 false

在您实现校准器后,您可以配置构建器以使用它

config->setInt8Calibrator(calibrator.get());

实现 writeCalibrationCache()readCalibrationCache() 方法以缓存校准表。

使用 Python 进行校准#

以下步骤说明了如何使用 Python API 创建 INT8 校准器对象。

  1. 导入 TensorRT。

    import tensorrt as trt
    
  2. 与测试/验证数据集类似,使用一组输入文件作为校准数据集。确保校准文件代表整体推理数据文件。为了使 TensorRT 使用校准文件,您必须创建一个 batchstream 对象。batchstream 对象用于配置校准器。

    NUM_IMAGES_PER_BATCH = 5
    batchstream = ImageBatchStream(NUM_IMAGES_PER_BATCH, calibration_files)
    
  3. 使用输入节点名称和批流创建 Int8_calibrator 对象。

    Int8_calibrator = EntropyCalibrator(["input_node_name"], batchstream)
    
  4. 设置 INT8 模式和 INT8 校准器。

    config.set_flag(trt.BuilderFlag.INT8)
    config.int8_calibrator = Int8_calibrator
    

量化噪声降低#

对于具有隐式量化的网络,TensorRT 尝试通过强制网络输出附近的一些层以 FP32 运行来降低输出中的量化噪声,即使 INT8 实现可用。

启发式方法尝试通过对多个量化值求和来确保 INT8 量化平滑。被视为“平滑层”的层是卷积、反卷积、全连接层或矩阵乘法,然后到达网络输出。例如,如果网络由一系列(卷积 + 激活 + shuffle)子图组成,并且网络输出具有 FP32 类型,则最后一个卷积将输出 FP32 精度,即使允许 INT8 并且速度更快。

启发式方法不适用于以下情况

  • 网络输出具有 INT8 类型。

  • 从最后一个平滑层到输出(包括最后一个平滑层)的路径上的操作受到 ILayer::setOutputTypeILayer::setPrecision 的约束,以输出 INT8。

  • 没有通往输出的平滑层路径,或者该路径具有介入的插件层。

  • 网络使用显式量化。

显式量化#

当 TensorRT 检测到网络中存在 Q/DQ 层时,它会使用显式精度处理逻辑构建引擎,并且不需要精度控制构建标志。

在显式量化中,网络表示与量化数据类型之间的转换是显式的;因此,INT8 和 FP8 不得用作类型约束。

对于强类型网络,既不需要也不允许构建器标志。

量化权重#

Q/DQ 模型的权重可以使用高精度数据类型(FP32、FP16 或 BF16)或低精度量化类型(INT8、FP8、INT4、FP4)指定。当 TensorRT 构建引擎时,高精度权重使用 IQuantizeLayer 比例进行量化,该比例作用于权重。量化的(低精度)权重存储在引擎计划文件中。当使用预量化权重(低精度)时,权重和使用权重的线性运算符之间需要 IDequantizeLayer

INT4 和 FP4 量化权重通过每个字节打包两个元素来存储。第一个元素存储在 4 个最低有效位中,第二个元素存储在 4 个最高有效位中,如下图所示。

4-bit packing (logical tensor on the left; physical layout on the right).

下图展示了打包 (2, 3) 4 位张量的示例。

An example packed 4-bit (2, 3) tensor.

ONNX 支持#

当在 PyTorch 或 TensorFlow 中使用量化感知训练 (QAT) 训练的模型导出到 ONNX 时,框架图中的每个伪量化操作都导出为一对 QuantizeLinearDequantizeLinear ONNX 运算符。当 TensorRT 导入 ONNX 模型时,ONNX QuantizeLinear 运算符作为 IQuantizeLayer 实例导入,而 ONNX DequantizeLinear 运算符作为 IDequantizeLayer 实例导入。

ONNX 在 opset 10 中引入了对 QuantizeLinearDequantizeLinear 的支持,并在 opset 13 中添加了 quantization-axis 属性(按通道量化需要)。PyTorch 1.8 引入了使用 opset 13 将 PyTorch 模型导出到 ONNX 的支持。

ONNX opset 19 添加了四种 FP8 格式,其中 TensorRT 支持 E4M3FN(在 ONNX 运算符架构中也称为 tensor (float8e4m3fn))。最新的 Pytorch 版本 (Pytorch 2.0) 不支持 FP8 格式,也不支持使用 opset 19 导出到 ONNX。

为了弥合差距,TransformerEngine 将其 FP 量化函数导出为属于 “trt” 域的自定义 ONNX Q/DQ 运算符(TRT_FP8 QuantizeLinearTRT_FP8 DequantizeLinear)。TensorRT 可以解析自定义运算符和标准 opset 19 Q/DQ 运算符;但是,需要注意的是,TensorRT 并非完全支持 opset 19。诸如 ONNX Runtime 之类的其他工具无法解析自定义运算符。ONNX opset 21 添加了对 INT4 数据类型和块量化的支持。ONNX opset 23 添加了对 FP4E2M1 类型的支持。

警告

ONNX GEMM 运算符是可以按通道量化的示例。PyTorch torch.nn.Linear 层导出为具有 (K, C) 权重布局且启用了 transB GEMM 属性的 ONNX GEMM 运算符(这会在执行 GEMM 运算之前转置权重)。另一方面,TensorFlow 在 ONNX 导出之前预先转置权重 (C, K)

  • PyTorch: \(y=xW^{T}\)

  • TensorFlow: \(y=xW\)

因此,TensorRT 转置 PyTorch 权重。TensorRT 在转置权重之前对其进行量化,因此源自从 PyTorch 导出的 ONNX QAT 模型的 GEMM 层使用维度 0 进行按通道量化(轴 K = 0),而源自 TensorFlow 的模型使用维度 1(轴 K = 1)。

TensorRT 不支持使用 INT8/FP8 量化运算符的预量化 ONNX 模型。具体来说,以下 ONNX 量化运算符受支持,并且如果在 TensorRT 导入 ONNX 模型时遇到它们,则会生成导入错误

TensorRT 处理 Q/DQ 网络#

当 TensorRT 在 Q/DQ 模式下优化网络时,优化过程仅限于不更改网络算术正确性的优化。位级精度很少可能实现,因为浮点运算的顺序会产生不同的结果(例如,将 \(\text{a}\ast s+b\ast s\) 重写为 \(\left(a+b \right)\ast s\) 是一种有效的优化)。允许这些差异是后端优化的一般基础,这也适用于将带有 Q/DQ 层的图转换为使用量化运算。

Q/DQ 层控制网络的计算和数据精度。IQuantizeLayer 实例通过采用量化将高精度浮点张量转换为量化张量,而 IDequantizeLayer 实例使用反量化将量化张量转换为高精度浮点张量。TensorRT 期望在可量化层的每个输入上都有一对 Q/DQ 层。可量化层是可以与 IQuantizeLayerIDequantizeLayer 实例融合以转换为量化层的深度学习层。当 TensorRT 执行这些融合时,它会将可量化层替换为使用适用于量化类型的计算操作在量化数据上运行的量化层。

对于本章中使用的图表,绿色表示低精度(量化),蓝色表示高精度。箭头表示网络激活张量,正方形表示网络层。

A quantizable ``AveragePool`` layer (in blue) is fused with DQ and Q layers. All three layers are replaced by a quantized ``AveragePool`` layer (in green).

在网络优化期间,TensorRT 在 Q/DQ 传播中移动 Q/DQ 层。传播的目标是最大化可以在低精度下处理的图的比例。因此,TensorRT 向后传播 Q 节点(量化尽早发生)并向前传播 DQ 节点(因此反量化尽可能晚发生)。Q 层可以与与量化交换的层交换位置,而 DQ 层可以与与反量化交换的层交换位置。

如果 \(\text{Q}\left(Op\left(x \right) \right)==\text{Op}\left(Q\left(x \right) \right)\),则层 \(\text{Op}\) 与量化交换

类似地,如果 \(\text{Op}\left(DQ\left(x \right) \right)==\text{DQ}\left(Op\left(x \right) \right)\),则层 \(\text{Op}\) 与反量化交换

下图说明了 DQ 向前传播和 Q 向后传播。这些是模型的合法重写,因为最大池化具有 INT8 实现并且与 DQ 和 Q 交换。

An illustration depicting a DQ forward-propagation and Q backward-propagation

为了理解最大池化交换,让我们看一下应用于某些任意输入的最大池化操作的输出。最大池化应用于输入系数组,并输出具有最大值的系数。对于由系数 \(\left\{x_{0}..x_{m} \right\}\) 组成的组 i

\(output_{i}:=max\left( \left\{ x_{0},x_{1},...x_{m} \right\} \right)=max\left( \left\{max\left( \left\{ max\left( \left\{ x_{0},x_{1} \right\} \right),x_{2} \right\} \right),...x_{3} \right\} \right)\)

因此,足以查看两个任意系数而不会失去一般性 (WLOG):\(x_{j}=max\left( \left\{ x_{j},x_{k} \right\} \right)for x_{j}\ge x_{k}\)

对于量化函数 \(\text{Q}\left( a,scale,x_{max},x_{min} \right):=truncate\left( round\frac{a}{scale} \right),x_{max},x_{min} scale\gt 0\),请注意(不提供证明并使用简化的符号):\(\text{Q}\left( x_{j},scale \right)\ge \text{Q}\left( x_{k},scale \right)for x_{j}\ge x_{k}\)

因此:\(\text{max}\left( \left\{ \text{Q}\left( x_{j},scale \right),\text{Q}\left( x_{k},scale \right) \right\} \right)=\text{Q}\left( x_{j},scale \right) for x_{j}\ge x_{k}\)

但是,根据定义:\(\text{Q}\left( max\left( \left\{ x_{j},x_{k} \right\} \right),scale \right)=\text{Q}\left( x_{j},scale \right) for x_{j}\ge x_{k}\)

函数 \(\text{max}\) 与量化交换,最大池化也是如此。

类似地,对于反量化,函数 \(\text{DQ}\left( a,scale \right):=a\ast scale\)\(\text{scale}\gt 0\) 可以证明:\(\text{max}\left( \left\{ \text{DQ}\left(x_{j},scale \right),\text{DQ}\left( x_{k},scale \right) \right\} \right)=\text{DQ}\left( x_{j},scale \right)=\text{DQ}\left( \text{max}\left( \left\{ x_{j},x_{k} \right\} \right),scale \right) for x_{j}\ge x_{k}\)

可量化层和交换层在处理方式上存在差异。这两个层都可以在 INT8/FP8 中计算,但可量化层也与 DQ 输入和 Q 输出层融合。例如,AveragePooling 层(可量化)不与 Q 或 DQ 交换,因此使用 Q/DQ 融合对其进行量化,如第一个图所示。这与最大池化(交换)的量化方式形成对比。

仅权重化#

仅权重化 (WoQ) 是一种优化,当内存带宽限制 GEMM 运算的性能或 GPU 内存稀缺时,此优化非常有用。在 WoQ 中,GEMM 权重被量化为 INT4 精度,而 GEMM 输入数据和计算操作保持高精度。TensorRT 的 WoQ 内核从内存中读取 4 位权重,并在高精度执行点积之前对其进行反量化。

Weight-only Quantization (WoQ)

WoQ 仅适用于带有 GEMM 层的 INT4 块量化。GEMM 数据输入以高精度(FP32、FP16、BF16)指定,并且权重像往常一样使用 Q/DQ 量化。TensorRT 创建一个具有 INT4 权重和高精度 GEMM 运算的引擎。引擎读取低精度权重,并在高精度执行 GEMM 运算之前对其进行反量化。

Q/DQ 层放置建议#

Q/DQ 层在网络中的放置会影响性能和准确性。激进的量化可能会由于量化引入的误差而导致模型准确性下降。但是量化还可以减少延迟。此处列出了一些在网络中放置 Q/DQ 层的建议。

请注意,较旧的设备可能没有所有层的低精度内核实现,并且在构建引擎时可能会遇到找不到任何实现的错误。要解决此问题,请删除量化失败层的 Q/DQ 节点。

量化加权运算(卷积、转置卷积和 GEMM)的所有输入。 量化权重和激活可以减少带宽需求,并启用 INT8 计算来加速带宽受限和计算受限的层。

Two examples of how TensorRT fuses convolutional layers. On the left, only the input is quantized. On the right, both the input and output are quantized.

默认情况下,不要量化加权运算的输出。 有时保留更高精度的反量化输出很有用。例如,如果线性运算后跟一个激活函数(在下图中为 SiLU),则它需要更高精度的输入才能产生可接受的准确性。

Example of a linear operation followed by an activation function

不要在训练框架中模拟批归一化和 ReLU 融合,因为 TensorRT 优化保证保留这些运算的算术语义。

Batch normalization is fused with convolution and ReLU while keeping the same execution order defined in the pre-fusion network. There is no need to simulate BN-folding in the training network.

量化跳跃连接中的残差输入。 TensorRT 可以融合加权层之后的逐元素加法,这对于具有像 ResNet 和 EfficientNet 这样的跳跃连接的模型很有用。逐元素加法层的第一输入的精度决定了融合输出的精度。

例如,在下图中,\(x_{f^{}}^{1}\) 的精度是浮点数,因此融合卷积的输出被限制为浮点数,并且尾随的 Q 层无法与卷积融合。

The precision of :math:`x_{f^{}}^{1}` is a floating point, so the output of the fused convolution is limited to the floating point, and the trailing Q-layer cannot be fused with the convolution.

相反,当 \(x_{f^{}}^{1}\) 量化为 INT8 时,如下下图所示,融合卷积的输出也是 INT8,并且尾随的 Q 层与卷积融合。

When :math:`x_{f^{}}^{1}` is quantized to INT8, the output of the fused convolution is also INT8, and the trailing Q-layer is fused with the convolution.

为了获得额外的性能,尝试量化不与 Q/DQ 交换的层。目前,具有 INT8 输入的非加权层也需要 INT8 输出,因此请同时量化输入和输出。

An example of quantizing a quantizable operation. An element-wise addition is fused with the input DQs and the output Q.

如果 TensorRT 无法将运算与周围的 Q/DQ 层融合,则性能可能会下降,因此在添加 Q/DQ 节点时要保守,并牢记准确性和 TensorRT 性能进行实验

下图显示了由于额外的 Q/DQ 运算而可能导致的次优融合(突出显示的浅绿色背景矩形)。卷积与逐元素加法分开融合,因为 Q/DQ 对相互包围。

An example of suboptimal quantization fusions: contrast the suboptimal fusion in A and the optimal fusion in B. The extra pair of Q/DQ operations (highlighted with a glowing green border) forces the separation of the convolution from the element-wise addition.

对激活使用按张量量化,对权重使用按通道量化。 经验证,此配置可以带来最佳的量化准确性。

您可以通过启用 FP16 来进一步优化引擎延迟。TensorRT 会尽可能尝试使用 FP16 而不是 FP32(目前并非所有层类型都支持此功能)。

Q/DQ 限制#

TensorRT 执行的一些 Q/DQ 图重写优化会比较两个或多个 Q/DQ 层之间的量化比例值,并且仅当比较的量化比例相等时才执行图重写。当重新拟合可重新拟合的 TensorRT 引擎时,可以为 Q/DQ 节点的比例分配新值。在 Q/DQ 引擎的重新拟合操作期间,TensorRT 检查参与比例相关优化的 Q/DQ 层是否被分配了破坏重写优化的新值,如果为真,则抛出异常。

An example showing scales of Q1 and Q2 are compared for equality, and if equal, they are allowed to propagate backward. If the engine is refitted with new values for Q1 and Q2 such that Q1 != Q2, then an exception aborts the refitting process.

Q/DQ 与插件的交互#

插件通过允许用自定义和专有实现替换一组层来扩展 TensorRT 的功能。您可以决定在插件中包含哪些功能,以及将哪些功能留给 TensorRT 处理。

这同样适用于带有 Q/DQ 层的 TensorRT 网络。当插件消耗量化输入(INT8/FP8)并生成量化输出时,输入 DQ 和输出 Q 节点必须包含在插件中并从网络中删除。

考虑一个简单的情况,即一个顺序图,该顺序图由夹在两个卷积层之间的单个 INT8 插件(恰当地命名为 MyInt8Plugin)组成(忽略权重量化)

\(\text{Input}\gt \text{Q}\to \text{DQ}\gt \text{Conv}\gt \text{Q}\to \text{DQ_i}\gt \text{MyInt8Plugin}\gt \text{Q_o}\to \text{DQ}\gt \text{Conv}\gt \text{Output}\)

\(\gt\) 箭头表示具有 FP32 精度的激活张量,而 \(\to\) 箭头表示 INT8 精度。

当 TensorRT 优化此图时,它会将层融合到以下图中(方括号表示 TensorRT 融合)

\(\text{Input}\gt \text{Q}\to \left[\text{DQ}\to \text{Conv} \to \text{Q} \right]\to \text{DQ_i}\gt \text{MyInt8Plugin}\gt \text{Q_o}\to \left[ \text{DQ}\to \text{Conv} \right]\gt \text{Output}\)

在上图中,插件消耗并生成 FP32 输入和输出。由于插件 MyInt8Plugin 使用 INT8 精度,因此后续过程涉及手动集成 DQ_iQ_oMyInt8Plugin,然后为此特定插件层调用 setOutputType(kINT8) 方法;TensorRT 将看到如下网络

\(\text{Input}\gt \text{Q}\to \text{DQ}\gt \text{Conv}\gt \text{Q}\to \text{MyInt8Plugin}\gt \text{DQ}\gt \text{Conv}\gt \text{Output}\)

它将融合到

\(\text{Input}\gt \text{Q}\to \left[ \text{DQ}\to \text{Conv}\to \text{Q} \right]\gt \text{MyInt8Plugin}\gt \left[ \text{DQ}\to \text{Conv} \right]\gt \text{Output}\)

当“手动融合” DQ_i 时,您获取输入量化比例并将其提供给您的插件,以便它知道如何反量化(如果需要)输入。使用 Q_o 中的比例来量化插件的输出也是如此。

使用 TensorFlow 的 QAT 网络#

我们提供了一个开源 TensorFlow 量化工具包,用于在 TensorFlow 2 Keras 模型中按照 NVIDIA 的 QAT 配方执行 QAT。这可以实现 NVIDIA GPU 和硬件加速器上 TensorRT 的最佳模型加速。NVIDIA TensorFlow 量化工具包用户指南提供了更多详细信息。

TensorFlow 1 不支持按通道量化 (PCQ),建议对权重使用 PCQ 以保留模型的准确性。

使用 PyTorch 的 QAT 网络#

PyTorch 1.8.0 及更高版本支持 ONNX QuantizeLinearDequantizeLinear 支持按通道比例。

您可以使用 TensorRT 模型优化器来校准 INT8,为 TensorRT 支持的各种精度执行 QAT 和 PTQ,并导出到 ONNX。

使用 TransformerEngine 的 QAT 网络#

我们提供 TransformerEngine,这是一个开源库,用于加速 Transformer 模型的训练、推理和导出。它包括用于构建 Transformer 层和框架无关的 C++ 库的 API,包括 FP8 支持所需的结构和内核。TransformerEngine 提供的模块在内部维护 FP8 训练所需的缩放因子和其他值。您可以使用 TransformerEngine 训练混合精度模型、导出 ONNX 模型,并使用 TensorRT 在此 ONNX 模型上运行推理。

动态量化#

动态量化是块量化的一种形式,其中比例在推理期间根据输入数据计算。它产生两个输出:量化数据和每块比例。

动态量化有两个主要优点

  1. 准确性:通过动态量化,选择比例以仅将单个块的动态范围映射到量化类型。由于单个块的动态范围通常远小于整个张量的动态范围,因此量化误差会减小。这对于子字节量化类型最为重要,因为这些数据类型中可表示的值范围很小。

  2. 减少 PTQ 开销:由于比例是在推理期间自动计算的,因此用户无需根据样本数据校准比例。

对于每个块,比例通过以下方式计算

\(\text{scale}=max_{i\in \left\{ 0...blockSize-1 \right\}}\left( \frac{abs\left( x_{i} \right)}{qTypeMax} \right)\)

其中

  • \(\text{qTypeMax}\) 是量化类型中的最大值(例如,FP4E2M1 为 6)。

TensorRT 支持一种称为动态双重量化的动态量化形式,其中计算出的比例也被量化。将单个块的比例计算和比例量化放在一起

\(scale_{quantized}=quantize\left( max_{i\in \left\{ 0...blockSize-1 \right\}}\left( \frac{abs\left( x_{i} \right)}{qTypeMax} \right),scale=globalSf \right)\)

其中

  • \(\text{globalSf}\) 是离线校准的每张量量化比例(标量)。

  • \(\text{qTypeMax}\) 是用于数据的量化类型中可表示的最大值。

比例计算针对每个块重复,总共计算 \(\frac{inputVolume}{blockSize}\) 个块比例。

TensorRT 目前仅支持 FP4 数据和 FP8 比例的动态双重量化。使用 \(qTypeMax=6\)[-448,448] 的 FP8 范围,量化比例可以写为

\(scale_{fp8}=castToFp8\left( \frac{max_{i\in \left( 0...blockSize-1 \right)}\left( abs\left(x_{i} \right) \right)}{6\ast globalSf} \right)\)

量化数据是使用块量化和计算出的比例来计算的。

要反量化使用动态双重量化量化的数据,必须发生两个连续的反量化操作:第一个是使用每张量量化反量化比例,第二个是反量化数据。

\(data_{DQ}=dequantize\left( data_{Q},dequantize\left( scale_{Q},scale=globalSf \right) \right)\)

An example showing the fusion of Dynamic Double Quantization with GEMM; where: ``hp`` is high precision and ``sf`` is scale factors

量化类型舍入模式#

后端

计算内核量化

权重量化 (FP32 to INT8/FP8/INT4/FP4)

量化网络 (QAT)

动态范围 API / 校准

GPU

round-to-nearest-with-ties-to-even (INT8, FP8, INT4, FP4)

round-to-nearest-with-ties-to-even

round-to-nearest-with-ties-to-positive-infinity (仅限 INT8)

DLA

round-to-nearest-with-ties-to-even

不适用

round-to-nearest-with-ties-to-even (仅限 INT8)