Обучение трансформаторов с использованием 4-битных целых чисел

Training Transformers with 4-bit Integers

June 21, 2023

Авторы: Haocheng Xi, Changhao Li, Jianfei Chen, Jun Zhu

cs.AI

Аннотация

Квантование активаций, весов и градиентов до 4 бит является перспективным подходом для ускорения обучения нейронных сетей. Однако существующие методы обучения с 4-битной точностью требуют использования специализированных числовых форматов, которые не поддерживаются современным оборудованием. В данной работе мы предлагаем метод обучения трансформеров, в котором все матричные умножения реализованы с использованием арифметики INT4. Обучение с ультранизкой точностью INT4 представляет собой сложную задачу. Для её решения мы тщательно анализируем специфические структуры активаций и градиентов в трансформерах и предлагаем специализированные квантователи для них. Для прямого распространения мы выявляем проблему выбросов и предлагаем квантователь Хаара для их подавления. Для обратного распространения мы используем структурную разреженность градиентов, предлагая методы битового разделения и сэмплирования на основе весов для точного квантования градиентов. Наш алгоритм демонстрирует конкурентоспособную точность на широком спектре задач, включая понимание естественного языка, машинный перевод и классификацию изображений. В отличие от предыдущих методов обучения с 4-битной точностью, наш алгоритм может быть реализован на текущем поколении графических процессоров. Наша прототипная реализация линейного оператора работает до 2,2 раз быстрее, чем аналоги с точностью FP16, и ускоряет обучение до 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%.