Блог Qwen: балансировка нагрузки MoE по глобальному батчу (конспект)

На чём основан конспект

Исходник — content/blog/global-load-balance/index.md в QwenLM/QwenLM.github.io. Лицензии на текст нет — пересказ своими словами; формулы перенесены из оригинала, графики описаны словами, пример с числами — наш.

Коротко

• В MoE-моделях к основной потере добавляют вспомогательную потерю балансировкивспомогательный лосс балансировки (aux loss)EN auxiliary load-balancing lossДобавка к основному лоссу, штрафующая неравномерную загрузку экспертов. Классическая форма из Switch Transformer — α·N·Σ f_i·P_i. (load balancing loss, LBL), чтобы маршрутизатормаршрутизатор (роутер), гейтEN router / gating network не загружал одних экспертов и не забывал про других.
• Почти все фреймворки (например, Megatron-core) считают её на каждом микробатче отдельно и потом усредняют.
• Беда: микробатч часто состоит из данных одного домена (например, только код). Потеря всё равно заставляет разбрасывать эти токены по всем экспертам — и эксперты не специализируются.
• Решение Qwen: считать частоты выбора экспертов по глобальному батчу. Синхронизировать нужно лишь вектор длины «число экспертов» — это почти бесплатно.
• Итог в экспериментах авторов: качество лучше во всех настройках, эксперты заметно специализируются по доменам.

Как читать вместе с курсом

Читайте после урока 5.3.1 — там выведена та же формула потери балансировки. Пост хорошо дополняет урок 5.3.2: DeepSeek в V3 пошёл другим путём — отказался от вспомогательной потери совсем (балансировка без вспомогательного лоссабалансировка без вспомогательного лоссаEN auxiliary-loss-free load balancingМетод DeepSeek (arXiv:2408.15664, использован в V3) — к оценкам роутера прибавляется смещение b_i, влияющее только на выбор экспертов; после каждого шага b_i сдвигается на ±γ по загрузке.), а Qwen оставил потерю, но поменял, на чём её считать. Понятия микробатча и глобального батча — в уроке 5.4.3.

1. Напоминание: MoE-слой и потеря балансировки

MoE-слой — маршрутизатор (обычно один линейный слой) и группа экспертов (в трансформере каждый эксперт — FFN). Для входа $\mathbf{x}$ активируются только $K$ экспертов с наибольшими оценками, и их выходы смешиваются с весами маршрутизатора:

$$\mathbf{y} = \sum_{i \in N_E,\ g_i \in \operatorname{topK}} g_i(\mathbf{x})\, E_i(\mathbf{x}).$$

Потеря балансировки:

$$L_\text{balance} = N_E \sum_{i=1}^{N_E} f_i\, p_i,$$

где $N_E$ — число экспертов, $f_i$ — частота, с которой выбирали эксперта $i$, а $p_i$ — средняя оценка маршрутизатора для эксперта $i$. Словами: штраф велик, когда эксперт и часто выбирается, и получает высокие оценки; минимум (равный 1) — при равномерной загрузке.

2. Где прячется ошибка

Ключевая мысль поста: важно, по какому набору токенов считать $f_i$ и $p_i$. Если по микробатчу — а именно так сделано в большинстве фреймворков, — то потеря требует равномерности внутри каждого микробатча. Пока данные в микробатче разнообразны, это безобидно. Но при обучении LLM микробатч часто целиком из одного домена: подряд идущие фрагменты одного длинного документа, один источник данных. Тогда потеря толкает маршрутизатор раскидывать, скажем, только код по всем экспертам — и «эксперт по коду» не может возникнуть. Авторы считают, что это отчасти объясняет, почему у большинства открытых MoE-моделей заметной специализации экспертов нет.

Пример с числами (наш)

Четыре эксперта, два микробатча. Первый — только код, маршрутизатор естественно отправляет его в экспертов 1 и 2; второй — только обычный текст, он уходит в экспертов 3 и 4.

	$f$ (частоты)	$p$ (средние оценки)	$N_E \sum f_i p_i$
микробатч 1 (код)	0,5; 0,5; 0; 0	0,45; 0,45; 0,05; 0,05	$4 \cdot (0{,}5 \cdot 0{,}45 \cdot 2) = 1{,}8$
микробатч 2 (текст)	0; 0; 0,5; 0,5	0,05; 0,05; 0,45; 0,45	$1{,}8$
среднее по микробатчам			1,8 — сильный штраф
глобальный батч	0,25; 0,25; 0,25; 0,25	0,25; 0,25; 0,25; 0,25	$4 \cdot 4 \cdot 0{,}25 \cdot 0{,}25 =$ 1,0 — минимум

Одно и то же поведение маршрутизатора — специализация «код → эксперты 1–2, текст → эксперты 3–4» — микробатчевая потеря штрафует, а глобальная считает идеально сбалансированным. Ведь в целом все четыре эксперта загружены одинаково.

3. Как считать по глобальному батчу

Рецепт из поста в три шага:

1. синхронизировать частоты выбора экспертов $f_i$ между всеми параллельными группами (всеми GPU);
2. посчитать потерю балансировки в каждой группе, но с глобальными частотами;
3. сложить потери по всем микробатчам.

Формула:

$$L_\text{global} = N_E \sum_{i=1}^{N_E} f_i^\text{global}\, p_i^\text{global} = \frac{1}{N_p}\sum_{j=1}^{N_p}\Big(N_E \sum_{i=1}^{N_E} \bar{f_i}\cdot P_i^{j}\Big),$$

где $N_p$ — число параллельных групп (микробатчей), $\bar{f_i}$ — частота эксперта $i$ по всему глобальному батчу, а $P_i^j$ — средняя оценка эксперта $i$ в группе $j$. Средние оценки усредняются как обычно, поэтому собирать по всем GPU нужно только вектор $f$ длиной $N_E$ (например, 128 чисел) — одна маленькая операция all-reduceall-reduce (суммирование со всех на все)EN all-reduceКоллективная операция — каждый GPU отдаёт свой тензор и получает сумму (или среднее) тензоров всех GPU. Так DDP усредняет градиенты. на шаг. Отсюда «почти бесплатный обед» в названии.

4. Результаты

Эксперименты авторов — три конфигурации MoE (3,4 млрд параметров при 0,6 млрд активных; 15 млрд при 2,54; 43 млрд при 6,6) и два объёма данных (120 и 400 млрд токенов):

• глобальная балансировка лучше микробатчевой во всех настройках — по модели, данным и задачам;
• на тепловых картах активаций по доменам при микробатчевой потере почти все эксперты включаются одинаково независимо от домена, а при глобальной — некоторые эксперты явно предпочитают свой домен;
• размер батча балансировки: на модели 3,4B (0,6B активных) перплексия при предобучении быстро падает при росте батча балансировки от 2 до 128 и дальше почти не меняется. Авторы отмечают, что в распространённых фреймворках даже с обменом между группами экспертного параллелизма этот размер обычно 8–16 — далеко до насыщения;
• цена: глобальная балансировка может ухудшить равномерность внутри микробатча, а значит, и скорость (эксперты на разных GPU загружены неравно). Если добавить сверху маленькую микробатчевую потерю (вес 0,01 от глобальной), шаг ускоряется с 1,64 до 1,59 с, а качество почти не меняется.

5. Вывод

Глобальная потеря балансировки улучшает качество и даёт экспертам специализироваться. Эксперименты — на языковых задачах, но авторы надеются, что подход поможет обучать более крупные и специализированные MoE и в других областях. Идея вошла в обучение MoE-моделей Qwen (в техотчёте Qwen3 упомянута глобальная балансировка по батчу).

Мораль для инженера

«Демоны в деталях» (так называется статья) — хороший пример того, как одна строка реализации (где усреднять потерю) меняет поведение модели. Когда читаете код MoE, проверяйте: по какому множеству токенов считается каждая статистика?

Что сделать руками

1. Повторите таблицу из примера в Python для 8 экспертов и 4 микробатчей разных доменов.
2. В учебной MoE из урока 5.3.1 замените микробатчевую потерю глобальной (накопите $f$ за несколько шагов накопления градиента) и посмотрите на тепловую карту «домен × эксперт».

Словарик

English	Русский
load balancing loss (LBL)	потеря балансировки нагрузки
micro-batch / global batch	микробатч / глобальный батч
expert specialization	специализация экспертов
gating score	оценка маршрутизатора
activation frequency	частота выбора эксперта