Multiplicações Rápidas de Matrizes para LLMs Quantizados em Tabela de Pesquisa

Fast Matrix Multiplications for Lookup Table-Quantized LLMs

July 15, 2024

Autores: Han Guo, William Brandon, Radostin Cholakov, Jonathan Ragan-Kelley, Eric P. Xing, Yoon Kim

cs.AI

Resumo

A implantação de grandes modelos de linguagem (LLMs) muitas vezes é limitada pela largura de banda de memória, onde o principal gargalo é o custo de transferir os parâmetros do modelo da memória global da GPU para seus registradores. Quando combinada com kernels personalizados que fundem as operações de desquantização e multiplicação de matriz, a quantização apenas de peso pode possibilitar inferências mais rápidas ao reduzir a quantidade de movimentação de memória. No entanto, desenvolver kernels de alto desempenho para LLMs quantizados por peso apresenta desafios substanciais, especialmente quando os pesos são comprimidos em larguras de bits não uniformemente divisíveis (por exemplo, 3 bits) com quantização de tabela de pesquisa não uniforme. Este artigo descreve o FLUTE, um mecanismo flexível de tabela de pesquisa para LLMs quantizados por LUT, que utiliza reestruturação offline da matriz de peso quantizada para minimizar manipulações de bits associadas à descompactação, e vetorização e duplicação da tabela de pesquisa para mitigar as restrições de largura de banda de memória compartilhada. Com tamanhos de lote < 32 e tamanho de grupo de quantização de 128 (típico na inferência de LLM), o kernel FLUTE pode ser de 2 a 4 vezes mais rápido do que os kernels GEMM existentes. Como aplicação do FLUTE, exploramos uma extensão simples à quantização NormalFloat baseada em tabela de pesquisa e a aplicamos para quantizar o LLaMA3 em várias configurações, obtendo desempenho de quantização competitivo em relação a bases sólidas, enquanto alcançamos um aumento de throughput de ponta a ponta de 1,5 a 2 vezes.

English

The deployment of large language models (LLMs) is often constrained by memory bandwidth, where the primary bottleneck is the cost of transferring model parameters from the GPU's global memory to its registers. When coupled with custom kernels that fuse the dequantization and matmul operations, weight-only quantization can thus enable faster inference by reducing the amount of memory movement. However, developing high-performance kernels for weight-quantized LLMs presents substantial challenges, especially when the weights are compressed to non-evenly-divisible bit widths (e.g., 3 bits) with non-uniform, lookup table (LUT) quantization. This paper describes FLUTE, a flexible lookup table engine for LUT-quantized LLMs, which uses offline restructuring of the quantized weight matrix to minimize bit manipulations associated with unpacking, and vectorization and duplication of the lookup table to mitigate shared memory bandwidth constraints. At batch sizes < 32 and quantization group size of 128 (typical in LLM inference), the FLUTE kernel can be 2-4x faster than existing GEMM kernels. As an application of FLUTE, we explore a simple extension to lookup table-based NormalFloat quantization and apply it to quantize LLaMA3 to various configurations, obtaining competitive quantization performance against strong baselines while obtaining an end-to-end throughput increase of 1.5 to 2 times.