Treinamento de Transformers com Inteiros de 4 bits

Training Transformers with 4-bit Integers

June 21, 2023

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

cs.AI

Resumo

Quantizar as ativações, pesos e gradientes para 4 bits é uma abordagem promissora para acelerar o treinamento de redes neurais. No entanto, os métodos existentes de treinamento em 4 bits exigem formatos numéricos personalizados que não são suportados pelo hardware contemporâneo. Neste trabalho, propomos um método de treinamento para transformadores com todas as multiplicações de matrizes implementadas usando aritmética INT4. Treinar com uma precisão ultrabaixa INT4 é desafiador. Para alcançar isso, analisamos cuidadosamente as estruturas específicas das ativações e gradientes em transformadores para propor quantizadores dedicados a eles. Para a propagação direta, identificamos o desafio dos valores discrepantes (outliers) e propomos um quantizador de Hadamard para suprimir esses valores. Para a retropropagação, aproveitamos a esparsidade estrutural dos gradientes propondo técnicas de divisão de bits e amostragem por pontuação de alavancagem para quantizar os gradientes com precisão. Nosso algoritmo alcança precisão competitiva em uma ampla gama de tarefas, incluindo compreensão de linguagem natural, tradução automática e classificação de imagens. Diferentemente dos métodos anteriores de treinamento em 4 bits, nosso algoritmo pode ser implementado na geração atual de GPUs. Nossa implementação prototípica de operador linear é até 2,2 vezes mais rápida que as contrapartes FP16 e acelera o treinamento em até 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%.