StableHLO の量子化

StableHLO の量子化タイプ

量子化は、浮動小数点数（元のモデルで使用されているものなど）を低精度の整数に変換して、ML モデルを最適化する手法です。これにより、メモリ使用量が減り、計算が高速化されるため、リソースが限られたデバイスへのデプロイでモデルの効率が向上します。

StableHLO 量子化は、LiteRT 量子化仕様に準拠しており、テンソル単位と軸単位の両方の量子化をサポートする均一な量子化スキームを使用します。MLIR の Quant 言語から型式を継承し、量子化されたデータ型を表す標準化された方法を提供します。

均一量子化では、均一なステップサイズを使用して浮動小数点値を整数にマッピングし、等間隔の量子化値が得られます。これは、2 つの主要な量子化パラメータを使用したアフィン関係によって実現されます。

均一量子化では、浮動小数点数を均等に配置された整数にマッピングすることで、浮動小数点数の表現を簡素化します。このマッピングは、スケールとゼロ点という 2 つのキーパラメータを使用するアフィン変換によって実現されます。スケールは、連続する量子化値間のステップサイズを決定します。スケールが小さいほど、量子化された値が近くなります。ゼロポイントは、元の浮動小数点空間でゼロを表す整数値を定義します。

均一量子化における元の浮動小数点値（real_value）と量子化された整数値（quantized_value）の関係は次のとおりです。

real_value = scale * (quantized_value - zero_point)

テンソルごとの量子化

テンソル単位の量子化では、テンソル内のすべての値に単一のスケールとゼロ点が使用されます。テンソルごとの量子化型は、StableHLO で次のように表されます。

quant.uniform scale:zero_point>

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

これは、0.01 のスケールと 50 のゼロポイントを使用して 32 ビットの浮動小数点数（f32）を保存するために使用される 8 ビットの整数（i8）を表します。

軸ごとの量子化

軸ごとの量子化は、テンソルごとの量子化よりもきめ細かいアプローチです。テンソル全体に単一のスケールとゼロ点を使用する代わりに、軸ごとの量子化では、テンソルの特定のディメンション quantized_dimension に沿ったスライスに個別のスケールとゼロ点が割り当てられます。これは、値がディメンション間で大きく異なる場合に特に便利で、情報と精度の維持に役立ちます。

ディメンションサイズが [4, 3, 2] のテンソル t を考えます。このテンソルを 2 番目のディメンション（quantized_dimension = 1）に沿って量子化することにします。つまり、3 つのスライス（2 番目のディメンションのサイズが 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`

このパスは、量子化モデルの一般的なパターン（逆量子化オペレーション、浮動小数点オペレーション、量子化オペレーション）を 1 つの量子化オペレーションに融合します。詳細

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

このパスは、TOSA のリスケールオペレーションを StableHLO プリミティブの算術オペレーションに書き換えます。このパスの主なユースケースの 1 つは、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 リファレンスインタープリタと HLO エバリュエータに準拠して、入力データに適用されたオペレーションの想定される出力。

これらのテストケースは、次の点で役立ちます。

StableHLO 量子化の検証: StableHLO オペレーションの量子化動作が期待される結果と一致することを確認します。
交差検証: StableHLO 量子化の動作を他の実装やフレームワークと比較します。
デバッグと開発: 新しい量子化機能や最適化の開発とデバッグを支援します。