Training von Transformern mit 4-Bit-Ganzzahlen

Training Transformers with 4-bit Integers

June 21, 2023

Autoren: Haocheng Xi, Changhao Li, Jianfei Chen, Jun Zhu

cs.AI

Zusammenfassung

Die Quantisierung von Aktivierungen, Gewichten und Gradienten auf 4 Bit ist vielversprechend, um das Training neuronaler Netzwerke zu beschleunigen. Allerdings erfordern bestehende 4-Bit-Trainingsmethoden benutzerdefinierte numerische Formate, die von aktueller Hardware nicht unterstützt werden. In dieser Arbeit schlagen wir eine Trainingsmethode für Transformer vor, bei der alle Matrixmultiplikationen mit INT4-Arithmetik implementiert sind. Das Training mit einer ultra-niedrigen INT4-Präzision ist jedoch herausfordernd. Um dies zu erreichen, analysieren wir sorgfältig die spezifischen Strukturen von Aktivierungen und Gradienten in Transformern und schlagen dedizierte Quantisierer für diese vor. Für die Vorwärtspropagation identifizieren wir das Problem von Ausreißern und schlagen einen Hadamard-Quantisierer vor, um diese zu unterdrücken. Für die Rückwärtspropagation nutzen wir die strukturelle Sparsity von Gradienten, indem wir Bit-Splitting und Leverage-Score-Sampling-Techniken vorschlagen, um Gradienten präzise zu quantisieren. Unser Algorithmus erreicht eine wettbewerbsfähige Genauigkeit bei einer Vielzahl von Aufgaben, einschließlich Natural Language Understanding, maschineller Übersetzung und Bildklassifizierung. Im Gegensatz zu früheren 4-Bit-Trainingsmethoden kann unser Algorithmus auf der aktuellen Generation von GPUs implementiert werden. Unsere prototypische Implementierung des linearen Operators ist bis zu 2,2-mal schneller als die FP16-Pendants und beschleunigt das Training um bis zu 35,1 %.

English

Quantizing the activation, weight, and gradient to 4-bit is promising to accelerate neural network training. However, existing 4-bit training methods require custom numerical formats which are not supported by contemporary hardware. In this work, we propose a training method for transformers with all matrix multiplications implemented with the INT4 arithmetic. Training with an ultra-low INT4 precision is challenging. To achieve this, we carefully analyze the specific structures of activation and gradients in transformers to propose dedicated quantizers for them. For forward propagation, we identify the challenge of outliers and propose a Hadamard quantizer to suppress the outliers. For backpropagation, we leverage the structural sparsity of gradients by proposing bit splitting and leverage score sampling techniques to quantize gradients accurately. Our algorithm achieves competitive accuracy on a wide range of tasks including natural language understanding, machine translation, and image classification. Unlike previous 4-bit training methods, our algorithm can be implemented on the current generation of GPUs. Our prototypical linear operator implementation is up to 2.2 times faster than the FP16 counterparts and speeds up the training by up to 35.1%.