Entraînement des Transformers avec des entiers sur 4 bits

Training Transformers with 4-bit Integers

June 21, 2023

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

cs.AI

Résumé

La quantification des activations, poids et gradients en 4 bits est prometteuse pour accélérer l'entraînement des réseaux de neurones. Cependant, les méthodes existantes d'entraînement en 4 bits nécessitent des formats numériques personnalisés qui ne sont pas supportés par le matériel contemporain. Dans ce travail, nous proposons une méthode d'entraînement pour les transformers où toutes les multiplications matricielles sont implémentées avec l'arithmétique INT4. L'entraînement avec une précision INT4 ultra-faible est un défi. Pour y parvenir, nous analysons attentivement les structures spécifiques des activations et des gradients dans les transformers afin de proposer des quantificateurs dédiés. Pour la propagation avant, nous identifions le problème des valeurs aberrantes et proposons un quantificateur de Hadamard pour les supprimer. Pour la rétropropagation, nous exploitons la sparsité structurelle des gradients en proposant des techniques de division de bits et d'échantillonnage par score de levier pour quantifier les gradients avec précision. Notre algorithme atteint une précision compétitive sur un large éventail de tâches, notamment la compréhension du langage naturel, la traduction automatique et la classification d'images. Contrairement aux méthodes précédentes d'entraînement en 4 bits, notre algorithme peut être implémenté sur la génération actuelle de GPU. Notre implémentation prototype d'opérateur linéaire est jusqu'à 2,2 fois plus rapide que les équivalents en FP16 et accélère l'entraînement jusqu'à 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%.