Cuantización de StableHLO

Tipos de cuantización en StableHLO

La cuantización es una técnica para optimizar los modelos de aprendizaje automático que convierte los números de punto flotante (como los que se usan en los modelos originales) en números enteros de menor precisión. Esto reduce el uso de memoria y acelera los cálculos, lo que hace que los modelos sean más eficientes para la implementación en dispositivos con recursos limitados.

La cuantización de StableHLO sigue la especificación de cuantización de LiteRT y usa un esquema de cuantización uniforme con compatibilidad para la cuantización por tensor y por eje. Hereda su expresión de tipo del dialecto Quant de MLIR, lo que proporciona una forma estandarizada de representar tipos de datos cuantificados.

La cuantificación uniforme asigna valores de punto flotante a números enteros con un tamaño de paso uniforme, lo que genera valores cuantificados espaciados de manera uniforme. Esto se logra a través de una relación afín que usa dos parámetros de cuantificación clave.

La cuantización uniforme simplifica la representación de números de punto flotante, ya que los asigna a números enteros que están espaciados de manera uniforme. Esta asignación se logra a través de una transformación afín que usa dos parámetros clave: escala y punto cero. La escala determina el tamaño del paso entre valores cuantificados consecutivos. Una escala más pequeña significa que los valores cuantificados están más cerca entre sí. El punto cero define el valor entero que representa el cero en el espacio de punto flotante original.

La relación entre el valor de punto flotante original (real_value) y el valor de número entero cuantificado (quantized_value) en la cuantificación uniforme es la siguiente:

real_value = scale * (quantized_value - zero_point)

Cuantización por tensor

En la cuantización por tensor, se usan un solo punto cero y una sola escala para todos los valores dentro del tensor. Un tipo cuantificado por tensor se expresa en StableHLO de la siguiente manera:

quant.uniform scale:zero_point>

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

Representa un número entero de 8 bits (i8) que se usa para almacenar un número de punto flotante de 32 bits (f32) con una escala de 0.01 y un punto cero de 50.

Cuantización por eje

La cuantización por eje ofrece un enfoque más detallado en comparación con la cuantización por tensor. En lugar de usar una sola escala y un solo punto cero para todo el tensor, la cuantificación por eje asigna escalas y puntos cero separados a las porciones a lo largo de una dimensión específica quantized_dimension del tensor. Esto es especialmente útil cuando los valores varían significativamente en las diferentes dimensiones, lo que permite conservar mejor la información y la precisión.

Considera un tensor t con tamaños de dimensiones [4, 3, 2]. Elegimos cuantificar este tensor a lo largo de la segunda dimensión (quantized_dimension = 1). Esto significa que tendremos tres segmentos (ya que la segunda dimensión tiene un tamaño de 3), cada uno con su propia escala y punto cero:

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].

En StableHLO, el tipo cuantificado por eje se expresa de la siguiente manera:

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

donde la longitud de scale:zero_point coincide con la cantidad de segmentos a lo largo de quantized_dimension del tensor que lo contiene.

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

Pases de cuantización en StableHLO

StableHLO proporciona varios pases del compilador que permiten diferentes transformaciones y optimizaciones relacionadas con la cuantificación, lo que te brinda flexibilidad en la forma en que controlas los modelos cuantificados. Estos pases son los siguientes:

stablehlo-legalize-qdq-to-quantized-op

Este paso fusiona un patrón común en los modelos cuantizados, una operación de decuantización seguida de una operación de punto flotante y, finalmente, una operación de cuantización, en una sola operación cuantizada. detalles

stablehlo-legalize-quantized-op-to-qdq

Este pase hace lo contrario del pase anterior. Descompone una operación cuantificada en su secuencia equivalente de operaciones de descuantización, punto flotante y cuantificación. detalles

stablehlo-legalize-quant-to-math

Este paso convierte las operaciones de StableHLO en tipos cuantificados en operaciones equivalentes en tipos de números enteros. Básicamente, implementa la aritmética de cuantización con operaciones matemáticas estándar. Esta descomposición es útil para los sistemas que no admiten la cuantización de forma nativa, pero que pueden usar la aritmética de cuantización para expresar la semántica de los modelos cuantizados. detalles

stablehlo-quant-legalize-to-tosa-rescale

StableHLO ofrece la capacidad de legalizar operaciones cuantificadas en sus representaciones correspondientes en el dialecto de TOSA. Esta legalización facilita la compatibilidad y la interoperabilidad entre StableHLO y TOSA. Este paso convierte estratégicamente las operaciones cuantificadas de StableHLO en una combinación de operaciones de StableHLO y TOSA, y el dialecto de TOSA se emplea principalmente para la operación rescale. La operación tosa.rescale desempeña un papel fundamental en el ajuste de la escala y el punto cero de los valores cuantificados, lo que permite una representación precisa de los datos cuantificados dentro del marco de TOSA. detalles

tosa-rescale-legalize-to-stablehlo

Este paso reescribe las operaciones de cambio de escala de TOSA en operaciones matemáticas primitivas de StableHLO. Uno de los principales casos de uso de este paso es permitir que el intérprete de StableHLO evalúe programas que contienen operaciones de cambio de escala de TOSA. detalles

Cómo evaluar programas cuantizados

El intérprete de referencia de StableHLO puede ejecutar de manera eficiente programas que contienen operaciones cuantificadas. Para lograr esto, primero reduce el programa a una representación equivalente que solo usa operaciones de números enteros. Este proceso de reducción implica una serie de pases del compilador que transforman el programa antes de la interpretación.

Básicamente, el intérprete aprovecha el pase stablehlo-legalize-quant-to-math para convertir las operaciones cuantificadas en sus implementaciones aritméticas de números enteros correspondientes. Este paso introduce operaciones de transmisión de CHLO para controlar la multiplicación o división de la escala y la suma del punto cero. Para garantizar la compatibilidad con el intérprete de StableHLO, estas operaciones de CHLO se legalizan en operaciones de StableHLO. Esto introduce operaciones relacionadas con la forma que luego se canonizan y optimizan con una serie de pases de canonización.

La secuencia completa de pases involucrados en este proceso de reducción es la siguiente:

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

Casos de prueba cuantificados

StableHLO proporciona un conjunto integral de casos de prueba cuantificados para validar la corrección y el comportamiento de las operaciones cuantificadas. Estos casos de prueba sirven como pruebas de unidad y abarcan varias operaciones de StableHLO en situaciones cuantificadas.

Un ejemplo típico de un caso de prueba cuantificado se ve de la siguiente manera:

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>
  }

y se incluye lo siguiente:

• Datos de entrada: Son valores de entrada representativos para la operación.
• Salida ideal: Es el resultado esperado de la operación cuando se aplica a los datos de entrada, de conformidad con el intérprete de referencia de StableHLO y el evaluador de HLO.

Estos casos de prueba son valiosos para lo siguiente:

• Validación de la cuantificación de StableHLO: Garantizar que el comportamiento de cuantificación de las operaciones de StableHLO se alinee con los resultados esperados
• Validación cruzada: Comparar el comportamiento de la cuantificación de StableHLO con otras implementaciones o frameworks
• Depuración y desarrollo: Ayuda en el desarrollo y la depuración de nuevas funciones o optimizaciones de cuantificación.