Effizientere Inference durch Model Quantization

In Teil eins dieser zweiteiligen Serie über effizientere Model Inference auf CPUs habe ich verschiedene Methoden verglichen, die Modelle optimieren, indem sie diese in Compute Graphs umwandeln, effizientere Operationen nutzen und mehrere Model Layers zu einer einzigen verschmelzen. Keine dieser Methoden ändert die Parameter des Modells. Daher beeinflussen sie nicht die Performance Metrics des Modells, wie etwa die Accuracy. Wenn man jedoch bereit ist, eine gewisse Verschlechterung der Performance in Kauf zu nehmen, gibt es andere Optimierungen zur Beschleunigung der Inference. In diesem Artikel werde ich eine solche Technik genauer unter die Lupe nehmen, nämlich die Quantization der Modellparameter auf 8 Bit.

Überblick über Optimierungsstrategien

Quantization ist der bekannteste Ansätz, die Modellparameter verändern, um eine schnellere Inference zu ermöglichen, aber ist nur einer von vielen. Ein weiterer beliebter Ansatz ist das Model Pruning. Wie der Name schon sagt, werden beim Model Pruning ganze Parameter entfernt, oft diejenigen mit der geringsten Magnitude. Dieser Ansatz funktioniert sowohl während des Trainings durch iteratives Entfernen von Parametern als auch nach dem Training, indem eine festgelegte Anzahl von Parametern entfernt wird, mit anschließendem Fine-Tuning, um den durch das Pruning verursachten Qualitätsverlust auszugleichen. Sogar komplexere Methoden, die Second-Order Derivatives in ihrem Pruning Criterion verwenden, schneiden beim Post-Training ohne die Notwendigkeit eines Fine-Tunings gut ab, sind jedoch mit erheblichen Kosten für die Iteration verbunden, um den nächsten zu entfernenden Parameter zu finden. Alle diese Pruning-Methoden haben gemeinsam, dass sie zu wesentlich kleineren Modellen führen, die bei nur geringem Performanceverlust schneller laufen.

Ein weiterer verbreiteter Ansatz ist die Model Distillation, bei der ein kleineres Modell auf den Softmax Outputs eines größeren (Ensemble) Teacher Models trainiert wird. Die Prämisse ist, dass die Wahrscheinlichkeitsverteilung der Outputs des Teachers wertvolle Informationen enthält, die dem destillierten Modell helfen, besser zu generalisieren als beim Training rein auf gelabelten Daten mit deren Bias zur korrekten Antwort. Daher wählt man üblicherweise eine hohe (>1) Temperature für den Softmax, um mehr Variabilität im Output des Teachers zu erhalten. Dieser Ansatz erfordert jedoch das Training eines komplett neuen, wenn auch kleineren Modells von Grund auf.

Es gibt noch weitere Methoden, wie die Low-Rank Approximation von Linear Layers und Kombinationen aus zwei oder mehr dieser Methoden. Ich habe mich jedoch entschieden, mich in meinen Experimenten ganz auf die Post-Training Quantization zu konzentrieren, da es sich um eine Ad-hoc-Methode handelt, die auf viele Model Architectures anwendbar ist.

Quantization

Ein Quantization-Primer

Ein einfacher, aber leistungsstarker Quantization-Algorithm ist die *Uniform Affine Quantization*. Er wandelt eine Floating-Point Number x in einen Integer p um:
$$p=round(\frac{x}{S})+Z$$ In der obigen Gleichung ist S der Scale und Z der Zero Point. Der Scale bestimmt im Wesentlichen die Auflösung der Quantization. Wenn der Scale beispielsweise auf 40 eingestellt ist, werden Floating-Point Numbers im Bereich [x - 20, x + 20] in denselben quantisierten Wert p umgewandelt. Daher ist eine sorgfältige Auswahl des Scales entscheidend für die Quantization Performance. Wenn Sie wissen, dass die meisten Ihrer Eingaben im Bereich [min, max] liegen, können Sie den Scale für die 8-Bit-Quantization nach dieser Formel festlegen:
$$S=\frac{(max-min)}{255}$$ Der zweite Parameter, der Zero Point, dient dazu, Null auf eine quantisierte Zahl abzubilden, was nützlich ist, um beispielsweise Sparsity beizubehalten oder Daten für sequentielle Daten wie Text Input mit Nullen aufzufüllen (Padding). Wenn Ihre Zahlen wieder im obigen Bereich liegen, kann Z wie folgt berechnet werden:
$$Z=-round(\frac{min}{S})$$ Für einen Bereich von [-2, 3] wäre S beispielsweise 0,0196 und Z wäre 102. Mit der Formel für Affine Quantization wird -1 zu 5 quantisiert. Das Konvertieren von Zahlen außerhalb dieses Bereichs ist nicht möglich; alle Werte unter -2 müssen auf 0 und alle Werte über 3 auf 255 abgebildet werden.

Herausforderungen der Quantization

Die Wahl einer guten Quantization Method für Floating-Point Numbers ist nur die erste von vielen Herausforderungen. Eine weitere ist das Multiplizieren von zwei 8-Bit-Zahlen, was zu einer 16-Bit-Zahl führt, oder das noch komplexere Szenario der Matrix Multiplication, die eine Akkumulation in 32-Bit-Zahlen erfordert. Die Ergebnisse solcher Operationen müssen durch die Berechnung neuer Scales und Zero Points wieder auf eine neue 8-Bit-Repräsentation herunterskaliert werden. Um zudem die meisten Operationen der Matrix Multiplication zwischen 8-Bit-Zahlen zu halten, müssen zusätzliche Tricks angewendet werden, wie das Umschreiben der Gleichungen. Glücklicherweise kümmern sich Frameworks wie PyTorch um diese Implementierungsdetails.

Static vs Dynamic Quantization

PyTorch bietet zwei Wege zur Quantization eines Modells: Dynamic und Static Quantization. Bei der ersteren werden die Weights quantisiert und als solche gespeichert. Die Normalization Layers und Activations verbleiben jedoch im Floating Point und werden mit Floating-Point Operations berechnet. Mit anderen Worten: Nur die rechenintensiven Aufgaben wie der Linear Layer werden quantisiert. Dabei müssen die Inputs eines quantisierten Linear Layers on-the-fly quantisiert werden, indem ein passender Scale und Zero Point berechnet werden. Die Outputs dieses Linear Layers sind 32-Bit-Zahlen, das Ergebnis vieler akkumulierter 8-Bit-Multiplikationen. Daher müssen sie zurück in Floating-Point Numbers konvertiert werden, um an die nächsten Activation Functions weitergeleitet zu werden. Insgesamt funktioniert dieser Ansatz gut für Modelle mit großen linearen Blöcken, die den Rechen-Engpass (Bottleneck) darstellen, wie Transformers.

Static Quantization kann nicht für Transformers verwendet werden, da Beobachtungen gezeigt haben, dass es gelegentlich Outlier Activations mit hohen Werten gibt, was die Auswahl eines Scales mit ausreichender Auflösung erschwert. Infolgedessen führt die Quantization der Activations zu einer Performanceverschlechterung. Jüngere Ansätze wie SmoothQuant konnten diese Probleme jedoch überwinden und ermöglichen die Quantization von Activation Functions ohne hohen Verlust an Accuracy.

Post Training Quantization vs. Quantization Aware Training

Die im vorherigen Abschnitt vorgestellten Methoden wurden alle nach dem Training angewendet und funktionieren gut für Quantizations bis hinunter zu 8 Bit. Um ein Modell auf 4 Bit oder weniger zu quantisieren, sind andere Ansätze erforderlich. Ein Ansatz ist Quantization Aware Training (QAT), das im Forward Pass des Trainings eine Fake Quantization anwendet. Fake Quantization führt die Affine Transformation durch, aber ihre Outputs werden nicht in Integers umgewandelt. Stattdessen bleiben sie Floating-Point Numbers (wenn die Quantization z. B. 17 ergeben würde, gibt die Fake Quantization 17.0 zurück). Solche Funktionen sind jedoch nicht differenzierbar. Daher müssen Tricks wie die Verwendung eines Straight-Through Estimators für die Gradients während des Backward Pass angewendet werden.

Neuere Ansätze wie AdaQuant erzielen eine bessere Post-Training Quantization Performance, indem sie die Quantization Layer für Layer durchführen, den optimalen Quantization Factor für jeden Layer auswählen und die Batch Normalization Layers nach der Quantization neu kalibrieren.

Experimente

Ich habe diese Quantization-Konzepte an den BERT- und ResNet-Modellen aus meinem vorherigen Artikel getestet. Für beide Modelle musste ich einen `QuantStub` und einen `DeQuantStub` am Anfang und Ende des Modells hinzufügen, um die Inputs zu quantisieren und zu dequantisieren. Da ich für beide Modelle die Static Quantization ausprobiert habe, musste ich zudem alle arithmetischen Operatoren in ihre quantisierten Äquivalente, die `FloatFunctional`-Operatoren, umwandeln. Dies war notwendig, da ich den Eager Quantization Mode verwendet habe. PyTorch ist derzeit dabei, die gesamte Arbeit zur Architekturoptimierung auf Torch AO zu verlagern, das eine graphbasierte Quantization bietet, die solche manuellen Änderungen überflüssig macht. Für das ResNet-Modell habe ich zudem versucht, alle Convolutional, Batch Normalization und ReLU Layers zu fusionieren. Dies führte jedoch zu einem erheblichen Verlust an Accuracy. Daher konnte ich nur die ersten beiden für alle außer der ersten Schicht fusionieren.

BERT

Für BERT zeige ich nur die Ergebnisse für die Dynamic Quantization, da die Static Quantization nicht funktionierte. In der Grafik unten können wir sehen, dass dies zu einer Beschleunigung von etwa 20 % über alle Batch Sizes hinweg führte.

Diese Beschleunigungen sind viel geringer als erwartet und führten auch dazu, dass das Modell Vorhersagen machte, die kaum genauer als zufälliges Raten waren. Die Ursache für beides ist wahrscheinlich ein Fehler, den ich bei der Eager Conversion von BERT in 8-Bit-Integers gemacht habe, da andere Arbeiten eine mehr als 3-fache Beschleunigung zeigten.

ResNet

Für ResNet habe ich Ergebnisse für die Static Quantization mit und ohne Layer Fusing. Die Beschleunigungen sind mit 39 % bzw. 34 % bei einer Batch Size von eins größer als bei BERT.

Wiederum sind die Beschleunigungen geringer als erwartet und stehen in keinem Verhältnis zu dem Aufwand, den die Einrichtung der Quantization dieser Modelle erforderte. Zumindest führte die Quantization beim ResNet-Modell nicht zu einer Performanceverschlechterung.

Fazit

PyTorch unterstützt die Quantization von Modellen. Die Verwendung der Eager-Implementierung erfordert jedoch die manuelle Anpassung des Modells für im schlimmsten Fall nur geringe Performancegewinne. Meine persönliche Empfehlung ist, zuerst die Modelllaufzeit zu optimieren, wie in Teil eins dieser Serie beschrieben. Nur wenn das Modell mit diesen Optimierungen noch nicht schnell genug läuft, würde ich eine Quantization in Betracht ziehen. Glücklicherweise kann die Quantization darauf aufbauend angewendet werden. Die Grafik unten zeigt beispielsweise, wie die Anwendung der Quantization auf die IPEX Optimization zu zusätzlichen Performancegewinnen führt.

Wenn Sie daran interessiert sind, diese Optimierungen selbst anzuwenden, werfen Sie einen Blick auf mein GitHub-Repository, das den Code für meine Experimente enthält: ML-Inference-Experiments.