Cómo abaratar la destilación de conocimiento para ejecutarla a gran escala

La destilación de conocimiento, donde un modelo estudiante más pequeño aprende a igualar a un modelo profesor más grande, es una técnica estándar para comprimir grandes modelos de lenguaje. Los recientes modelos de lenguaje abiertos de gran escala como Kimi-K3, con 2.8 billones de parámetros, requieren alrededor de 3TB de memoria VRAM, lo que hace que la compresión sea esencial. El artículo 'Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss' aborda los altos costos de la destilación al almacenar en caché los logits top-K del profesor fuera de línea, eliminando la necesidad de mantener al profesor en memoria. También introduce una pérdida KL por trozos fusionada que calcula la pérdida en trozos, evitando la matriz completa de vocabulario por secuencia. Esto reduce el uso máximo de VRAM de 250GB a 128GB en una sola H200, lo que permite la destilación en un contexto de 32K en una sola GPU en lugar de cuatro nodos, con tiempos de paso 5 veces más rápidos. El modelo estudiante, destilado de Llama 3.1 8B Instruct a 3.2B parámetros, conserva la mayor parte de la precisión en puntos de referencia como BoolQ y HellaSwag, y se mantiene dentro de nueve puntos en MMLU.