DisaggKD traslada al profesor a nodos independientes y acelera la destilación de Qwen3.5-35B-A3B a 9B

El equipo detrás de DisaggKD propone eliminar una restricción habitual de la destilación de conocimiento: que el profesor y el alumno tengan que convivir en las mismas GPU. En el par Qwen3.5-35B-A3B (profesor) y 9B (estudiante), el esquema desagregado obtuvo el mismo LM loss que una ejecución colocada con la misma función de pérdida, y además terminó antes. En pares más grandes, las ejecuciones colocadas se quedan sin memoria, mientras que las desagregadas siguen entrenando.

Por qué la colocación choca con el límite de memoria

En la mayoría de los frameworks de KD, el profesor y el estudiante comparten los mismos aceleradores. En cada tarjeta hay que alojar los pesos del profesor, los pesos del estudiante, los gradientes y los estados del optimizador. Cuanto más potente es el profesor, menos espacio queda para todo lo demás, y el tamaño del profesor queda limitado por la memoria de una sola GPU.

Hay un segundo problema, oculto en la propia función de pérdida. El KD clásico compara las distribuciones del profesor y del estudiante sobre todo el vocabulario. Para ello hacen falta logits de tamaño «tokens × vocabulario», y los modelos actuales tienen vocabularios grandes. Esa matriz se convierte en uno de los mayores consumidores de memoria en el pico.

Cómo funciona DisaggKD

Los autores dividen el esquema en tres partes.

Primera parte: separación por nodos. Las réplicas del profesor trabajan por separado y pasan al estudiante los estados ocultos (hidden states). Los logits no viajan por la red.

Segunda parte: fused loss. La KD-loss se calcula por ventanas de filas y shards del vocabulario. Con aritmética exacta, el resultado es igual a la pérdida sobre el vocabulario completo, pero la matriz de logits completos nunca se forma. La equivalencia se declara precisamente para aritmética exacta. En la práctica, el orden de suma en float puede generar pequeñas discrepancias, por lo que los autores comparan por separado las curvas de entrenamiento reales.

Tercera parte: prefetch ring. Es un búfer circular que sincroniza el ritmo del productor (el profesor) con el del consumidor (el trainer). Mientras el estudiante calcula un paso, los siguientes estados ocultos ya están en camino, y ninguna de las dos partes queda ociosa.

Además, los autores aportan modelos cerrados de memoria y de tráfico de red. Con ellos se puede estimar de antemano dónde está el límite de memoria para cada configuración (colocada o desagregada).

Qué mostraron las mediciones

• Par Qwen3.5-35B-A3B → 9B: el esquema desagregado alcanza el mismo LM loss que el colocado con la misma función de pérdida y es más rápido que él.
• El fused loss es útil por sí solo. Una ejecución colocada con él puede hacerse en menos nodos y con mayor rapidez.
• Si el estudiante ocupa varios nodos, el esquema desagregado ahorra más memoria.
• En pares más grandes, las ejecuciones colocadas fallan por OOM incluso con fused loss: los pesos del profesor y los estados del estudiante consumen todo el presupuesto. El esquema desagregado sí entrena en esos pares.

El resumen del artículo no incluye cifras absolutas de aceleración, número de GPU ni tamaños de los pares grandes. La magnitud de la mejora habrá que consultarla en el propio paper y contrastarla con el hardware de cada equipo.

A quién le puede servir

Las conclusiones de los autores plantean una bifurcación clara. Si el profesor es pequeño y el estudiante cabe en uno o dos nodos, basta con el esquema colocado con fused loss: ya ahorra memoria y acelera el entrenamiento. La configuración desagregada merece consideración cuando el profesor tiene decenas de miles de millones de parámetros (como 35B-A3B, con expertos) o cuando el estudiante se distribuye entre varios nodos. En ese caso, el equilibrio de memoria decide si la ejecución es siquiera posible.

También hay un coste. El esquema desagregado exige una conexión rápida entre nodos: los estados ocultos deben transmitirse sin pausas. Por eso resultan prácticos los modelos de tráfico del artículo: antes de alquilar un clúster, se pueden incorporar al cálculo para comprobar si la red se convertirá en un cuello de botella.

La idea del fused loss merece mención aparte. Se puede trasladar a pipelines de KD ya existentes sin cambiar la topología del clúster y eliminar de inmediato el pico de memoria de los logits. Para los equipos que destilan grandes modelos abiertos en otros compactos, es la forma más barata de probar el enfoque con sus propios datos.