Entrenamiento de Transformers con enteros de 4 bits

Training Transformers with 4-bit Integers

June 21, 2023

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

cs.AI

Resumen

La cuantización de las activaciones, pesos y gradientes a 4 bits es prometedora para acelerar el entrenamiento de redes neuronales. Sin embargo, los métodos existentes de entrenamiento en 4 bits requieren formatos numéricos personalizados que no son compatibles con el hardware actual. En este trabajo, proponemos un método de entrenamiento para transformadores en el que todas las multiplicaciones de matrices se implementan con aritmética INT4. Entrenar con una precisión ultra baja de INT4 es un desafío. Para lograrlo, analizamos cuidadosamente las estructuras específicas de las activaciones y los gradientes en los transformadores y proponemos cuantizadores dedicados para ellos. Para la propagación hacia adelante, identificamos el desafío de los valores atípicos y proponemos un cuantizador de Hadamard para suprimir estos valores. Para la retropropagación, aprovechamos la dispersión estructural de los gradientes mediante técnicas de división de bits y muestreo por puntuación de influencia para cuantizar los gradientes con precisión. Nuestro algoritmo logra una precisión competitiva en una amplia gama de tareas, incluyendo comprensión del lenguaje natural, traducción automática y clasificación de imágenes. A diferencia de los métodos anteriores de entrenamiento en 4 bits, nuestro algoritmo puede implementarse en la generación actual de GPUs. Nuestra implementación prototípica del operador lineal es hasta 2.2 veces más rápida que sus contrapartes en FP16 y acelera el entrenamiento hasta en un 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%.