Дистилляция GPT-OSS 20B уместилась в одну GPU-ноду вместо четырёх: память на loss упала в 15,6 раза

Дистилляция GPT-OSS 20B на контексте 32 768 токенов теперь идёт на одной GPU-ноде вместо четырёх. Шаг обучения упал с 57,0 до 12,23 секунды, пропускная способность на GPU выросла с 74,2 до 345,7 TFLOP/s (блог Multiverse Computing на Hugging Face, 10.08). Пиковая память на подсчёте loss ужалась с 85,2 до 5,45 GiB.

Раньше платили словарём. Плотный KL-loss держит в памяти матрицу на весь словарь. У gpt-oss-120b это 201 088 токенов и около 250 GB VRAM на контексте 32K, у чанкованной версии - 128 GB. На 256K расход становится линейным по длине последовательности, а не по размеру словаря: 11,6 GiB против 134,2 GiB у следующей по эффективности реализации.

Второй кусок экономии даёт сам учитель. Авторы кэшируют его top-100 логитов один раз и учат студента против кэша: итерация быстрее примерно на 29%, пропускная способность на одной H200 выше до 41%. Учитель при этом уходит из памяти целиком, а training loss почти не отличается от online-дистилляции.

Студент проседает не везде. Модель на 3,2B, дистиллированная из Llama 3.1 8B Instruct, держит точность учителя на BoolQ и HellaSwag и отстаёт примерно на девять пунктов на MMLU. Даже на 8K экономия видна: online-дистилляция берёт 102,8 GB пиковой памяти и 25,9 секунды на итерацию, fused chunked KL - 58,3 GB и 20,2 секунды.

Взять можно сейчас. Код лежит в репозитории CompactifAI/Full-Chunked-KL-Loss на GitHub. Нужны Python 3.10+, PyTorch с CUDA и хотя бы одна NVIDIA GPU, бенчмарк трёх реализаций запускается командой `python benchmark.py`, проверка - `python test_losses.py`.

Готового плагина под TRL и HF Trainer нет, отдельного pip-пакета тоже. В своём стеке авторы подключали метод руками через Megatron-Bridge с бэкендом Megatron-LM и NVIDIA ModelOpt. Код открытый, на сам метод подана патентная заявка EP26382987.1.

Реализацию для стандартных библиотек авторы обещают выложить отдельно - пока метод вживляется в тренировочный стек руками.