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%.
翻译:将激活值、权重和梯度量化至4位,有望加速神经网络训练。然而,现有的4位训练方法需要定制数值格式,而当代硬件并不支持这些格式。本文提出一种训练方法,使Transformer模型中的所有矩阵乘法均采用INT4算术实现。以超低精度的INT4进行训练极具挑战性。为达成目标,我们通过细致分析Transformer中激活值和梯度的特定结构,提出了针对性的量化器。在前向传播中,我们识别出异常值带来的挑战,并提出一种哈达玛量化器来抑制异常值。在反向传播中,我们利用梯度的结构稀疏性,通过引入位拆分与杠杆分数采样技术来精确量化梯度。我们的算法在自然语言理解、机器翻译和图像分类等广泛任务上取得了具有竞争力的精度。与先前的4位训练方法不同,本算法可在当前世代的GPU上实现。我们实现的线性算子原型相比FP16算子最高加速2.2倍,并将训练速度提升高达35.1%。