Resumen Técnico: PagedWeight

Declaración del Problema

Los modelos de lenguaje de gran escala (LLM) de Mezcla de Expertos (MoE) ofrecen una alta eficiencia y precisión activando solo un subconjunto de expertos por token. Sin embargo, en escenarios de servicio, los modelos MoE enfrentan una tensión crítica de memoria: la GPU debe almacenar tanto los masivos pesos del modelo como el creciente caché de Clave-Valor (KV). Mientras que técnicas como PagedAttention gestionan la fragmentación del caché KV, no abordan la significativa huella de memoria de los pesos MoE, que pueden ocupar más del 60% de la memoria de la GPU.

Los métodos de cuantización existentes son mayoritariamente estáticos, fijando la precisión antes de que comience la inferencia. Aunque existen algunos métodos de adaptación de precisión en tiempo de ejecución, estos no utilizan los cambios de precisión para reasignar dinámicamente la memoria de la GPU entre los pesos y el caché KV. En consecuencia, existe una falta de sistemas que puedan liberar dinámicamente memoria de los pesos de los expertos menos críticos a medida que el caché KV se expande, sin sacrificar la precisión de la tarea ni interrumpir la inferencia.

Metodología: PagedWeight

Los autores proponen PagedWeight, un novedoso sistema de gestión para el servicio de MoE LLM que cuantiza dinámicamente los pesos del modelo en tiempo de ejecución para equilibrar la precisión de los pesos de los expertos con los tamaños del caché KV. El sistema trata los pesos de los expertos de manera similar a los bloques de caché KV paginados, permitiendo su descarga (offloading) y recarga dinámica.

Componentes Principales

1. Representación de Pesos Paginados (Paged Weight Representation):

   • PagedWeight opera a la granularidad de bloques lineales por experto (identificados por capa, experto y tipo de módulo: gate_up o down).
   • Utiliza la tecnología Any-Precision LLM (APL), representando los pesos en un formato de plano de bits superpuesto con tablas de búsqueda (LUTs) asociadas.
   • Una Tabla de Páginas de Pesos (Weight Page Table) rastrea el ancho de bits deseado ($d_i$) y el ancho de bits comprometido ($q_i$) para cada bloque lineal. Las páginas pueden existir en estados residentes en GPU, residentes en CPU o en estado de transferencia.

2. Planificador de Tiempo de Ejecución Consciente de la Calidad (Quality-Aware Runtime Planner):
   El planificador determina qué pesos cuantizar (reducir el ancho de bits) basándose en tres factores para minimizar la pérdida de precisión:

   • Sensibilidad Global Offline: Utiliza puntuaciones ponderadas por Hessiano para estimar la sensibilidad intrínseca de cada bloque lineal a la cuantización.
   • Estadísticas de Enrutamiento Online: Rastrea la "masa de enrutamiento" (frecuencia y peso de selección) de los expertos. Los expertos frecuentemente enrutados ("hot") son protegidos con suelos de ancho de bits más altos y multiplicadores de daño.
   • Residuales de Prompt: Un término de corrección específico del prompt predicho mediante cabezales de regresión lineal. Esto ajusta la estimación de sensibilidad global basada en la secuencia de entrada actual, reconociendo que el impacto de la cuantización varía según el prompt.

   El planificador calcula una puntuación de "daño" para las posibles reducciones de ancho de bits. Cuando la presión del caché KV activa un objetivo de memoria, el planificador selecciona de forma codiciosa las reducciones de menor daño (mayor calidad por byte ahorrado) para cumplir con el objetivo.

3. Ejecución Asíncrona y Optimización de Kernel:

   • Movimiento de Páginas Asíncrono: Para ocultar la latencia, el sistema descarga las páginas de pesos a la RAM de la CPU y las recarga de forma asíncrona. Los compromisos de ancho de bits solo se actualizan en límites seguros (por ejemplo, después de que un plan se ha ejecutado completamente), asegurando que la inferencia no sea interrumpida.
   • Kernel de Precisión Mixta Fusionado: Un kernel CUDA personalizado lee directamente los planos de bits y las LUTs de Any-Precision, soportando diferentes anchos de bits para cada bloque lineal dentro de una sola operación fusionada para reducir la sobrecarga.

Contribuciones Clave

1. Diseño del Sistema: La propuesta de PagedWeight, un sistema de pesos paginados para pesos MoE de Any-Precision en línea que gestiona anchos de bits deseados/comprometidos, estados de página y consumo de memoria de forma dinámica.
2. Algoritmo de Planificación: Un planificador de tiempo de ejecución consciente de la calidad que sintetiza la sensibilidad offline, las estadísticas de enrutamiento online y los residuales del prompt para seleccionar estrategias de reducción de pesos de bajo daño bajo presión de memoria.
3. Estrategia de Ejecución: Implementación de movimiento de páginas asíncrono para ocultar la latencia de descarga/recarga con una pérdida mínima de rendimiento (throughput), junto con un kernel MoE de precisión mixta fusionado.
4. Validación Empírica: Evaluación exhaustiva en tres modelos MoE (Qwen1.5-MoE-A2.7B, Mixtral-8×7B, Gemma-4-26B-A4B) demostrando un rendimiento superior sobre las líneas base existentes.

Resultados

Los autores evaluaron PagedWeight frente a FP16, Any-Precision LLM (APL), DP-LLM y MxMoE en tareas de modelado de lenguaje, razonamiento y contexto largo.

• Compromiso Calidad-Memoria: PagedWeight logra una precisión equivalente a FP16 con hasta un 72.0% de ahorro de memoria de GPU y una mejora de rendimiento (throughput) de 1.94× en comparación con las líneas base.
• Calidad con Presupuesto Fijo: Con presupuestos de memoria similares, PagedWeight mejora la calidad de la tarea sobre los métodos de cuantización hasta en un 39.3%, con una pérdida de rendimiento de como máximo un 4.1%.
• Rendimiento de Contexto Largo: Bajo presupuestos de memoria estrictos (por ejemplo, 10 GB), PagedWeight mantiene una calidad de contexto largo (Recuperación de Pasajes, NarrativeQA) comparable a FP16, mientras que las líneas base de cuantización estática (como APL) sufren una degradación significativa.
• Rendimiento (Throughput): PagedWeight iguala el rendimiento de las líneas base de cuantización uniforme, siendo la mayor caída observada del 4.1% con un tamaño de lote (batch size) de 4.
• Ablación: Eliminar componentes como las estadísticas de enrutamiento, los residuales del prompt o el movimiento de páginas degrada consistentemente el rendimiento, confirmando la necesidad del diseño completo del sistema.

Significancia

El artículo afirma que PagedWeight navega eficazmente el complejo compromiso entre la precisión de la tarea del modelo, el consumo de memoria y el rendimiento/latencia. Al tratar los pesos de los expertos como unidades gestionables y paginadas en lugar de activos estáticos, PagedWeight desbloquea un espacio de compromiso inexplorado entre la precisión de la tarea y la asignación de memoria de la GPU entre los pesos y el caché KV. El sistema demuestra que la cuantización dinámica y consciente de la calidad es una estrategia viable y superior para el servicio de MoE en comparación con la cuantización estática o los métodos que no se adaptan a la presión de memoria en tiempo de ejecución.