Un Nuevo Truco Reduce Drásticamente el Costo de Comprimir Enormes Modelos de IA

Puntos clave

• La destilación de conocimiento, copiar el comportamiento de un gran modelo de IA en uno más pequeño, típicamente requiere cientos de chips especializados y coordinación cuidadosa.
• Un nuevo artículo almacena en caché las salidas del modelo grande una sola vez por adelantado, por lo que nunca tiene que ejecutarse de nuevo durante el entrenamiento, reduciendo drásticamente la memoria máxima.
• Una segunda técnica, llamada pérdida KL fusionada en fragmentos, procesa datos de entrenamiento en pequeños fragmentos en lugar de todos a la vez, manteniendo la memoria máxima casi plana incluso con entradas muy largas.
• Con 32K 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 para 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 parámetros como los 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 de veinte veces el espacio de tarjetas gráficas de consumo de alta gama.

Porque ejecutar modelos de ese tamaño es tan costoso, los investigadores routinariamente los comprimen en versiones más pequeñas. El método estándar es la destilación de conocimiento: ejecutar el modelo "maestro" grande en un montón de texto, registrar cómo respondería, luego entrenar un modelo "estudiante" más pequeño para copiar esas respuestas lo más fielmente posible. Más económico de ejecutar, casi tan capaz.

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

El gasto tiene dos fuentes. Mantener tanto el maestro como el estudiante cargados en memoria al mismo tiempo es la primera. La segunda es que comparar sus salidas requiere construir una malla enorme: una fila para cada palabra en el vocabulario del modelo, una columna para cada posición en la entrada. Multiplique eso a lo largo de un texto largo y llenará cada chip que tenga.

Un ejemplo concreto: gpt-oss-120b tiene un vocabulario de 201.088 palabras. Con una longitud de secuencia de 32K tokens y un tamaño de lote de cuatro ejemplos a la vez, el tensor de salida del maestro solo ocupa aproximadamente 50 gigabytes. Agregue el modelo estudiante, gradientes, estados del optimizador y todo lo demás que el proceso de entrenamiento necesita, y un solo paso de entrenamiento puede aumentar a aproximadamente 250 gigabytes, más de lo que puede contener un H200, uno de los chips de IA más poderosos disponibles.

¿Qué cambiaron los investigadores?

Dos cosas, funcionando juntas. Ambas aparecen en un artículo de CompactifAI, compartido por primera vez en Hugging Face.

La primera es la destilación sin conexión. En lugar de ejecutar de nuevo el maestro en cada paso de entrenamiento, los investigadores lo ejecutan una vez y guardan solo sus 100 predicciones de palabras más probables por posición. Ese caché va al disco; el maestro se apaga completamente. Porque el caché es reutilizable en muchos experimentos, la costosa ejecución del maestro se convierte en un costo único.

La segunda es la pérdida KL fusionada en fragmentos. La divergencia KL es una medida matemática de cuán diferentes son dos distribuciones de probabilidad, aquí significando cuánto se desvían las predicciones de palabras del estudiante de las del maestro. El cálculo estándar construye esa malla completa de vocabulario por secuencia de una sola vez y mantiene todo. El nuevo enfoque funciona un fragmento de la secuencia a la vez, incorpora cada fragmento en un total acumulativo, luego lo descarta antes de pasar al siguiente. El paso hacia atrás recomputa lo que necesita 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 con fragmentación 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. La pérdida de calidad al usar predicciones en caché de los 100 principales en lugar de la distribución completa del maestro es insignificante.

¿Importa esto para documentos más largos?

Sí, y la brecha crece rápidamente. La ventaja de memoria del enfoque fusionado fragmentado es modesta en 8K tokens pero dramática más allá de eso. Con 32K tokens, la memoria máxima cae de 85.2 gigabytes con el enfoque estándar a 5.45 gigabytes con el nuevo. Los experimentos que anteriormente necesitaban un rack de chips especializados ahora pueden ejecutarse en una sola tarjeta de alta gama. Esa es una apertura realista para equipos de investigación más pequeños y empresas que quieran construir modelos comprimidos a partir de los sistemas de IA de código abierto más grandes disponibles hoy. Cubrimos cómo el marco Orchard de Microsoft persigue un objetivo similar, ajustando agentes capaces en modelos más pequeños, en nuestro artículo del 5 de agosto.

La advertencia honesta: el punto de referencia aquí utiliza Llama 3.1 8B como estudiante. Qué tan bien se mantiene la técnica con tamaños de estudiante más grandes es una pregunta que el artículo aún no resuelve completamente.

Preguntas comunes

¿Termina el modelo más pequeño peor que el original?

En los puntos de referencia reportados, la pérdida de entrenamiento 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 fragmentado es de código abierto, publicado públicamente en GitHub por CompactifAI. Cualquiera con un chip de IA moderno único y una copia de un modelo grande de código abierto puede, en principio, ejecutar sus propios experimentos de destilación.

¿Quién se beneficia primero?

Los pequeños equipos de investigación y empresas que desean comprimir modelos grandes de código abierto para un despliegue más económico sentirán la diferencia más pronto. Los usuarios finales pueden beneficiarse más adelante cuando modelos comprimidos más asequibles lleguen a productos y servicios.