StableHLO 量化

StableHLO 中的量化类型

量化是一种通过将浮点数(例如原始模型中使用的浮点数)转换为低精度整数来优化机器学习模型的技术。这样可以减少内存用量并加快计算速度,从而使模型更高效地部署在资源有限的设备上。

StableHLO 量化遵循 LiteRT 量化规范,使用统一的量化方案,支持按张量和按轴量化。它从 MLIR 的 Quant 方言继承其类型表达式,从而提供了一种表示量化数据类型的标准化方式。

均匀量化使用统一的步长将浮点值映射到整数,从而生成均匀间隔的量化值。这是通过使用两个关键量化参数的仿射关系来实现的。

均匀量化通过将浮点数映射到均匀间隔的整数来简化浮点数的表示。此映射通过使用两个关键形参(即缩放比例零点)的仿射转换来实现。缩放比例决定了连续量化值之间的步长。缩放比例越小,量化值越接近。零点定义了表示原始浮点空间中零的整数值。

在均匀量化中,原始浮点值 (real_value) 与量化后的整数值 (quantized_value) 之间的关系为:

real_value = scale * (quantized_value - zero_point)

按张量量化

在按张量量化中,单个缩放比例和零点用于张量中的所有值。在 StableHLO 中,每个张量的量化类型表示为:

quant.uniform scale:zero_point>

示例!quant.uniform<i8:f32, 0.01:50>

这表示一个 8 位整数 (i8),用于使用 0.01 的缩放比例和 50 的零点存储 32 位浮点数 (f32)。

按轴量化

与按张量量化相比,按轴量化可提供更精细的方法。与为整个张量使用单一缩放比例和零点不同,按轴量化会为张量的特定维度 quantized_dimension 上的切片分配单独的缩放比例和零点。当不同维度之间的值差异很大时,此功能特别有用,可更好地保留信息并提高准确性。

假设有一个维度大小为 [4, 3, 2] 的张量 t。我们选择沿第二个维度 (quantized_dimension = 1) 量化此张量。这意味着我们将有三个切片(因为第二个维度的尺寸为 3),每个切片都有自己的缩放比例和零点:

t[:, 0, :]: This slice gets scale[0] and zero_point[0].
t[:, 1, :]: This slice gets scale[1] and zero_point[1].
t[:, 2, :]: This slice gets scale[2] and zero_point[2].

在 StableHLO 中,按轴量化类型表示为:

quant.uniform {scale0:zero_point0, scale1:zero_point1, ...}>

其中,scale:zero_point 的长度与包含张量的 quantized_dimension 上的切片数量一致。

示例tensor<4x3x2x!quant.uniform<i8:f32:1, {0.2:20, 0.1:10, 0.3:30}>>

StableHLO 中的量化传递

StableHLO 提供了多个编译器传递,可实现与量化相关的不同转换和优化,让您可以灵活地处理量化模型。这些卡包括:

stablehlo-legalize-qdq-to-quantized-op

此 pass 会将量化模型中的常见模式(即先进行反量化操作,然后进行浮点运算,最后进行量化操作)融合为单个量化操作。详情

stablehlo-legalize-quantized-op-to-qdq

此遍处理与上一遍处理相反。它将量化操作分解为等效的去量化、浮点运算和量化操作序列。 详情

stablehlo-legalize-quant-to-math

此遍次会将量化类型上的 StableHLO 操作转换为整数类型上的等效操作。它本质上是使用标准数学运算来实现量化算术。对于不支持原生量化但仍可使用量化算术来表达量化模型语义的系统,此分解非常有用。详情

stablehlo-quant-legalize-to-tosa-rescale

StableHLO 能够将量化操作合法化为 TOSA 方言中的相应表示形式。这种合法化有助于实现 StableHLO 和 TOSA 之间的兼容性和互操作性。此传递会策略性地将 StableHLO 量化操作转换为 StableHLO 和 TOSA 操作的组合,其中 TOSA 方言主要用于 rescale 操作。tosa.rescale 操作在调整量化值的缩放比例和零点方面发挥着至关重要的作用,从而能够在 TOSA 框架内准确表示量化数据。详情

tosa-rescale-legalize-to-stablehlo

此 pass 会将 TOSA 重新缩放操作重写为 StableHLO 基元数学运算。此传递的主要使用场景之一是允许 StableHLO 解释器评估包含 TOSA 重新缩放操作的程序。详情

评估量化程序

StableHLO 参考解释器可以高效执行包含量化运算的程序。为此,它首先使用仅包含整数运算的等效表示形式来降低程序。此降级过程涉及一系列编译器传递,这些传递会在解释之前转换程序。

从本质上讲,解释器利用 stablehlo-legalize-quant-to-math 传递将量化操作转换为相应的整数算术实现。此传递引入了 CHLO 广播操作,用于处理缩放乘法/除法和零点加法。为了确保与 StableHLO 解释器的兼容性,这些 CHLO 操作随后会被合法化为 StableHLO 操作。这会引入形状相关操作,这些操作随后会通过一系列规范化传递进行规范化和优化。

此降级过程涉及的完整遍历序列如下:

stablehlo-legalize-quant-to-math
chlo-legalize-to-stablehlo
canonicalize
shape-legalize-to-stablehlo
stablehlo-canonicalize-dynamism

量化测试用例

StableHLO 提供了一套全面的量化测试用例,用于验证量化操作的正确性和行为。这些测试用例充当单元测试,涵盖量化场景中的各种 StableHLO 操作。

量化测试用例的典型示例如下所示

func.func @main() -> tensor<11xf32> {
    %operand_0 = stablehlo.constant dense<...> : tensor<11xf32>
    %operand_1 = stablehlo.constant dense<...> : tensor<11xf32>
    %golden = stablehlo.constant dense<...> : tensor<11xf32>

    %0 = stablehlo.uniform_quantize %operand_0 : (tensor<11xf32>) -> tensor<11x!quant.uniform<i8:f32, 0.3>>
    %1 = stablehlo.uniform_quantize %operand_1 : (tensor<11xf32>) -> tensor<11x!quant.uniform<i8:f32, 0.3>>

    %2 = stablehlo.add %1, %0 : tensor<11x!quant.uniform<i8:f32, 0.3>>

    %result = stablehlo.uniform_dequantize %2 : (tensor<11x!quant.uniform<i8:f32, 0.3>>) -> tensor<11xf32>

    %4 = stablehlo.custom_call @check.eq(%golden, %result) : (tensor<11xf32>, tensor<11xf32>) -> tensor<i1>

    return %3 : tensor<11xf32>
  }

且包含:

这些测试用例对于以下情况非常有用:

  • 验证 StableHLO 量化:确保 StableHLO 操作的量化行为与预期结果一致。
  • 交叉验证:将 StableHLO 量化的行为与其他实现或框架进行比较。
  • 调试和开发:帮助开发和调试新的量化功能或优化。