Во время стажировки в команде GigaChat Pretrain мы исследовали диффузионные языковые модели (dLLM) — относительно новое направление в LLM. Наша цель состояла не в том, чтобы тратить много ресурсов на обучение с нуля, а в том, чтобы перевести базовую авторегрессионную модель GigaChat3-10B-A1.8B-base в диффузионный режим. Так появились экспериментальные GFusion-10B-A1.8B-base и GFusion-10B-A1.8B.
Главные результаты проекта:
* Ускорение генерации: в режиме одного пользователя GFusion в среднем на 70 % быстрее GigaChat3 и превосходит вариант с дополнительной MTP-головой на 39 %.
* Сохранение качества: GFusion достигает ускорения при просадке качества всего на 2–4 п. п. в сравнении с GigaChat3; этот компромисс можно контролировать через параметры генерации диффузии.
* Открытый код: мы поддержали GFusion в SGLang и добавили новый алгоритм семплирования, который ускоряет генерацию других диффузионных LLM.
* Ускорение обучения: оптимизировали attention для текстовой диффузии — это дало +60 % к средней скорости обучения GFusion по сравнению с Flex-Attention.
Большинство современных LLM являются авторегрессионными (AR): каждый следующий токен предсказывается по предыдущему контексту. Диффузионные модели ослабляют это ограничение — они работают не с одним следующим токеном, а с частично замаскированным блоком фиксированного размера. Модель итеративно восстанавливает маски, постепенно превращая блок в осмысленный текст. Токены внутри одного блока могут декодироваться не только слева направо, но и в произвольном порядке, по несколько за один проход.
Полный цикл обучения GFusion включал несколько этапов:
* базовый претрейн — адаптация AR-модели к диффузионной генерации;
* мидтрейн — прокачивание способностей и знаний модели;
* расширение контекста — увеличение контекстного окна до 32 тыс. токенов;
* SFT — обучение модели следовать инструкциям и формату;
* confidence tuning — короткая стадия после SFT для ускорения генерации.
В ходе экспериментов мы протестировали разные подходы: варьировали размер блока (от 16 до 64 токенов), пробовали разные методы генерации (предсказание следующего замаскированного токена или текущего замаскированного), использовали гибридную постановку (добавляли AR-регуляризацию к диффузионной функции потерь). В итоге выбрали стандартный подход с предсказанием текущего замаскированного токена, постепенный рост размера блока и стандартную диффузионную loss-функцию.
Для декодирования мы сравнили два алгоритма семплирования — threshold sampling и entropy-bounded sampling. Последний показал себя лучше: генерация стала быстрее, а метрики выросли. В дальнейшем мы везде используем EB-семплирование.
Также мы расширили контекст модели до окна в 32 тыс. токенов, дообучив её на 30 млрд токенов длинных примеров. Для интерполяции RoPE с 4 тыс. до 32 тыс. использовали YaRN.
На стадии SFT мы применили идею complementary masking: для каждого обучающего примера добавляли парный пример с обратной маской (видимые токены маскировались, и наоборот). Лучший результат мы получили после шести эпох обучения — это вдвое больше, чем обычно требуется для авторегрессии.
На стадии confidence tuning мы добавили к функции потерь компонент, который штрафует модель за высокую энтропию распределения для корректно предсказанных токенов. Это позволило увеличить TPF (Tokens-Per-Forward) и сократить число вычислений, необходимых для генерации ответа.
Итоговые варианты GFusion мы сравнивали с GigaChat3 и другими диффузионными языковыми моделями похожих размеров (LLaDA-MoE-7B, LLaDA2.0-mini preview). GFusion-10B-A1.8B в среднем даёт +70 % к скорости генерации по сравнению с GigaChat3-10B-A1.8B и обгоняет вариант с MTP на +39 %.
Для удобства использования GFusion мы добавили её поддержку в SGLang. Также реализовали недостающую логику для бэкенда Flash-Attention для работы MLA в диффузионном режиме с attention-головой размера 192. Кроме того, написали свою реализацию ядер на TileLang, которая учитывает структуру диффузионной маски. Это позволило получить ускорение в +41,7 % и +77,4 % end-to-end при обучении на контексте в 4 тыс. и 32 тыс. токенов соответственно. Все реализации мы выкладываем в открытый доступ.
Подписаться в Макс
