Quantização do StableHLO

Tipos de quantização em StableHLO

A quantização é uma técnica para otimizar modelos de machine learning convertendo números de ponto flutuante (como os usados em modelos originais) em números inteiros de menor precisão. Isso reduz o uso de memória e acelera os cálculos, tornando os modelos mais eficientes para implantação em dispositivos com recursos limitados.

A quantização do StableHLO segue a especificação de quantização do LiteRT (link em inglês), usando um esquema de quantização uniforme com suporte para quantização por tensor e por eixo. Ele herda a expressão de tipo do dialeto Quant do MLIR, oferecendo uma maneira padronizada de representar tipos de dados quantizados.

A quantização uniforme mapeia valores de ponto flutuante para números inteiros usando um tamanho de etapa uniforme, resultando em valores quantizados uniformemente espaçados. Isso é feito por uma relação afim usando dois parâmetros de quantização principais.

A quantização uniforme simplifica a representação de números de ponto flutuante mapeando-os para números inteiros com espaçamento uniforme. Esse mapeamento é feito por uma transformação afim que usa dois parâmetros principais: escala e ponto zero. A escala determina o tamanho da etapa entre valores quantizados consecutivos. Uma escala menor significa que os valores quantizados estão mais próximos. O ponto zero define o valor inteiro que representa zero no espaço de ponto flutuante original.

A relação entre o valor de ponto flutuante original (real_value) e o valor inteiro quantizado (quantized_value) na quantização uniforme é:

real_value = scale * (quantized_value - zero_point)

Quantização por tensor

Na quantização por tensor, uma única escala e um único ponto zero são usados para todos os valores dentro do tensor. Um tipo quantizado por tensor é expresso em StableHLO como:

quant.uniform scale:zero_point>

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

Isso representa um número inteiro de 8 bits (i8) usado para armazenar um número de ponto flutuante de 32 bits (f32) usando uma escala de 0.01 e um ponto zero de 50.

Quantização por eixo

A quantização por eixo oferece uma abordagem mais refinada em comparação com a quantização por tensor. Em vez de usar uma única escala e ponto zero para todo o tensor, a quantização por eixo atribui escalas e pontos zero separados a fatias ao longo de uma dimensão específica quantized_dimension do tensor. Isso é particularmente útil quando os valores variam muito em diferentes dimensões, permitindo uma melhor preservação das informações e da precisão.

Considere um tensor t com tamanhos de dimensões [4, 3, 2]. Optamos por quantizar esse tensor ao longo da segunda dimensão (quantized_dimension = 1). Isso significa que teremos três intervalos (já que a segunda dimensão tem um tamanho de 3), cada um com sua própria escala e ponto zero:

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

Em StableHLO, o tipo quantizado por eixo é expresso como:

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

em que o comprimento de scale:zero_point corresponde ao número de fatias ao longo da quantized_dimension do tensor que contém.

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

Passagens de quantização em StableHLO

O StableHLO oferece várias transmissões de compilador que permitem diferentes transformações e otimizações relacionadas à quantização, oferecendo flexibilidade na forma como você lida com modelos quantizados. São eles:

`stablehlo-legalize-qdq-to-quantized-op`

Essa transmissão combina um padrão comum em modelos quantizados, uma operação de desquantização seguida por uma operação de ponto flutuante e, por fim, uma operação de quantização em uma única operação quantizada. detalhes

stablehlo-legalize-quantized-op-to-qdq

Essa transmissão faz o oposto da anterior. Ela decompõe uma operação quantizada na sequência equivalente de operações de desquantização, ponto flutuante e quantização. detalhes

stablehlo-legalize-quant-to-math

Essa transmissão converte operações do StableHLO em tipos quantizados em operações equivalentes em tipos de números inteiros. Ele implementa essencialmente a aritmética de quantização usando operações matemáticas padrão. Essa decomposição é útil para sistemas que não oferecem suporte à quantização de forma nativa, mas ainda podem usar a aritmética de quantização para expressar a semântica dos modelos quantizados. detalhes

stablehlo-quant-legalize-to-tosa-rescale

O StableHLO oferece a capacidade de legalizar operações quantizadas para as representações correspondentes no dialeto TOSA. Essa legalização facilita a compatibilidade e a interoperabilidade entre o StableHLO e o TOSA. Essa transmissão converte estrategicamente operações quantizadas do StableHLO em uma combinação de operações do StableHLO e da TOSA, com o dialeto da TOSA usado principalmente para a operação rescale. A operação tosa.rescale desempenha um papel crucial no ajuste da escala e do ponto zero dos valores quantizados, permitindo uma representação precisa dos dados quantizados na estrutura TOSA. detalhes

tosa-rescale-legalize-to-stablehlo

Essa transmissão reescreve operações de redimensionamento do TOSA para operações matemáticas primitivas do StableHLO. Um dos principais casos de uso dessa transmissão é permitir que o intérprete do StableHLO avalie programas que contêm operações de reescala da TOSA. Detalhes

Como avaliar programas quantizados

O interpretador de referência StableHLO pode executar com eficiência programas que contêm operações quantizadas. Para isso, primeiro ele reduz o programa a uma representação equivalente usando apenas operações de números inteiros. Esse processo envolve uma série de transmissões do compilador que transformam o programa antes da interpretação.

Basicamente, o intérprete usa a transmissão stablehlo-legalize-quant-to-math para converter operações quantizadas nas implementações de aritmética de números inteiros correspondentes. Essa transmissão apresenta operações de transmissão CHLO para lidar com multiplicação/divisão de escala e adição de ponto zero. Para garantir a compatibilidade com o interpretador StableHLO, essas operações CHLO são legalizadas para operações StableHLO. Isso introduz operações relacionadas a formas que são posteriormente canonizadas e otimizadas usando uma série de transmissões de canonização.

A sequência completa de transmissões envolvidas nesse processo de redução é a seguinte:

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

Casos de teste quantizados

O StableHLO oferece um conjunto abrangente de casos de teste quantizados para validar a correção e o comportamento das operações quantizadas. Esses casos de teste servem como testes de unidade, abrangendo várias operações do StableHLO em cenários quantizados.

Um exemplo típico de caso de teste quantizado é

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

e inclui:

Dados de entrada:valores de entrada representativos para a operação.
Saída de referência:a saída esperada da operação quando aplicada aos dados de entrada, em conformidade com o interpretador de referência StableHLO e o avaliador de HLO.

Esses casos de teste são úteis para:

Validação da quantização do StableHLO:garante que o comportamento de quantização das operações do StableHLO esteja alinhado aos resultados esperados.
Validação cruzada:comparar o comportamento da quantização do StableHLO com outras implementações ou frameworks.
Depuração e desenvolvimento:ajuda no desenvolvimento e na depuração de novos recursos ou otimizações de quantização.