DisaggKD выносит учителя на отдельные узлы и ускоряет дистилляцию Qwen3.5-35B-A3B в 9B

Команда, стоящая за DisaggKD, предлагает убрать из дистилляции знаний привычное ограничение: учитель и ученик больше не обязаны жить на одних и тех же GPU. На паре Qwen3.5-35B-A3B (учитель) и 9B (студент) раздельная схема показала тот же LM loss, что и совмещённый запуск с тем же лоссом, и при этом отработала быстрее. На более крупных парах совмещённые запуски упираются в нехватку памяти, а раздельные продолжают обучаться.

Почему совмещение упирается в память

В большинстве фреймворков для KD учитель и студент делят одни и те же ускорители. На каждой карте приходится держать веса учителя, веса студента, градиенты и состояния оптимизатора. Чем сильнее учитель, тем меньше остаётся места под всё остальное, и размер учителя ограничивает память отдельной GPU.

Вторая проблема скрыта в самом лоссе. Классический KD сравнивает распределения учителя и студента по всему словарю. Для этого нужны логиты размером «токены × словарь», а у современных моделей словари большие. Эта матрица становится одним из главных потребителей памяти на пике.

Как устроен DisaggKD

Авторы разбирают схему на три части.

Первая часть: разделение по узлам. Реплики учителя работают отдельно и передают студенту скрытые состояния (hidden states). Логиты по сети не ходят.

Вторая часть: fused loss. KD-loss считается по окнам строк и шардам словаря. В точной арифметике результат равен лоссу по полному словарю, но сама матрица полных логитов не формируется. Эквивалентность заявлена именно для точной арифметики. На практике порядок суммирования в float может давать мелкие расхождения, поэтому реальные кривые обучения авторы сравнивают отдельно.

Третья часть: prefetch ring. Это кольцевой буфер, который согласует темп производителя (учителя) и потребителя (тренера). Пока студент считает шаг, следующие скрытые состояния уже летят к нему, и ни одна из сторон не простаивает.

Кроме того, авторы приводят замкнутые модели памяти и сетевого трафика. По ним можно заранее прикинуть, где для каждой раскладки (совмещённой или раздельной) проходит граница по памяти.

Что показали замеры

• Пара Qwen3.5-35B-A3B → 9B: раздельная схема достигает того же LM loss, что и совмещённая с тем же лоссом, и работает быстрее неё.
• Fused loss полезен и сам по себе. Совмещённый запуск с ним можно вести на меньшем числе узлов и быстрее.
• Если студент занимает несколько узлов, раздельная схема экономнее по памяти.
• На более крупных парах совмещённые запуски даже с fused loss падают с OOM: веса учителя и состояния студента съедают весь бюджет. Раздельная схема на этих парах обучается.

В аннотации нет абсолютных цифр ускорения, числа GPU и размеров крупных пар. Величину выигрыша придётся смотреть в самой статье и сверять с собственным железом.

Кому это пригодится

Выводы авторов дают понятную развилку. Если учитель маленький, а студент умещается на одном-двух узлах, хватит совмещённой схемы с fused loss: она уже экономит память и ускоряет обучение. Раздельную раскладку стоит рассматривать, когда учитель в десятки миллиардов параметров (как 35B-A3B с экспертами) или студент распределён по нескольким узлам. Тогда баланс памяти решает, будет ли запуск вообще.

Есть и цена. Раздельная схема требует быстрой связи между узлами: скрытые состояния нужно стримить без пауз. Поэтому модели трафика из статьи практичны: перед арендой кластера их можно подставить в расчёт и проверить, не станет ли сеть узким местом.

Отдельно полезна идея fused loss. Её можно перенести в уже существующие KD-пайплайны, не меняя топологию кластера, и сразу снять пик памяти на логитах. Для команд, которые дистиллируют крупные открытые модели в компактные, это самый дешёвый способ проверить подход на своих данных.