Alibaba Qwen представила TRACE: Rollout в FP4 ускоряется до 5,4 раза

Команда Alibaba Qwen представила TRACE — метод для обучения больших моделей с подкреплением, который переводит фазу Rollout в FP4: и веса, и активации (W/A), и KV-кеш. По заявлению авторов, Rollout ускоряется до 5,4 раза по сравнению с BF16, а обучение остаётся стабильным. Последнее здесь важнее скорости, потому что гнать генерацию в четырёх битах умеют многие, а удержать при этом RL от разваливания получается далеко не у всех.

Почему Rollout стал самой прожорливой частью RL

В RL для языковых моделей цикл выглядит так: модель генерирует множество ответов, часто длинных цепочек рассуждений, их оценивают, и по этим траекториям обновляют веса. Генерация идёт токен за токеном, упирается в пропускную способность памяти, а KV-кеш растёт вместе с длиной ответа. Чем длиннее reasoning и больше сэмплов на промпт, тем больше времени и видеопамяти съедает именно Rollout, а не шаг градиентного спуска.

Отсюда логичное желание сжать этот этап. Переход с BF16 на FP4 теоретически режет и объём вычислений, и расход памяти. Меньше байт на вес и на элемент кеша означает больше параллельных запросов на одну карту и меньше ожидания данных.

Где FP4 ломает обучение

Цена такого сжатия в грубости числовой сетки. В четырёх битах помещается всего 16 значений, поэтому две очень близкие величины легко оказываются по разные стороны границы квантования и падают в разные кодовые точки. Если тренер и Rollout-движок считают в разной точности, то на почти одинаковых входах они получают заметно разные результаты. Крошечная разница на уровне числа превращается в расхождение между тем, что модель «сгенерировала», и тем, что она «думает» при обновлении весов.

Для RL это опасная вещь. Алгоритм предполагает, что траектории порождены политикой, близкой к обучаемой. Когда генератор и тренер расходятся, градиент начинает опираться на данные, которые текущая модель выдала бы иначе. Обычно это проявляется как шум, дрейф и нестабильность обучения. С FP4 проблема острее, чем с FP8, потому что граничных эффектов больше, а запас точности меньше. Именно эту проблему TRACE и пытается снять.

Что известно о TRACE

Ключевые параметры работы:

• Разработчик: команда Qwen (Alibaba).
• Область применения: этап Rollout в RL-обучении больших моделей.
• Формат: FP4 для весов и активаций (W/A) плюс FP4 для KV-кеша.
• Базовая линия: BF16.
• Заявленный эффект: ускорение Rollout до 5,4 раза.
• Цель: сохранить стабильность обучения при низкой точности генерации.

Формулировка «до 5,4 раза» означает максимум, а не средний результат. Реальный выигрыш зависит от длины генерации, размера батча, архитектуры модели и железа. Чем длиннее контекст и тяжелее KV-кеш, тем заметнее эффект от квантования кеша, потому что именно он часто становится узким местом на длинных рассуждениях.

Что проверять командам, которые гонят RL на своих кластерах

Для тех, кто сам обучает reasoning-модели, главный вывод в том, что Rollout стоит рассматривать как отдельную оптимизационную задачу, а не как побочный процесс. Если генерация занимает основную часть времени шага, то сжатие точности даёт больше, чем тюнинг оптимизатора.

Прежде чем переходить на FP4, нужно измерить расхождение между вероятностями токенов у тренера и у Rollout-движка. Если оно уже заметно на BF16 или FP8, то на четырёх битах оно вырастет. Отдельно стоит проверять качество на длинных траекториях: ошибка KV-кеша копится с длиной ответа.

Наконец, выигрыш упирается в аппаратную поддержку. Нативный FP4 есть в новых поколениях ускорителей, например в архитектуре Blackwell, а на более старых картах 4-битный путь не даст такого ускорения. Для лабораторий с ограниченным парком GPU работы вроде TRACE интересны тем, что позволяют получать больше RL-итераций с тем же железом. Следующий вопрос — насколько результат повторяется на моделях другого размера и в других задачах, от математики до кода.