SpenseGPT: Practical One-shot Pruning Enabling Sparse and Dense GEMMs for LLM Inference

Resumen Técnico: SpenseGPT: Poda de un solo paso práctica para GEMMs dispersos y densos en la inferencia de LLM

1. Planteamiento del Problema

La dispersión semiestructurada 2:4 (donde dos de cada cuatro pesos consecutivos se podan a cero) cuenta con un amplio soporte en aceleradores modernos (por ejemplo, NVIDIA B200, GPUs AMD, Meta MTIAs) y ofrece una aceleración teórica de 2× al permitir que el hardware omita las computaciones de pesos con valor cero. Sin embargo, la adopción práctica para la inferencia de Modelos de Lenguaje Extensos (LLM) enfrenta dos obstáculos principales:

1. Degradación Estricta de la Precisión: La restricción rígida de dispersión 2:4 a menudo conduce a una pérdida significativa de precisión en la poda post-entrenamiento, particularmente en tareas de razonamiento y seguimiento de instrucciones, a menos que se realice un reentrenamiento extensivo.
2. Impracticabilidad de las Relajaciones Existentes: Trabajos recientes que intentan relajar la restricción 2:4 (por ejemplo, PATCH, SlideSparse) sufren limitaciones críticas. PATCH requiere soporte de compiladores especializados (Triton) que aún no está ampliamente disponible, mientras que SlideSparse requiere la expansión de la activación de entrada, lo que introduce sobrecargas de tiempo de ejecución que niegan las aceleraciones de extremo a extremo, especialmente para longitudes de secuencia largas ($M$).

2. Metodología

El artículo propone Spense, un formato híbrido denso-disperso práctico, y SpenseGPT, un método de poda post-entrenamiento de un solo paso construido sobre este.

A. El Formato Spense
En lugar de imponer una dispersión 2:4 en toda la matriz de pesos, Spense divide cada matriz de pesos en dos regiones contiguas:

• Una región densa (que retiene el $p\%$ de los pesos).
• Una región dispersa 2:4 (podada con la dispersión 2:4 estándar).

Este diseño aprovecha las bibliotecas GEMM altamente optimizadas existentes (por ejemplo, cuSPARSELt para dispersos y cuBLASLt para densos en GPUs NVIDIA) sin requerir compiladores personalizados o expansión de la activación de entrada. Los autores implementan dos variantes:

• División de salida (Output-split): Divide el tensor de salida, utilizando GEMM disperso para una parte y GEMM denso para la otra.
• División de entrada (Input-split): Divide el tensor de entrada, procesando la primera parte con GEMM disperso y la segunda con GEMM denso (con acumulación).
  Para las variantes de GLU en los MLPs, los autores utilizan la división de salida para $W_{up}$ y $W_{gate}$, y la división de entrada para $W_{down}$.

B. Estrategias de Selección de Índices Densos
Un conocimiento crítico es que qué índices intermedios se mantienen densos impacta significativamente la calidad del modelo. El artículo identifica que los índices intermedios en los MLPs de estilo GLU pueden permutarse sin cambiar la salida matemática, lo que permite la selección de índices específicos para que sean densos antes de reorganizarlos en un bloque contiguo para su ejecución. Se proponen dos estrategias:

1. SpenseGPT (Estrategia Simple): Selecciona el último $p\%$ de los índices intermedios para que permanezcan densos. Esto se alinea con el orden de procesamiento por columnas de SparseGPT, donde las columnas procesadas se congelan y los errores de reconstrucción se compensan actualizando las columnas restantes hacia la derecha. Mantener las columnas densas a la derecha proporciona un "buffer de compensación".
2. SpenseGPT+ (Estrategia Basada en Pérdida): Selecciona los índices densos basándose en una estimación de la pérdida de reconstrucción. Inspirado en Wanda, calcula un puntaje de importancia para cada índice intermedio agregando contribuciones a través de las proyecciones de puerta (gate), up y down. El puntaje combina la magnitud del peso y la magnitud de la activación de entrada ($S_{ij} = |W_{ij}|\sqrt{H_{j,j}}$). Los índices con los puntajes más altos se seleccionan para permanecer densos, y luego se permutan al bloque contiguo antes de aplicar SparseGPT a las regiones dispersas restantes.

3. Contribuciones Clave

• Formato Spense: Un formato híbrido denso-disperso que relaja la estricta dispersión 2:4 manteniendo la compatibilidad con las bibliotecas de GEMM de alto rendimiento existentes, evitando el soporte de compiladores personalizados y la sobrecarga de expansión de activación.
• Importancia de la Selección de Densidad: Demostración de que la selección de los índices densos es un factor crítico para mantener la calidad del modelo, ya que una selección ingenua conduce a caídas severas de rendimiento.
• Estrategias de Poda: Introducción de SpenseGPT (selección de cola contigua) y SpenseGPT+ (selección basada en la pérdida de reconstrucción) para la poda de un solo paso.
• Aceleración en el Mundo Real: La primera demostración de la poda de un solo paso que logra aceleraciones de decodificación de LLM de extremo a extremo en el mundo real utilizando núcleos tensoriales dispersos semiestructurados en GPUs recientes (B200) preservando la calidad del modelo.

4. Resultados Experimentales

Los métodos fueron evaluados en Qwen3-32B y Seed-OSS-36B utilizando precisión FP8 en GPUs B200 con el framework VLLM. Las pruebas de rendimiento incluyeron AIME (matemáticas), GPQA (conocimiento), IFEval (seguimiento de instrucciones) y LiveCodeBench (codificación).

• Preservación de la Precisión:
  • La poda estricta 2:4 (50% de dispersión) causó caídas significativas de precisión (por ejemplo, el puntaje promedio de Seed-OSS-36B cayó de 73.86 a 65.18 con SparseGPT).
  • SpenseGPT+ con 25% de dispersión logró un puntaje promedio de 73.12 en Seed-OSS-36B, igualando casi al baseline denso (73.86).
  • En Qwen3-32B, SpenseGPT+ con 25% de dispersión alcanzó 70.76, cercano al baseline denso de 71.68.
  • Las comparaciones mostraron que SpenseGPT+ superó significativamente a SpenseGPT y otros baselines (HyperPrune, Wanda) al seleccionar correctamente los índices densos.
• Aceleración de Extremo a Extremo:
  • Spense logró hasta un 1.2× de aceleración de decodificación de extremo a extremo sobre el baseline denso en GPUs B200.
  • En contraste, SlideSparse no logró proporcionar una aceleración significativa debido a las sobrecargas de no-GEMM derivadas de la expansión de la activación de entrada.
• Eficiencia:
  • La poda de modelos de 32B–36B tomó aproximadamente 1.33–2.15 horas en una sola GPU B200, comparable a SparseGPT y mucho más rápido que HyperPrue. ProxSparse se quedó sin memoria en estos modelos.

5. Significado y Reivindicaciones

El artículo afirma presentar la primera demostración de poda de un solo paso de aceleraciones de decodificación de LLM de extremo a extremo en el mundo real derivadas de núcleos tensoriales dispersos semiestructurados en hardware reciente (específicamente GPUs B200) manteniendo la calidad del modelo.

Los autores enfatizan que su enfoque cierra la brecha entre la aceleración teórica de hardware y el despliegue práctico. Al evitar la necesidad de compiladores especializados o expansión de entrada, Spense permite la utilización inmediata de las bibliotecas optimizadas existentes. El trabajo destaca que simplemente relajar la restricción de dispersión no es suficiente; la estrategia para seleccionar qué pesos permanecen densos es primordial para preservar las capacidades de razonamiento y seguimiento de instrucciones del modelo.

Limitaciones Reconocidas:

• La aceleración depende de la eficiencia de las bibliotecas GEMM subyacentes y varía con el tamaño del lote (batch size), la longitud de la secuencia y el hardware.
• SpenseGPT+ depende de datos de calibración; si la distribución de despliegue difiere significativamente del conjunto de calibración, los índices densos seleccionados pueden no ser óptimos.