Un nuevo truco reduce drásticamente el costo de comprimir grandes modelos de IA

Puntos clave

• La destilación de conocimiento, que copia el comportamiento de un modelo de IA grande en uno más pequeño, típicamente requiere cientos de chips especializados y una coordinación cuidadosa.
• Un nuevo artículo almacena en caché los resultados del modelo grande de una sola vez al inicio, por lo que nunca tiene que ejecutarse nuevamente durante el entrenamiento, reduciendo significativamente el uso de memoria.
• Una segunda técnica, llamada pérdida KL fusionada por fragmentos, procesa datos de entrenamiento en porciones pequeñas en lugar de todo a la vez, manteniendo la memoria máxima casi plana incluso con entradas muy largas.
• Con 32.000 tokens de contexto, la memoria máxima cae de 85,2 gigabytes a 5,45 gigabytes, una reducción de 15,6 veces.
• El código completo de la función de pérdida eficiente en memoria está disponible públicamente en GitHub.

Los modelos de IA más grandes del mundo son casi imposiblemente enormes. Kimi-K3, lanzado a principios de este año, tiene 2,8 billones de parámetros (piense en los parámetros como los millones de pequeños controles dentro de un modelo que se ajustan durante el entrenamiento) y necesita aproximadamente tres terabytes de memoria especializada solo para cargarse. Eso es más espacio del que ocupan veinte tarjetas gráficas de consumo de gama alta.

Debido a que ejecutar modelos de ese tamaño es tan costoso, los investigadores routinariamente los comprimen en versiones más pequeñas que cuestan menos de operar. El método estándar se llama destilación de conocimiento: ejecuta el modelo "maestro" grande en un conjunto de texto, registra cómo respondería, luego entrena un modelo "estudiante" más pequeño para copiar esas respuestas lo más fielmente posible. El estudiante termina siendo mucho más económico de ejecutar mientras mantiene la mayor parte de la calidad del maestro.

¿Por qué es tan difícil hacerlo de manera económica?

El gasto proviene de dos lugares. Primero, el enfoque estándar mantiene tanto el maestro como el estudiante cargados en memoria al mismo tiempo. Segundo, comparar sus resultados requiere construir una cuadrícula enorme: una fila para cada palabra en el vocabulario del modelo (a menudo más de 100.000 palabras), una columna para cada posición en la entrada. Multiplique eso en un texto largo y rápidamente llenará cada chip que tenga.

Como ejemplo concreto: un modelo abierto popular, gpt-oss-120b, tiene un vocabulario de 201.088 palabras. Con una longitud de secuencia de 32.000 tokens (aproximadamente el texto de una novela corta) y un tamaño de lote modesto de cuatro ejemplos a la vez, solo el tensor de salida del maestro ocupa unos 50 gigabytes. Agregue el modelo del estudiante, gradientes y todo lo demás que el proceso de entrenamiento necesita, y un solo paso de entrenamiento puede alcanzar aproximadamente 250 gigabytes, más de lo que puede contener un H200, uno de los chips de IA más potentes disponibles actualmente.

¿Qué cambiaron los investigadores?

Dos cosas, trabajando juntas. Ambas se describen en un artículo publicado por el equipo de CompactifAI, compartido por primera vez en Hugging Face.

El primer cambio es la destilación sin conexión. En lugar de ejecutar el maestro nuevamente en cada paso de entrenamiento, los investigadores lo ejecutan una sola vez y guardan solo sus 100 predicciones de palabras más probables por posición. Ese caché se escribe en el disco. Luego el maestro se apaga completamente, y el estudiante entrena contra el caché guardado únicamente. Debido a que el caché se puede reutilizar en muchos experimentos diferentes, la costosa ejecución del maestro se convierte en un costo único.

El segundo cambio es la pérdida KL fusionada por fragmentos. La divergencia KL es una medida matemática de cuán diferentes son dos distribuciones de probabilidad, en este caso, cuán diferentes son las predicciones de palabras del estudiante respecto a las del maestro. La forma estándar de calcularlo construye esa cuadrícula completa de vocabulario por secuencia de una sola vez y mantiene todo en memoria. El nuevo método divide la secuencia en fragmentos pequeños, calcula y descarta cada fragmento antes de pasar al siguiente, y recomputa lo que necesita durante el pase hacia atrás en lugar de almacenarlo.

Método	Memoria máxima (contexto 8K)	Tiempo de iteración
Destilación en línea	102,8 GB	25,9 s
Sin conexión, KL denso	78,3 GB	18,5 s
Sin conexión, KL fragmentado hacia adelante	61,8 GB	18,4 s
Sin conexión, KL fusionado fragmentado	58,3 GB	20,2 s

Los cuatro métodos producen resultados de entrenamiento casi idénticos, lo que significa que la pérdida de calidad al usar las predicciones almacenadas en caché de los 100 mejores en lugar de la distribución completa del maestro es insignificante.

¿Importa esto para documentos más largos?

Sí, y la diferencia crece rápidamente. La ventaja de memoria del enfoque fusionado fragmentado es modesta a 8.000 tokens pero dramática en longitudes más largas. A 32.000 tokens, la memoria máxima cae de 85,2 gigabytes con el enfoque estándar a 5,45 gigabytes con el nuevo. La pérdida densa falla completamente a 64.000 tokens. A 256.000 tokens, el método fusionado fragmentado usa 11,6 gigabytes contra 134,2 gigabytes de la siguiente mejor alternativa.

En términos prácticos: los experimentos que anteriormente requerían un bastidor de cientos de chips ahora pueden ejecutarse en una sola tarjeta de gama alta. Esto hace que sea realista para equipos de investigación más pequeños y empresas construir modelos comprimidos a partir de los sistemas de IA de código abierto más grandes disponibles actualmente.

Preguntas frecuentes

¿El modelo más pequeño termina siendo peor que el original?

En los puntos de referencia reportados, la pérdida de entrenamiento del modelo estudiante fue casi idéntica en los cuatro métodos, incluido el nuevo enfoque económico. La calidad se mantuvo a pesar de usar solo las 100 predicciones de palabras principales por posición en lugar de la salida completa del maestro.

¿Puede cualquiera usar esta técnica?

El código de pérdida fragmentada es de código abierto y está publicado públicamente en GitHub por CompactifAI. Cualquiera con un solo chip de IA moderno y una copia de un modelo de código abierto grande puede, en principio, ejecutar sus propios experimentos de destilación con él.

¿Quién se beneficia primero?

Los pequeños equipos de investigación en IA y las empresas que desean comprimir grandes modelos de código abierto para una implementación más económica sentirán la diferencia más rápidamente. Los usuarios finales pueden beneficiarse más tarde cuando modelos comprimidos más asequibles lleguen a productos y servicios.