Implementación de esparcidad 2:4 con FP8 en Hopper: Un Recetario de Producción

En GPUs de clase Hopper, la combinación de esparcidad estructurada 2:4 y precisión FP8 ofrece habitualmente mejoras de rendimiento de 1.5–2.0× para la decodificación de LLM, mientras reduce la energía por token en un 20–40% cuando se implementa en motores de producción como TensorRT-LLM. La perspectiva a nivel de kernel es simple: los núcleos de tensor dispersos de NVIDIA duplican el rendimiento matemático efectivo para GEMMs elegibles bajo el patrón N:M 2:4, y el motor Transformer con canalización FP8 reduce el ancho de banda y escala las activaciones de manera robusta. Pero hacer que esas mejoras se mantengan en producción—sin comprometer la calidad—requiere una canalización cuidadosa: las capas correctas, calibración bien curada, poda segura, rápida recuperación con adaptadores y un motor bien afinado.

Este artículo es un manual práctico para llevar un modelo denso a una construcción podada, recuperada con adaptadores y calibrada en FP8 en TensorRT-LLM. Te mostraremos cómo establecer una línea base y crear un perfil, seleccionar capas y formatos compatibles con canalizaciones 2:4, aplicar poda estructurada en caminos lineales/FFN, recuperar con LoRA/AdaLoRA, integrar en un motor de producción con formas estáticas y atención fusionada, y validar con una matriz robusta de métricas, distribuciones de latencia, memoria y registro de consumo de energía. Al final, tendrás un proceso repetible y seguro de retroceso que puedes operar en una flota real.

Detalles de Arquitectura/Implementación

Resumen de la canalización de extremo a extremo

1. Línea de base en condiciones de producción

• Construye una línea base densa FP16 o FP8 en TensorRT-LLM con tus parámetros de decodificación de producción (lotes, longitudes máximas de indicación/respuesta, tokenizador, programador).
• Habilita la atención fusionada (por ejemplo, núcleos al estilo FlashAttention) para que las multiplicaciones de matrices MLP se conviertan en el camino de cálculo dominante que luego esparsificarás.
• Mide tokens/s, latencia p50/p99, memoria de activación/pico y consumo de energía de la GPU (nvidia-smi dmon o NVML). Este es tu control.

2. Elige capas elegibles para 2:4 y orientación

• Dirige matrices lineales/FFN con grandes dimensiones K: las proyecciones QKV/proj de atención y las proyecciones hacia arriba/abajo del MLP son generalmente elegibles.
• Evita tablas de embedding, normalizaciones de capa y la cabecera final del LM para la poda.
• Respeta la agrupación 2:4 a lo largo de la dimensión de cálculo más interna esperada por cuSPARSELt (grupos de 4 valores contiguos a lo largo de K) para que los núcleos de tensor dispersos puedan activarse.

3. Calibra para FP8 y esparcidad

• Cura un conjunto de calibración de 2-5k indicaciones que reflejen la producción: instrucciones, razonamiento y código. Incluye un segmento de largo contexto.
• Ejecuta escalamiento de activación/peso FP8 con el motor Transformer (escalamiento dinámico por tensor) para evitar desbordamientos mientras recoges estadísticas de activación.

4. Aplica poda estructurada 2:4 de manera segura

• Usa puntuaciones basadas en magnitud o activación dentro de cada grupo de 4 valores para eliminar los dos elementos menos importantes.
• Comienza con una esparcidad global del 30-40% (2:4 es 50% dentro de matrices elegibles, pero solo se aplica a capas seleccionadas). Evita podas agresivas en capas de atención tardías que pueden ser críticas para el comportamiento a largo alcance.

5. Recuperación basada en adaptadores (LoRA/AdaLoRA)

• Conecta adaptadores de bajo rango a los módulos podados y afina durante unos pocos miles de pasos en una mezcla de tareas similar a tu conjunto de calibración.
• Usa detención temprana basada en una pequeña cesta de validación (por ejemplo, segmentos de MMLU/GSM8K/MT-Bench/HumanEval) para recuperar 0.5–2 puntos sin costoso ajuste fino completo.

6. Canalización de precisión y recalibración post-poda

• Establece primero una línea base FP8 estable, luego poda y recalibra las escalas FP8 posterior a la poda.
• Si prefieres INT8 (W8A8 o solo peso), crea una línea base pre-poda (LLM.int8(), GPTQ), poda, luego recalibra y valida.

7. Integración en el motor para producción

• Exporta pesos podados con máscaras 2:4.
• Construye un motor TensorRT-LLM con FP8 y esparcidad habilitada, perfiles de forma estática cubriendo tu envolvente de indicación/respuesta y atención fusionada habilitada.
• Verifica el compromiso de núcleos dispersos en todas las formas configuradas; desajustes de forma deshabilitarán silenciosamente la esparcidad.

8. Valida e iterar

• Mide de nuevo tokens/s, latencia p50/p99, memoria y consumo de energía. Espera ganancias de 1.3–1.8× solo de 2:4 y de 1.5–2.0× con FP8 en H100/H200, con reducciones del 20–40% en energía por token.
• Evalúa MMLU, GSM8K, HumanEval, MT-Bench, BBH y una tarea de largo contexto con parámetros de decodificación fijos para comparaciones justas.

Diagrama (conceptual): Los flujos de datos van del tokenizador → embeddings calibrados FP8 → bloques de atención (núcleos fusionados FlashAttention-2) → bloques FFN (matrices dispersas 2:4 vía cuSPARSELt) → cabecera LM. Los adaptadores de recuperación se sitúan en caminos lineales podados. Los perfiles de forma estática aseguran que los núcleos permanezcan en caminos optimizados.

Notas de rendimiento práctico en Hopper

• La elegibilidad es binaria: solo las matrices 2:4 correctamente enmascaradas obtienen la elevación de núcleos de 2×. Un único eje mal agrupado o un diseño no soportado te baja a núcleos densos.
• Las mayores ganancias de FP8 provienen del ahorro de ancho de banda y de mantener matrices en caminos rápidos; mantén los valores atípicos de escala bajo control con calibración actualizada después de la poda.
• La atención fusionada/optimizada (por ejemplo, FlashAttention-2) reduce la sobrecarga de atención, amplificando la aceleración realizada al dispersar los MLP.

Tablas Comparativas

Recetas de precisión y esparcidad en Hopper (TensorRT-LLM)

Receta	Esperado aumento de velocidad e2e (decodificación)	Impacto en Memoria/BW	Perfil de riesgo	Notas
Denso FP16	1.0×	Alto	Bajo	Línea base; más fácil de validar
Denso FP8	1.2–1.4×	BW de activación ↓	Bajo–moderado	Requiere escalado cuidadoso; buen primer paso
2:4 + FP16	1.3–1.8×	BW de peso ↓	Moderado	Asegúrate de la elegibilidad/orientación de la capa
2:4 + FP8	1.5–2.0×	BW de peso + activación ↓	Moderado	Punto ideal en H100/H200
INT8 W8A8	1.2–1.6×	BW de peso/act ↓	Moderado	Ampliamente soportado; recalibrar tras la poda

Qué podar (LLMs densos)

Módulo	Elegible para 2:4	Riesgo para la calidad	Guía
Atención Q/K/V	Sí	Medio	Podar de forma conservadora en capas tardías; preservar cabezas críticas de KV
Proyección de salida de atención	Sí	Bajo–moderado	Generalmente seguro; validar en largo contexto
FFN arriba/abajo (puerta)	Sí	Bajo	Principal objetivo para aceleraciones 2:4
Embeddings	No	Alto	No podar
RMSNorm/LayerNorm	No	Alto	No podar
Cabecera LM	Generalmente evitar	Alto	Opcional solo con validación sólida

Pros y contras (2:4 vs alternativas en NVIDIA)

Enfoque	Pros	Contras
2:4 N:M	2× rendimiento matemático a nivel de kernel; listo para producción en TensorRT-LLM	Restricciones de patrón; diseño/orientación estrictos
Solo FP8	Fácil de adoptar; portable a través de capas	Ganancias menores; aún limitado por ancho de banda en lugares
Esparcidad no estructurada	Alta capacidad de compresión	Poco/no aumento de velocidad sin núcleos especializados
Esparcidad estructurada por bloques	Buena localidad; núcleos más fáciles que no estructurados	Requiere núcleos personalizados afinados o cobertura específica

Mejores Prácticas

Curación del conjunto de calibración e higiene del escalado

• Mezcla instrucciones, razonamiento (matemáticas), código e indicaciones de largo contexto. 2–5k indicaciones son suficientes para un escalado estable y puntuaciones de poda.
• Fija los parámetros de decodificación (temperatura/top‑p/máx‑nuevos‑tokens) tanto para calibración como para evaluación para evitar confusiones.
• Para FP8, usa el escalamiento dinámico por tensor del motor Transformer y vuelve a ejecutar la calibración después de la poda para tener en cuenta los cambios de distribución.

Poda estructurada segura

• Podar solo capas con soporte probado 2:4 en tu motor. Sigue la agrupación de cuSPARSELt a lo largo de la dimensión K interna y mantén los pesos en diseños soportados.
• Incrementa la esparcidad en etapas (por ejemplo, 20% → 30% → 40%) con evaluaciones rápidas entre pasos. Detente cuando MMLU/GSM8K cambie alrededor de 1–2 puntos.
• Trata las capas de atención tardías como críticas para KV: poda menos agresivamente allí y preserva las cabezas conocidas por llevar señales de largo alcance.

Recuperación basada en adaptadores

• Empieza con rango LoRA 8–16 en caminos de atención y FFN; aumenta el rango solo si la validación no se recupera dentro de 1 punto en tus métricas.
• AdaLoRA puede asignar rango dinámicamente entre módulos; útil cuando los presupuestos de poda difieren por la profundidad de la capa.
• Entrena con tu mezcla de calibración y un peso ligero en código/matemáticas para estabilizar GSM8K/HumanEval.

Integración en el motor y disciplina en formas

• Construye motores con formas estáticas que cubran cargas reales (cubos de indicación/respuesta). Si caes fuera de una forma perfilada en tiempo de ejecución, la esparcidad puede deshabilitarse.
• Habilita atención fusionada (estilo FlashAttention-2) para que aparezca el cuello de botella del MLP y maximizar las ganancias de extremo a extremo.
• Verifica el compromiso de esparcidad inspeccionando trazas de kernel y registros de TensorRT-LLM; un pico repentino en p99 a menudo señala una vuelta a caminos densos.

Matriz de validación y salvaguardas

• Informa: tokens/s, latencia p50/p90/p99, memoria pico/activación, energía por token (vía registro de consumo de energía). Compara igualdad.
• Suite de tareas: MMLU, GSM8K, HumanEval, MT-Bench/BBH, además de al menos un benchmark de largo contexto.
• Pruebas canarias: indicaciones cortas orientadas a seguridad, seguimiento de instrucciones e integridad de largo contexto.
• Plan de reversión: artefactos versionados por etapa (línea de base → podado → recuperado → recalibrado). Escalación progresiva de tráfico con reversión automática en presupuestos de error o regresiones p99. ✅

Soluciones rápidas

• Desbordamiento/NaNs de FP8: ajusta clip‑max o vuelve a recoger escalas con indicaciones con valores atípicos.
• Aceleración <1.2×: verifica que las máscaras 2:4 se alineen al eje correcto y que todas las formas de producción estén perfiladas con esparcidad habilitada.
• Regresiones de atención: revierte la poda de Q/K en capas tardías o aumenta el rango LoRA en esos módulos.
• Picos en p99: asegura que el programador no exceda las longitudes máximas perfiladas y que el loteo no introduzca desviación de forma.

Ejemplos Prácticos

1) Aplicación de máscara 2:4 (PyTorch, ilustrativo)

import torch

@torch.no_grad()
def apply_2_4_mask(weight: torch.Tensor, group_dim: int = -1):
 # Remodelar para que los grupos de 4 estén en la dimensión más interna
 assert weight.shape[group_dim] % 4 == 0
 w = weight.transpose(group_dim, -1)
 g = w.reshape(*w.shape[:-1], w.shape[-1] // 4, 4)
 # Puntuación de magnitud dentro de cada grupo de 4
 scores = g.abs()
 # Mantén los dos mejores por grupo
 top2 = scores.topk(k=2, dim=-1).indices
 mask = torch.zeros_like(g, dtype=torch.bool)
 mask.scatter_(-1, top2, True)
 w_pruned = (g * mask).reshape_as(w)
 return w_pruned.transpose(group_dim, -1)

# Ejemplo: aplicar a una proyección hacia arriba de MLP
for name, mod in model.named_modules():
 if isinstance(mod, torch.nn.Linear) and is_eligible(name):
 mod.weight.copy_(apply_2_4_mask(mod.weight, group_dim=1)) # agrupar a lo largo de K

Nota: La agrupación/diseño exacto debe coincidir con las expectativas de cuSPARSELt para el compromiso de los núcleos de tensor dispersos.

2) Calibración FP8 con el motor Transformer (simplificado)

import transformer_engine.pytorch as te

class FP8Block(te.fp8_autocast):
 def __init__(self, enabled=True):
 super().__init__(enabled=enabled)

# Pase de calibración
with te.fp8_autocast(enabled=True):
 for batch in calib_loader:
 _ = model(**batch)
# Guarda las escalas recogidas (el framework gestiona estadísticas por tensor)

Referencia: El motor Transformer de NVIDIA proporciona casting FP8, recetas de escalado y guías de integración.

3) Recuperación LoRA (pseudocódigo estilo PEFT)

from peft import LoraConfig, get_peft_model

lora = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj","k_proj","v_proj","o_proj","up_proj","down_proj"], lora_dropout=0.05)
model = get_peft_model(model, lora)
# Entrena durante 3–10k pasos en instrucción/razonamiento/código mixto; detención temprana en segmento MMLU/GSM8K

Antecedentes: LoRA/AdaLoRA recuperan calidad con bajo cómputo tras la poda.

4) Construcción TensorRT-LLM (ejemplo de extracto de configuración)

{
 "precision": "fp8",
 "enable_sparse_weights": true,
 "fused_attention": "flash_v2",
 "profiles": [
 {"prompt": [1, 2048], "response": [1, 512]},
 {"prompt": [1, 4096], "response": [1, 1024]}
 ],
 "plugins": {"kv_cache": {"static": true}}
}

Construye y ejecuta (CLI varía según la versión; consulta la documentación de TensorRT-LLM):

trtllm-build --model./pruned_lora_recovered --config./trt_config.json --output./engine
trtllm-run --engine./engine --dataset./eval.jsonl --metrics tokens_per_s,latency_p50,latency_p99

Documentación: Repositorio y documentos de TensorRT-LLM para habilitar FP8, esparcidad y núcleos fusionados.

5) Registro de consumo de energía y latencia

# Energía
nvidia-smi dmon -s pucmt -i 0 -o DT >> power.log &
# Distribución de latencia a través de tu ejecutor
trtllm-run... --metrics latency_histogram

6) Alternativa INT8 (solo peso)

Si tu pila prefiere INT8, establece una línea base GPTQ/LLM.int8(), poda, luego recuantiza/recalibra antes de la construcción del motor.

Conclusión

La esparcidad 2:4 más FP8 en Hopper ya no es un truco de laboratorio—es una receta desplegable que consistentemente ofrece un rendimiento de 1.5–2.0× y ahorros materiales de energía al ser ejecutada con disciplina en TensorRT-LLM. El camino crítico es operacional: establece tu línea base, poda solo capas elegibles con la orientación correcta, recalibra la precisión tras los cambios estructurales, recupera con adaptadores ligeros y mantiene los motores en núcleos optimizados con perfiles de forma estática. Una matriz de validación rigurosa, pruebas canarias y salvaguardas de retroceso convierten esos beneficios en algo en lo que puedes confiar a escala.

Puntos clave

• Establece una línea base densa FP16/FP8 en TensorRT-LLM con atención fusionada antes de cualquier poda.
• Aplica 2:4 solo a matrices lineales/FFN elegibles, respetando la agrupación de cuSPARSELt; poda de forma conservadora en capas de atención tardías.
• Recalibra FP8 (o INT8) después de la poda y realiza una breve recuperación con LoRA/AdaLoRA para mantener las deltas métricas dentro de ~1–2 puntos.
• Valida tokens/s, latencia p50/p99, memoria y energía; usa MMLU, GSM8K, HumanEval, MT-Bench/BBH y largo contexto para detectar regresiones.
• Protege con canarios y retrocesos versionados; la disciplina en las formas es esencial para mantener los núcleos dispersos comprometidos.

Próximos pasos

• Prototipa la canalización en un modelo de tamaño medio (por ejemplo, 7–13B) con un pequeño conjunto de calibración para validar las herramientas.
• Muévete a tu escala objetivo, establece esparcidad en incrementos y codifica umbrales de aprobación/rechazo para la calidad y la latencia.
• Automatiza construcciones de motor para cada cubo de perfil e integra telemetría de energía/latencia en tu canalización de despliegue.

Mirando hacia adelante: a medida que la cobertura de TensorRT-LLM se amplía y las herramientas FP8 maduran, se esperan caminos más fáciles y automatizados para mezclar esparcidad 2:4 con escalado de precisión, y más de tu flota operando cómodamente en el carril rápido.

Fuentes y Referencias