Новый метод балансировки нагрузки в MoE-моделях
Новый метод балансировки нагрузки в MoE-моделях Команда Qwen разработала новый метод LBL-лосса для обучения моделей Mixture-of-Experts (MoE). В таких моделях роутер распределяет токены между экспертами. Новый LBL-лосс помогает сделать это распределение более равномерным и избежать перегрузки одних экспертов и недогрузки других. Традиционно LBL рассчитывают для каждого микробатча отдельно, что может привести к потере специализации экспертов. Авторы предлагают рассчитывать LBL на уровне глобального батча, где больше разнообразия данных. Для этого они добавляют шаг коммуникации: синхронизируют статистику выбора экспертов со всей DP-группы. Это позволяет роутеру более гибко распределять токены и сохранять специализацию экспертов. Исследования показали, что модели с глобальной балансировкой демонстрируют лучшие результаты и чёткую специализацию экспертов по доменам (код, математика, разные языки). При этом метод не замедляет процесс обучения.