Дистилляция LLM на одной GPU: offline-кэш и chunked KL loss

Дистилляция большой LLM в компактную модель обычно требует сотен гигабайт VRAM: при контексте 32K и словаре 201 088 токенов у gpt-oss-120b один тензор вероятностей учителя — уже около 50GB, а пик итерации доходит примерно до 250GB. Multiverse Computing показала, как уложить тот же прогон на одну H200 — offline-кэш top-100 логитов и fused chunked KL loss без полной матрицы «словарь × длина контекста».

Где утекает память

Классическая online-дистилляция держит учителя и студента в памяти и на каждом шаге пересчитывает полное распределение по словарю — хотя поведение teacher в рамках одного run не меняется. Для gpt-oss-120b при batch 4 и sequence 32K dense KL в реализациях вроде PyTorch даёт пик выше 141GB H200 — обучение просто не стартует.

Два приёма

• Offline-кэш — учитель считается один раз, сохраняются 100 самых вероятных токенов на позицию; во время training teacher не сидит в VRAM и тот же кэш можно гонять в разных ablation.
• Fused chunked KL — loss считается по кускам контекста: hidden states проецируются в логиты только для chunk, результат входит в running loss, chunk выбрасывается. Backward пересчитывает chunk на лету — пик памяти растёт линейно с длиной контекста, а не как vocabulary × sequence.

Реализацию выложили в github.com/CompactifAI/Full-Chunked-KL-Loss.

Что на одной карте

На одной H200 с Llama 3.1 8B Instruct как teacher и студентом 3.2B при контексте 8K все четыре схемы — online, dense KL, forward-chunked и fused chunked — дают почти одинаковый training loss. Пик памяти: online 102.8 GB, fused chunked 58.3 GB.

На 32K в изолированном бенчмарке dense loss — 85.2 GiB, fused chunked — 5.45 GiB (в 15.6 раз меньше). Дистилляция GPT-OSS 20B при 32 768 токенов: с четырёх GPU-нод до одной, step time с 57.0 до 12.23 с, throughput с 74.2 до 345.7 TFLOP/s на GPU.

Студент 3.2B из Llama 3.1 8B Instruct удерживает большую часть accuracy на BoolQ и HellaSwag, на MMLU — в пределах ~9 пунктов от teacher при менее половине параметров.

Источник: Making Knowledge Distillation Cheap Enough to Run at Scale.