TRACE: Обучение с учетом квантования на основе роллаута для FP4 метода усиления для моделей Mixture-of-Experts (MoE)

В этой статье представлено TRACE, фреймворк квантования FP4 для обучения методом усиления (RL) моделей Mixture-of-Experts (MoE). TRACE решает ограничение существующих методов FP4 RL, используя результаты квантования на стороне роллаута для направления решений округления FP4 на стороне обучения, уменьшая разницу между тренировочным и роллаут-путями. Фреймворк также включает эффективную схему кэширования информации о квантовании для уменьшения объема хранения и коммуникации. Оценка на четырех масштабных моделях MoE показала, что TRACE обеспечивает производительность RL, сравнимую с BF16, с ускорением роллаута до 5,4 раза и хорошей конечной производительностью FP4 по сравнению с пост-головным квантованием FP4 политик, обученных на BF16.