Despliegue eficiente de modelos de lenguaje grandes de baja precisión en CPU
La innovación clave de T-MAC radica en adoptar un paradigma de cálculo basado en tablas de búsqueda (LUT) en lugar del tradicional paradigma de multiplicación-acumulación (MAC). T-MAC utiliza tablas de búsqueda para soportar directamente cálculos de baja precisión, eliminando así las operaciones de descuantificación necesarias en otros sistemas y reduciendo significativamente el número de operaciones de multiplicación y suma.
Los experimentos han demostrado un rendimiento excepcional de T-MAC: en un PC Surface AI equipado con el último chipset Qualcomm Snapdragon X E lite, el modelo BitNet-b1.58 de 3B puede generar 48 tokens por segundo, el modelo llama de 2 bits y 7B puede generar 30 tokens por segundo, y el modelo llama de 4 bits y 7B puede generar 20 tokens por segundo.
¡Esto incluso supera el rendimiento de las NPU!
Al desplegar el modelo llama-2-7b-4bit, aunque la NPU puede generar 10.4 tokens por segundo, la CPU con la ayuda de T-MAC puede alcanzar 12.6 tokens por segundo usando solo dos núcleos, llegando incluso a 22 tokens por segundo en su máximo rendimiento.
Todo esto supera con creces la velocidad de lectura promedio humana y representa una mejora de 4 a 5 veces en comparación con el marco original llama.cpp.
Incluso en dispositivos de gama más baja como la Raspberry Pi 5, T-MAC puede alcanzar una tasa de generación de 11 tokens por segundo para el BitNet-b1.58 de 3B. T-MAC también tiene ventajas significativas en cuanto al consumo de energía: para lograr la misma tasa de generación, T-MAC requiere solo de 1/4 a 1/6 del número de núcleos en comparación con el llama.cpp original, reduciendo el consumo de energía y dejando recursos computacionales para otras aplicaciones.
Es importante destacar que el rendimiento computacional de T-MAC aumenta linealmente a medida que disminuye el número de bits, un fenómeno difícil de observar en GPUs y NPUs basadas en descuantificación. Sin embargo, T-MAC puede lograr 10 tokens por segundo en un solo núcleo y 28 tokens por segundo en cuatro núcleos con 2 bits, superando significativamente el rendimiento de las NPU.
La multiplicación de matrices no necesita multiplicación, solo búsqueda en tabla (LUT)
Para parámetros de baja precisión (pesos), T-MAC agrupa cada bit por separado (por ejemplo, un grupo de 4 bits), multiplica estos bits con el vector de activación, precalcula todas las sumas parciales posibles y luego las almacena usando LUT.
Posteriormente, T-MAC utiliza operaciones de desplazamiento y acumulación para soportar números de bits escalables de 1 a 4. Con este método, T-MAC abandona las instrucciones FMA (multiplicación-suma) ineficientes en CPU y en su lugar utiliza instrucciones TBL/PSHUF (búsqueda en tabla) de menor consumo y mayor eficiencia.
Cálculo centrado en bits, reemplazando el cálculo centrado en tipos de datos
El cálculo tradicional basado en descuantificación es en realidad un cálculo centrado en tipos de datos, lo que requiere personalización para cada tipo de dato diferente.
Cada combinación de ancho de bits de activación y peso, como W4A16 (peso int4, activación float16) y W2A8, requiere un diseño específico de layout de pesos y núcleo de cálculo.
Por ejemplo, el layout W3 necesita empaquetar 2 bits y 1 bit por separado, y utilizar diferentes métodos de entrelazado o mezcla para alineación de memoria o decodificación rápida. Luego, el núcleo de cálculo correspondiente necesita desempaquetar este layout específico a tipos de datos soportados por el hardware para su ejecución.
T-MAC, por otro lado, observa el cálculo de multiplicación de matrices de baja precisión desde la perspectiva de los bits, diseñando una estructura de datos óptima para un solo bit y luego expandiéndola a 2/3/4 bits mediante apilamiento.
Al mismo tiempo, para vectores de activación de diferentes precisiones (float16/float32/int8), solo el proceso de construcción de la tabla necesita cambiar, sin necesidad de considerar diferentes estructuras de datos durante la búsqueda en tabla.
Además, en los métodos tradicionales basados en descuantificación, al reducir de 4 bits a 3/2/1 bits, aunque el uso de memoria es menor, la cantidad de cálculos no disminuye, y el rendimiento puede incluso empeorar debido al aumento de la sobrecarga de descuantificación.
Sin embargo, la cantidad de cálculos de T-MAC disminuye linealmente con la reducción del número de bits, proporcionando así una mejor aceleración con menos bits y ofreciendo una solución de despliegue eficiente para los últimos modelos de 1 bit/2 bits como BitNet y EfficientQAT.
Implementación altamente optimizada de operadores
El cálculo centrado en bits tiene muchas ventajas, pero implementarlo en CPU sigue siendo un desafío considerable:
(1) En comparación con el acceso secuencial a datos de activación y pesos, el acceso a las tablas es aleatorio. La residencia de las tablas en la memoria rápida on-chip es crucial para el rendimiento final de inferencia;
(2) Sin embargo, la memoria on-chip es limitada, y el método de tabla de búsqueda (LUT) aumenta el uso de memoria on-chip en comparación con el mpGEMV tradicional. Esto se debe a que la tabla de búsqueda necesita guardar los resultados de la multiplicación del vector de activación con todos los posibles patrones de bits, lo cual es mucho más que la activación en sí.
Para abordar esto, los investigadores del Microsoft Research Asia exploraron en profundidad el flujo de datos de cálculo basado en tablas de búsqueda y diseñaron estructuras de datos eficientes y flujos de cálculo para este paradigma de cálculo, incluyendo:
-
Almacenar LUT en memoria on-chip para mejorar el rendimiento de acceso aleatorio a memoria utilizando instrucciones vectoriales de búsqueda en tabla (TBL/PSHUF) en CPU.
-
Cambiar el orden de cálculo del eje de la matriz para maximizar la reutilización de datos de las LUT limitadas que se pueden colocar en la memoria on-chip.
-
Diseñar un método óptimo de división de matrices (Tiling) específicamente para la búsqueda en tabla, combinado con la búsqueda de parámetros de división óptimos mediante autotvm.
-
Optimización del layout de los parámetros (pesos):
a) Reordenamiento de pesos para maximizar el acceso secuencial y mejorar la tasa de aciertos de caché.
b) Entrelazado de pesos para mejorar la eficiencia de decodificación.
- Optimizaciones específicas para CPUs Intel/ARM, incluyendo:
a) Reordenamiento de registros para construir rápidamente tablas de búsqueda.
b) Uso de instrucciones de promedio para realizar acumulaciones rápidas de 8 bits.
Los investigadores aplicaron gradualmente varias optimizaciones sobre una implementación básica, logrando finalmente una aceleración significativa en comparación con los operadores de baja precisión del estado del arte: