Destilación de LLMs: Multiverse Computing reduce memoria 15× con 1 GPU

La destilación de conocimiento se vuelve asequible: cómo reducir 15× el uso de memoria en LLMs

Multiverse Computing ha publicado una investigación que reduce el uso de memoria en la destilación de modelos de lenguaje grandes (LLMs) en hasta 15.6×, permitiendo entrenar modelos comprimidos en un solo GPU H200 en lugar de cuatro nodos. El paper «Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss» introduce dos cambios sistémicos que transforman la destilación de conocimiento de una operación que requiere cientos de GPUs a un proceso accesible para equipos más pequeños.

La destilación de conocimiento, técnica que entrena un modelo estudiantil más pequeño para igualar el rendimiento de un maestro más grande, se ha vuelto práctica estándar para comprimir LLMs masivos como Kimi-K3 (2.8 billones de parámetros, 3TB de VRAM solo para cargar). Sin embargo, el paso de destilación era tradicionalmente la parte más cara del pipeline: mantener tanto el maestro como el estudiantil cargados y producir distribuciones de probabilidad sobre todo el vocabulario para cada token requería cantidades enormes de memoria VRAM.

El problema de memoria que paralizaba la experimentación

En el setup estándar de destilación online, un modelo como gpt-oss-120b (vocabulario de 201,088 tokens) a una longitud de secuencia de 32K y batch size 4 produce un tensor de probabilidades del maestro con forma 4 × 201,088 × 32,768. En bfloat16, eso ya son 50GB de VRAM para un solo tensor. Sumando gradientes, activaciones, pesos del modelo y estados del optimizador, una sola iteración de entrenamiento podía alcanzar 250GB de VRAM, más de lo que incluso un GPU H200 o B200 puede proporcionar.

🤖 La IA no es solo para leer sobre ella

En la comunidad la aplicamos: automatización, agentes IA y herramientas reales para emprender, no solo para informarte.

👥 Aplicarla en la comunidad

Esta barrera de memoria hacía que la experimentación a gran escala fuera prohibitiva para la mayoría de los equipos, limitando la innovación en compresión de LLMs a organizaciones con recursos de supercomputación.

Dos innovaciones que cambian las reglas del juego

Multiverse Computing abordó este problema con dos cambios sistémicos fundamentales:

1. Destilación offline: En lugar de recomputar el maestro en cada paso de entrenamiento, el equipo computa su salida una vez, cachea los top-100 tokens más probables por posición, y entrena el estudiantil contra ese cache. El maestro nunca tiene que estar en memoria durante el entrenamiento y no necesita ejecutarse nuevamente una vez que el cache existe.

2. Pérdida KL fusionada y chunked: La implementación tradicional de la pérdida de divergencia Kullback-Leibler construía una grilla completa vocabulario × secuencia antes de producir un solo número. La versión fusionada procesa un chunk de la secuencia a la vez de principio a fin, proyectando estados ocultos a logits para ese chunk, incorporando el resultado a la pérdida acumulada, y descartando el chunk antes de moverse al siguiente.

Resultados concretos: de 4 nodos a 1 GPU

Los números hablan por sí mismos. En un benchmark con Llama 3.1 8B Instruct como maestro y un modelo estudiantil de 3.2B parámetros a contexto de 8K tokens:

Método (contexto 8K, single H200)	Memoria máxima	Tiempo por iteración	Throughput
Destilación online	102.8 GB	25.9 s	237 TFLOP/s
Offline, KL densa	78.3 GB	18.5 s	331 TFLOP/s
Offline, KL forward-chunked	61.8 GB	18.4 s	335 TFLOP/s
Offline, KL fused chunked	58.3 GB	20.2 s	304 TFLOP/s

La reducción es aún más dramática a longitudes de contexto mayores. En un benchmark aislado a 32K tokens, la memoria máxima cae de 85.2 GiB con la pérdida densa a 5.45 GiB con la versión completamente chunked: una reducción de 15.6×.

Para destilar un modelo GPT-OSS 20B a contexto de 32,768 tokens, la memoria liberada por la pérdida fusionada permitió reducir el setup de cuatro nodos GPU a uno solo. El tiempo por paso cayó de 57.0 a 12.23 segundos (5× más rápido), y el throughput por GPU aumentó de 74.2 a 345.7 TFLOP/s.

Qué significa esto para tu startup de IA

1. Experimentación a escala sin supercomputación: Si tu equipo trabaja en compresión de modelos, fine-tuning especializado o desarrollo de LLMs eficientes, ahora puedes ejecutar campañas de destilación completas con hardware mucho más accesible. En lugar de necesitar clusters de cientos de GPUs, puedes comenzar con un solo H200 para prototipado y escalar gradualmente.

2. Costos de entrenamiento reducidos drásticamente: La destilación offline elimina la necesidad de mantener el maestro en memoria durante el entrenamiento, reduciendo tanto los requisitos de hardware como los costos asociados. Para startups que operan con presupuestos ajustados, esta eficiencia puede ser la diferencia entre poder experimentar o no.

3. Acceso a técnicas de vanguardia: La implementación open-source disponible en github.com/CompactifAI/Full-Chunked-KL-Loss significa que no necesitas desarrollar estas optimizaciones desde cero. Puedes integrar directamente las mejoras de memoria en tu pipeline existente.

Implementación práctica: cómo comenzar

Para equipos que ya trabajan con destilación:

Reemplaza tu implementación actual de pérdida KL con la versión chunked del repositorio open-source
Implementa caching offline de los logits del maestro antes del entrenamiento
Ajusta tu pipeline para procesar secuencias en chunks en lugar de enteras

Para equipos nuevos en destilación:

Comienza con modelos más pequeños (3-8B parámetros) para validar el pipeline
Utiliza el caching offline desde el inicio para maximizar eficiencia
Experimenta con diferentes tamaños de chunk para optimizar memoria vs velocidad

El resultado final: calidad preservada con menos recursos

Lo más impresionante es que estas optimizaciones no comprometen calidad. El estudiantil comprimido de 3.2B parámetros destilado desde Llama 3.1 8B Instruct retiene la mayoría de la precisión del maestro en benchmarks como BoolQ y HellaSwag, manteniéndose dentro de aproximadamente nueve puntos en MMLU, con menos de la mitad del conteo de parámetros.

Esta investigación es parte del esfuerzo continuo de Multiverse Computing para hacer la destilación y «healing» práctica para ejecutar a escala, no solo como una receta única, sino como algo en lo que los equipos pueden iterar económicamente.