Учёные ускорили спекулятивное декодирование при RL-обучении моделей до 122 млрд параметров

Дообучение языковых моделей с подкреплением (RL post-training) тратит основную часть времени и вычислений на генерацию rollout, последовательностей, которые модель сама производит и по которым потом считается сигнал обучения. Спекулятивное декодирование ускоряет этот процесс: небольшая «черновая» модель предлагает несколько токенов вперёд, а основная модель проверяет их за один проход. Если черновую модель дополнительно совместно дообучать онлайн вместе с основной (online co-training), её предсказания становятся точнее и ускорение растёт ещё сильнее.

Авторы указывают, что перенести такое совместное дообучение на большие модели с длинным контекстом мешают два технических препятствия. Первое: механизм branch attention (внимание, которое одновременно учитывает и черновые, и основную ветки токенов) не поддерживается стандартными реализациями причинного контекстного параллелизма (context parallelism, CP), техники, которая разбивает обработку длинного контекста между несколькими устройствами. Второе: промежуточные признаки основной модели, нужные черновой для дообучения, оказываются раскиданы по разным стадиям конвейерного параллелизма (pipeline parallelism, PP), техники, которая делит саму модель на последовательные стадии по разным устройствам.

Для решения первой проблемы авторы расширили упакованное, балансированное по нагрузке зигзагообразное кольцевое внимание (zigzag ring attention): они объединили branch attention, вычисляемое локально на каждом ранге (вычислительном узле), с обычным причинным вниманием по основной последовательности. Для решения второй проблемы предложен отдельный механизм передачи данных TapChannel: он переносит промежуточные признаки основной модели между стадиями конвейера по отдельному каналу, не трогая при этом сам график выполнения конвейера.

По словам авторов, в экспериментах черновые модели, дообученные таким способом, точно следуют за поведением основной («политики») модели и при этом дают существенное ускорение, как на этапе генерации rollout, так и в целом по циклу обучения, на моделях вплоть до 122 млрд параметров. Предложенная схема контекстного параллелизма хорошо масштабируется при длине контекста 256 тысяч токенов и заметно экономит память по сравнению с прежними подходами, а накладные расходы на передачу данных через TapChannel в конвейерном параллелизме авторы называют небольшими. Конкретных цифр по величине ускорения, экономии памяти или накладных расходов в тексте не приведено, оценки даны только качественно («существенное», «значительное», «небольшие»). Код системы опубликован в репозитории проекта NeMo RL на GitHub.

Ключевые факты

Предложена система для совместного онлайн-дообучения черновой модели в спекулятивном декодировании во время RL-дообучения больших языковых моделей с длинным контекстом
Решены два препятствия: branch attention не поддерживалось стандартным контекстным параллелизмом (CP), а нужные черновой модели признаки были раскиданы по стадиям конвейерного параллелизма (PP)
Для CP расширили зигзагообразное кольцевое внимание, объединив branch attention с причинным вниманием по основной последовательности
Для PP создали механизм TapChannel, отдельный канал передачи промежуточных признаков между стадиями конвейера без изменения его расписания
Эксперименты показали ускорение rollout-генерации и всего цикла обучения на моделях вплоть до 122 млрд параметров, устойчивое масштабирование при контексте 256 тыс. токенов и экономию памяти; конкретные цифры ускорения и экономии в тексте не названы

Почему это важно

RL-дообучение больших языковых моделей, один из самых дорогих этапов их создания, и основную часть этой стоимости съедает генерация rollout. Работа показывает, как ускорить именно этот этап за счёт совместного дообучения черновой модели для спекулятивного декодирования, причём не в лабораторном масштабе, а на моделях вплоть до 122 млрд параметров и на контексте длиной 256 тысяч токенов, что требует решения нетривиальных проблем распределённых вычислений.

Кому это важно

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

Как это применить

Код системы выложен в репозитории проекта NeMo RL на GitHub, то есть подход можно изучить и опробовать напрямую, а не только по описанию в тексте. Применимость ограничена командами, у которых уже есть инфраструктура контекстного и конвейерного параллелизма для обучения моделей такого масштаба.

Можно ли доверять

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

Риски и подводные камни

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