Unweight: Compresión Lossless de Pesos de LLM para Inferencia en GPU H100

La inferencia de modelos de lenguaje grandes (LLMs) en GPUs de alto rendimiento como las NVIDIA H100 está limitada no por la capacidad de cómputo de los tensor cores, sino por el ancho de banda de la memoria (HBM). Generar un token requiere leer la matriz completa de pesos del modelo de la HBM, y los tensor cores pueden procesar datos significativamente más rápido de lo que la memoria puede entregarlos. Este problema fundamental de la computación, la disparidad entre la velocidad de cómputo y la velocidad de acceso a la memoria, se agrava con el tamaño creciente de los LLMs.

Unweight aborda este cuello de botella mediante la compresión lossless de los pesos del modelo, específicamente los componentes de exponente de los valores BF16, que exhiben una alta redundancia. La innovación clave es realizar la descompresión directamente en la memoria compartida (SMEM) on-chip de la GPU y alimentar los datos descomprimidos a los tensor cores sin un viaje de ida y vuelta adicional a la HBM. Esto reduce el tráfico de datos a través del bus de memoria, acelerando la generación de tokens y permitiendo que más modelos se ajusten a una sola GPU, lo que se traduce en una inferencia más económica y rápida a escala de hyperscaler.

La relevancia de esta solución radica en la necesidad crítica de optimizar el uso de recursos de GPU, que son costosos en términos de hardware, memoria y consumo energético. A diferencia de la cuantificación (compresión con pérdida), Unweight mantiene la exactitud bit a bit de las salidas del modelo, lo cual es crucial para aplicaciones de producción que requieren un comportamiento predecible y de alta calidad.

Arquitectura del Sistema

Unweight se centra en la compresión de los pesos de las capas Multi-Layer Perceptron (MLP) de los LLMs, que constituyen la mayoría de los parámetros y el tráfico de memoria durante la generación de tokens. La arquitectura explota la redundancia en los 8 bits del exponente de los valores BF16, dejando el signo (1 bit) y la mantisa (7 bits) sin comprimir. Utiliza la codificación Huffman, un algoritmo clásico de compresión de entropía, para asignar códigos cortos a los exponentes más comunes (los 16 principales cubren >99% de los pesos) y códigos más largos a los raros. Las filas de pesos con exponentes fuera de la paleta de los 16 principales se almacenan verbatim para evitar ramificaciones por elemento en el hot path.

El sistema ofrece cuatro pipelines de ejecución adaptativos para la descompresión y multiplicación de matrices, seleccionados en tiempo de ejecución por un autotuner basado en el tamaño del batch y la forma de la matriz de pesos: (1) Decodificación completa + cuBLAS: reconstruye BF16 en HBM y usa cuBLAS estándar; (2) Decodificación solo de exponente + matmul reconstructivo: reduce el tráfico de preprocesamiento a la mitad; (3) Transcodificación a paleta de 4 bits + matmul reconstructivo: reduce el tráfico a un cuarto; (4) Paleta directa: pre-transcodifica a 4 bits en tiempo de carga del modelo, el kernel de matmul reconstruye BF16 al vuelo. Los pipelines (2), (3) y (4) utilizan un kernel de multiplicación de matrices reconstructivo personalizado que fusiona la descompresión con el cómputo. Este kernel carga datos comprimidos de HBM, reconstruye los valores BF16 en la memoria compartida (SMEM) y los alimenta directamente a las instrucciones WGMMA de los tensor cores de Hopper, evitando que los pesos reconstruidos existan en la memoria principal.

Dentro del kernel reconstructivo, los grupos de hilos de la GPU se dividen en productores (cargan entradas comprimidas de HBM a SMEM usando hardware de copia de memoria TMA, llenando un buffer circular) y consumidores (reconstruyen valores BF16 y los alimentan a los tensor cores). La asignación de unidades de cómputo (SMs) entre el kernel de preprocesamiento (decodificación Huffman o transcodificación de paleta) y el kernel de matmul reconstructivo es un parámetro tunable, ya que compiten por la SMEM. Unweight también emplea un pipelining entre capas, decodificando capas 'duras' (con Huffman) en segundo plano mientras se computan capas 'fáciles' (con datos de paleta pre-transcodificados), utilizando CUDA streams separados y doble buffering.

Fundamentos Teóricos

El principio subyacente a Unweight se basa en la teoría de la información y la codificación de entropía, específicamente la codificación Huffman, un algoritmo desarrollado por David A. Huffman en 1952. Este algoritmo es óptimo para la compresión de datos sin pérdida cuando se conocen las frecuencias de los símbolos, asignando códigos de longitud variable donde los símbolos más frecuentes tienen códigos más cortos. La observación de que los exponentes de los pesos BF16 en LLMs tienen una distribución altamente sesgada, con solo un puñado de valores dominantes, es la base empírica que permite la aplicación efectiva de Huffman.

La idea de fusionar la descompresión con el cómputo para evitar viajes de ida y vuelta a la memoria principal se alinea con principios de optimización de localidad de datos y reducción de la latencia de memoria, conceptos explorados en la arquitectura de computadoras y sistemas de alto rendimiento durante décadas. La gestión de la jerarquía de memoria, desde registros hasta caché L1/L2/L3 y memoria principal (HBM en GPUs), es un tema central en la optimización del rendimiento. La estrategia de Unweight de usar la memoria compartida (SMEM) como una caché de descompresión justo antes de los tensor cores es una aplicación directa de estos principios para mitigar el 'memory wall' o 'memory bottleneck' que limita el rendimiento de los procesadores modernos, especialmente en cargas de trabajo intensivas en datos como la inferencia de LLMs.