StableHLO-Quantisierung

Quantisierungstypen in StableHLO

Die Quantisierung ist eine Technik zur Optimierung von Machine-Learning-Modellen, bei der Gleitkommazahlen (wie sie in Originalmodellen verwendet werden) in Ganzzahlen mit geringerer Genauigkeit umgewandelt werden. Dadurch wird der Speicherverbrauch reduziert und die Berechnungen werden beschleunigt, sodass Modelle effizienter auf Geräten mit begrenzten Ressourcen bereitgestellt werden können.

Die StableHLO-Quantisierung folgt der LiteRT-Quantisierungsspezifikation und verwendet ein einheitliches Quantisierungsschema, das sowohl die Quantisierung pro Tensor als auch pro Achse unterstützt. Der Typausdruck wird vom Quant-Dialekt von MLIR abgeleitet und bietet eine standardisierte Möglichkeit zur Darstellung quantisierter Datentypen.

Bei der gleichmäßigen Quantisierung werden Gleitkommawerte mithilfe einer gleichmäßigen Schrittweite in Ganzzahlen umgewandelt. Das Ergebnis sind gleichmäßig verteilte quantisierte Werte. Dies wird durch eine affine Beziehung mit zwei wichtigen Quantisierungsparametern erreicht.

Die einheitliche Quantisierung vereinfacht die Darstellung von Gleitkommazahlen, indem sie auf gleichmäßig verteilte Ganzzahlen abgebildet werden. Diese Zuordnung wird durch eine affine Transformation erreicht, die zwei wichtige Parameter verwendet: Skalierung und Nullpunkt. Die Skalierung bestimmt die Schrittweite zwischen aufeinanderfolgenden quantisierten Werten. Bei einem kleineren Maßstab liegen die quantisierten Werte näher beieinander. Der Nullpunkt definiert den ganzzahligen Wert, der im ursprünglichen Gleitkomma-Bereich für null steht.

Die Beziehung zwischen dem ursprünglichen Gleitkommawert (real_value) und dem quantisierten Ganzzahlwert (quantized_value) bei der einheitlichen Quantisierung ist:

real_value = scale * (quantized_value - zero_point)

Quantisierung pro Tensor

Bei der Quantisierung pro Tensor werden ein einzelner Skalierungsfaktor und ein einzelner Nullpunkt für alle Werte im Tensor verwendet. Ein quantisierter Typ pro Tensor wird in StableHLO so ausgedrückt:

quant.uniform scale:zero_point>

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

Dies entspricht einer 8-Bit-Ganzzahl (i8), die zum Speichern einer 32-Bit-Gleitkommazahl (f32) mit einem Skalierungsfaktor von 0.01 und einem Nullpunkt von 50 verwendet wird.

Achsenspezifische Quantisierung

Die achsenweise Quantisierung bietet einen detaillierteren Ansatz als die tensorweise Quantisierung. Bei der Quantisierung pro Achse werden nicht eine einzelne Skalierung und ein einzelner Nullpunkt für den gesamten Tensor verwendet, sondern den Slices entlang einer bestimmten Dimension quantized_dimension des Tensors werden separate Skalierungen und Nullpunkte zugewiesen. Das ist besonders nützlich, wenn die Werte über verschiedene Dimensionen hinweg erheblich variieren, da so Informationen und Genauigkeit besser erhalten bleiben.

Angenommen, Sie haben einen Tensor t mit den Dimensionen [4, 3, 2]. Wir quantisieren diesen Tensor entlang der zweiten Dimension (quantized_dimension = 1). Das bedeutet, dass wir drei Slices haben (da die zweite Dimension eine Größe von 3 hat), die jeweils einen eigenen Skalierungsfaktor und Nullpunkt haben:

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

In StableHLO wird der quantisierte Typ pro Achse so ausgedrückt:

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

Dabei entspricht die Länge von scale:zero_point der Anzahl der Slices entlang der quantized_dimension des enthaltenden Tensors.

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

Quantisierungsdurchläufe in StableHLO

StableHLO bietet mehrere Compiler-Durchläufe, die verschiedene Transformationen und Optimierungen im Zusammenhang mit der Quantisierung ermöglichen. So haben Sie Flexibilität bei der Verarbeitung quantisierter Modelle. Folgende Pässe sind verfügbar:

stablehlo-legalize-qdq-to-quantized-op

Bei diesem Durchlauf wird ein häufiges Muster in quantisierten Modellen zusammengeführt: ein Dequantisierungsvorgang, gefolgt von einem Gleitkommavorgang und schließlich einem Quantisierungsvorgang, in einen einzelnen quantisierten Vorgang. Details

stablehlo-legalize-quantized-op-to-qdq

Dieser Durchlauf macht das Gegenteil des vorherigen Durchlaufs. Sie zerlegt einen quantisierten Vorgang in die entsprechende Sequenz von Vorgängen zum Dequantisieren, Gleitkomma-Vorgang und Quantisieren. Details

stablehlo-legalize-quant-to-math

Bei diesem Durchlauf werden StableHLO-Vorgänge für quantisierte Typen in entsprechende Vorgänge für Ganzzahltypen konvertiert. Dabei wird die Arithmetik der Quantisierung mithilfe von mathematischen Standardoperationen implementiert. Diese Zerlegung ist nützlich für Systeme, die die Quantisierung nicht nativ unterstützen, aber dennoch die Quantisierungsarithmetik verwenden können, um die Semantik quantisierter Modelle auszudrücken. Details

stablehlo-quant-legalize-to-tosa-rescale

StableHLO bietet die Möglichkeit, quantisierte Vorgänge in ihre entsprechenden Darstellungen im TOSA-Dialekt zu legalisieren. Diese Legalisierung erleichtert die Kompatibilität und Interoperabilität zwischen StableHLO und TOSA. In diesem Durchgang werden StableHLO-quantisierte Vorgänge strategisch in eine Kombination aus StableHLO- und TOSA-Vorgängen umgewandelt, wobei der TOSA-Dialekt hauptsächlich für den rescale-Vorgang verwendet wird. Der tosa.rescale-Vorgang spielt eine entscheidende Rolle bei der Anpassung des Maßstabs und des Nullpunkts quantisierter Werte, sodass quantisierte Daten im TOSA-Framework genau dargestellt werden können. Details

tosa-rescale-legalize-to-stablehlo

In diesem Durchgang werden TOSA-Rescale-Vorgänge in StableHLO-Primitive-Math-Vorgänge umgeschrieben. Einer der Hauptanwendungsfälle für diesen Pass ist es, dem StableHLO-Interpreter zu ermöglichen, Programme mit TOSA-Rescale-Operationen auszuwerten. Details

Quantisierte Programme bewerten

Der StableHLO-Referenzinterpreter kann Programme mit quantisierten Vorgängen effizient ausführen. Dazu wird das Programm zuerst in eine äquivalente Darstellung mit nur Ganzzahloperationen umgewandelt. Dieser Prozess umfasst eine Reihe von Compilerdurchläufen, die das Programm vor der Interpretation transformieren.

Der Interpreter nutzt den stablehlo-legalize-quant-to-math-Pass, um quantisierte Vorgänge in die entsprechenden Implementierungen für die Ganzzahlarithmetik zu konvertieren. In diesem Durchlauf werden CHLO-Broadcast-Vorgänge zum Verarbeiten von Skalarmultiplikation/-division und Zero-Point-Addition eingeführt. Um die Kompatibilität mit dem StableHLO-Interpreter zu gewährleisten, werden diese CHLO-Vorgänge dann in StableHLO-Vorgänge umgewandelt. Dadurch werden formbezogene Vorgänge eingeführt, die anschließend mithilfe einer Reihe von Kanonisierungsdurchläufen kanonisiert und optimiert werden.

Die vollständige Abfolge der Durchgänge, die bei diesem Prozess zum Senken des Rangs erforderlich sind, ist wie folgt:

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

Quantisierte Testläufe

StableHLO bietet eine umfassende Suite quantisierter Testläufe, um die Richtigkeit und das Verhalten quantisierter Vorgänge zu validieren. Diese Testläufe dienen als Einheitentests und decken verschiedene StableHLO-Operationen in quantisierten Szenarien ab.

Ein typisches Beispiel für einen quantisierten Testlauf sieht so aus:

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

und umfasst:

  • Eingabedaten:Repräsentative Eingabewerte für den Vorgang.
  • Goldene Ausgabe:Die erwartete Ausgabe des Vorgangs, wenn er auf die Eingabedaten angewendet wird, in Übereinstimmung mit dem StableHLO-Referenzinterpreter und dem HLO-Evaluator.

Diese Testläufe sind nützlich für:

  • StableHLO-Quantisierung validieren:Sicherstellen, dass das Quantisierungsverhalten von StableHLO-Operationen mit den erwarteten Ergebnissen übereinstimmt.
  • Kreuzvalidierung:Vergleichen des Verhaltens der StableHLO-Quantisierung mit anderen Implementierungen oder Frameworks.
  • Debugging und Entwicklung:Unterstützung bei der Entwicklung und dem Debugging neuer Quantisierungsfunktionen oder ‑optimierungen.