Глобальный батч-баланс: почти бесплатный способ улучшить обучение MoE LLM

Исследователи из Alibaba Qwen представили метод глобальной балансировки батчей для обучения языковых моделей на архитектуре MoE. Стандартная практика в существующих фреймворках, таких как Megatron-core, использует микро-батч-баланс, при котором потери балансировки вычисляются внутри каждого микробатча и усредняются по глобальному батчу. Проблема в том, что если микробатч содержит однородные данные (например, только код), то потери балансировки заставляют распределять токены равномерно по всем экспертам, что мешает их специализации и может ухудшать производительность. Авторы предлагают синхронизировать частоту выбора экспертов между всеми параллельными группами и вычислять потери на уровне глобального батча. Такой подход почти не требует дополнительных вычислительных затрат и позволяет экспертам специализироваться по доменам. Эксперименты проводились с тремя конфигурациями MoE (3.4B/0.6B активированных, 15B/2.54B, 43B/6.6B) и двумя объёмами данных (120B и 400B токенов). Во всех случаях глобальный баланс показал лучшую производительность, чем микро-батч-баланс, и обеспечил заметную доменную специализацию экспертов. Дополнительные эксперименты показали, что увеличение размера батча для балансировки с 2 до 128 резко снижает perplexity предобучения, а затем эффект насыщается. В современных фреймворках размер батча для балансировки обычно составляет 8-16, что подтверждает значимость предложенного метода. Авторы также выяснили, что добавление небольшого веса микро-батч-баланса поверх глобального не ухудшает качество, но ускоряет обучение (с 1,64 до 1,59 секунды на шаг). Работа опубликована в arXiv и может быть полезна для обучения более эффективных MoE-моделей в различных областях.