TRACE: метод FP4-обучения MoE-нейросетей с подкреплением ускоряет rollout до 5,4 раза
При дообучении больших языковых моделей с подкреплением (RL) самая затратная часть, rollout, то есть генерация ответов моделью: она требует много вычислений и памяти. Поэтому для rollout пытаются использовать низкую точность чисел, в частности формат FP4.
Авторы статьи указывают на слабое место существующих FP4-методов для RL: они оптимизируют точность квантования на пути обучения и на пути rollout независимо друг от друга, а не сокращают расхождение между двумя квантованными путями выполнения.
В ответ предложен TRACE (Train-Rollout Quantization Alignment via Compact GuidancE), фреймворк FP4-квантования для RL-обучения языковых моделей архитектуры Mixture-of-Experts (MoE, «смесь экспертов»). Он состоит из двух частей. Первая, квантизационно-осведомлённое обучение (quantization-aware training), управляемое rollout: результаты квантования на стороне rollout используют, чтобы направлять решения об округлении в FP4 на стороне обучения. Так напрямую уменьшается расхождение между обучением и rollout. Вторая, схема кэширования информации о квантовании: она выборочно сохраняет мантиссу и масштаб (scale) только из более глубоких слоёв, чтобы снизить расходы на хранение и передачу данных, которые добавляет такое управление.
TRACE проверили на четырёх крупных MoE-моделях в задачах рассуждения, программирования и длинных (long-horizon) RL-задачах. По результатам авторов, метод позволяет вести rollout сразу с весами и активациями в FP4 и с KV-кэшем в FP4, а качество RL остаётся сопоставимым с rollout в BF16. Заявлено ускорение rollout до 5,4 раза и сильный итоговый результат в FP4 по сравнению с последующим (post-hoc) FP4-квантованием политик, обученных в BF16.
Ключевые факты
- TRACE, FP4-фреймворк для RL-обучения языковых моделей архитектуры Mixture-of-Experts.
- Идея: результаты квантования на стороне rollout направляют округление в FP4 на стороне обучения, что сокращает расхождение между двумя путями.
- Кэширование выборочно хранит мантиссу и масштаб из глубоких слоёв, чтобы снизить расходы на хранение и передачу данных.
- Проверка на четырёх крупных MoE-моделях в задачах рассуждения, программирования и длинных RL-задачах.
- Заявлено ускорение rollout до 5,4 раза при качестве RL, сопоставимом с BF16 rollout.
Почему это важно
Генерация ответов в RL-дообучении, дорогая часть процесса по вычислениям и памяти. Авторы утверждают, что существующие FP4-подходы не сводят напрямую расхождение между квантованными путями обучения и rollout, и предлагают это исправить. Если результаты подтвердятся, RL-дообучение MoE-моделей в низкой точности станет дешевле без заметной потери качества.
Кому это важно
Командам, которые дообучают большие языковые модели (в том числе MoE) с подкреплением и упираются в стоимость rollout. Также исследователям, работающим над квантованием и низкоточным обучением.
Как это применить
Пока это описание метода в статье. О выпуске кода или моделей в тексте источника ничего не сказано. Практически идеи можно взять на заметку: направлять округление в FP4 при обучении по результатам квантования на стороне rollout и хранить мантиссу и масштаб только из глубоких слоёв.
Можно ли доверять
Все утверждения, собственные результаты авторов, приведённые в аннотации. В тексте нет названий четырёх моделей, названий тестов и числовых показателей качества, кроме ускорения. «Сопоставимо с BF16» не выражено в цифрах. Ускорение «до 5,4 раза», максимум, а не типичное значение, и с чем именно оно сравнивается, источник не уточняет. Независимой проверки в материале нет.
Риски и подводные камни
«До 5,4 раза» не означает такой выигрыш во всех сценариях. Выводы сделаны на четырёх MoE-моделях и описанных типах задач, перенос на другие модели не показан. Метод добавляет собственные накладные расходы на хранение и передачу данных о квантовании, хотя схема кэширования призвана их снизить.