Multiplicaciones Rápidas de Matrices para LLMs Cuantificados con Tablas de Búsqueda.

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

Resumen

La implementación de modelos de lenguaje grandes (LLMs, por sus siglas en inglés) suele estar limitada por el ancho de banda de memoria, donde el cuello de botella principal es el costo de transferir los parámetros del modelo desde la memoria global de la GPU a sus registros. Al combinar esto con núcleos personalizados que fusionan las operaciones de descuantificación y multiplicación de matrices, la cuantificación solo de pesos puede permitir una inferencia más rápida al reducir la cantidad de movimiento de memoria. Sin embargo, desarrollar núcleos de alto rendimiento para LLMs cuantificados por pesos presenta desafíos sustanciales, especialmente cuando los pesos se comprimen a anchos de bits no uniformemente divisibles (por ejemplo, 3 bits) con cuantificación de tabla de búsqueda no uniforme. Este artículo describe FLUTE, un motor de tabla de búsqueda flexible para LLMs cuantificados por LUT, que utiliza una reestructuración offline de la matriz de pesos cuantificados para minimizar las manipulaciones de bits asociadas con el desempaquetado, y la vectorización y duplicación de la tabla de búsqueda para mitigar las restricciones de ancho de banda de memoria compartida. Con tamaños de lote < 32 y un tamaño de grupo de cuantificación de 128 (típico en la inferencia de LLM), el núcleo FLUTE puede ser de 2 a 4 veces más rápido que los núcleos GEMM existentes. Como aplicación de FLUTE, exploramos una extensión simple a la cuantificación NormalFloat basada en tabla de búsqueda y la aplicamos para cuantificar LLaMA3 en varias configuraciones, obteniendo un rendimiento de cuantificación competitivo frente a bases sólidas y un aumento de rendimiento de extremo a extremo de 1.5 a 2 veces.

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.