Что такое expert parallelism для MoE моделей (Mixtral)?

Краткий тезис

Expert parallelism — это техника распределённого обучения и инференса, специально разработанная для Mixture-of-Experts (MoE) моделей. Вместо того чтобы реплицировать все эксперты на каждом ускорителе (GPU), разные эксперты размещаются на разных устройствах, и при активации токена нужные эксперты динамически подгружаются по сети. Такой подход позволяет обучать колоссальные модели с триллионами параметров, но требует тщательной балансировки нагрузки и вносит значительный overhead explosion|communication overhead. В архитектуре Mixtral 8x7B каждый слой содержит 8 экспертов, из которых для каждого токена активируются только 2, что делает этот параллелизм особенно эффективным.

1. Термин: MoE (Mixture of Experts)

Mixture of Experts — это архитектура нейронной сети, в которой вместо одного большого плотного слоя используется множество специализированных подсетей (экспертов) и маршрутизатор (router). Маршрутизатор для каждого входного токена выбирает, к каким экспертам отправить данные, и комбинирует их выходы. Ключевая идея: каждый токен обрабатывается лишь небольшим подмножеством экспертов, что экономит вычислительные ресурсы при огромном общем количестве параметров.

В MoE-слое обычно:

• E экспертов (например, 8 или 64).
• Каждый эксперт — это, как правило, FFN (Feed-Forward Network) с двумя линейными преобразованиями и нелинейностью.
• Router — небольшой линейный слой + softmax, выдающий вероятности выбора экспертов.
• Топ-2 (top-k) эксперта с наибольшими весами активируются, их выходы взвешенно суммируются.

Пример: Mixtral 8x7B имеет 8 экспертов в каждом MoE-слое. Несмотря на 8×7 = 56 млрд параметров, на инференсе для каждого токена используются только 2 эксперта (≈14 млрд параметров), поэтому скорость сравнима с моделью 14B, а качество — с плотной моделью 56B.

2. Проблема: стандартные стратегии параллелизма для MoE

Традиционный data parallelism (и его варианты) предполагает, что каждый GPU хранит полную копию модели (все параметры). Для MoE это означает, что на каждом устройстве содержатся все эксперты. При размере эксперта 7B и 8 экспертах это 56B параметров на каждом GPU. Уже на одном GPU это не влезает, а на нескольких — катастрофически неэффективно по памяти.

Model parallelism (разрезание одного большого слоя по параметрам) тоже не решает проблему, потому что каждый экспертный слой по-прежнему требует присутствия всех экспертов в одном месте.

Expert parallelism решает эту проблему, распределяя эксперты по разным GPU, а не реплицируя их.

3. Что такое expert parallelism

Expert parallelism — это форма специализированного параллелизма для MoE, при котором разные эксперты размещаются на разных ускорителях (или группах ускорителей). Каждый GPU отвечает только за свои эксперты. Когда маршрутизатор решает, что токену нужен эксперт, находящийся на другом устройстве, активации (скрытые состояния) пересылаются по сети к этому эксперту, обрабатываются, и результат возвращается.

Процесс инференса с expert parallelism (упрощённо):

1. Входной батч разбивается на микро-батчи.
2. Каждый микро-батч проходит через общие слои (эмбеддинги, self-attention) — эти части могут быть реплицированы или тоже распределены.
3. На MoE-слое: маршрутизатор (он реплицирован на всех GPU) для каждого токена определяет, к каким экспертам отправить. Номер эксперта маппится на GPU, где этот эксперт хранится.
4. Токен пересылается (all-to-all коммуникация) на целевые GPU.
5. Целевые GPU выполняют forward pass для своих экспертов на полученных токенах.
6. Результаты отправляются обратно (ещё одна all-to-all).
7. Полученные на исходном GPU взвешенные выходы экспертов суммируются.

4. Сравнение expert parallelism и стандартного параллелизма для MoE

5. Балансировка нагрузки (load balancing)

При expert parallelism критична балансировка нагрузки (load balancing). Если маршрутизатор отправляет большинство токенов к паре экспертов, их GPU будут перегружены, а остальные простаивают. Это снижает утилизацию и увеличивает задержку.

Методы борьбы

• Auxiliary loss (балансировочная потеря): к основной функции потерь добавляется штраф, поощряющий равномерное распределение токенов по экспертам. Например, в Mixtral используется load-balancing loss с коэффициентом 0.01.
• Capacity factor: ограничение на максимальное количество токенов, которое может обработать один эксперт за шаг. Лишние токены сбрасываются или перераспределяются. Обычно capacity factor = 1.25 (на 25% больше среднего).
• Dynamic routing: адаптивные алгоритмы маршрутизации, учитывающие текущую загрузку экспертов.
• Expert placement: стратегическое размещение экспертов на GPU с учётом их популярности.

6. Преимущества expert parallelism

• Масштабирование до триллионных моделей: модель может содержать тысячи экспертов, каждый — на отдельном GPU. Параметры модели растут линейно с числом GPU. Например, Switch Transformer (1.6T параметров) обучался с expert parallelism.
• Эффективность вычислений: при top-2 активации на каждый токен тратится лишь малая часть всех параметров, что сильно экономит FLOPs.
• Гибкость: можно добавлять новых экспертов (дообучать) без изменения остальной модели.
• Возможность использовать expert parallelism в комбинации с другими видами параллелизма (data, tensor, pipeline) — так называемый 4D-параллелизм.

7. Недостатки и вызовы

• Communication overhead: all-to-all коммуникации на каждом MoE-слое могут составлять значительную долю времени, особенно при большом числе экспертов и медленной межсоединительной сети (например, Ethernet вместо NVLink). В Mixtral 8x7B это менее критично, так как экспертов немного, но для сотен экспертов проблема существенна.
• Неравномерная нагрузка: даже с aux loss возможны пиковые нагрузки, вызывающие простаивание части GPU.
• Сложность реализации: требуется специальный runtime (DeepSpeed, Megablocks, Tutel), управляющий динамической маршрутизацией и коммуникацией.
• Пропускная способность памяти: при пересылке активаций может упираться в PCIe/NVLink bandwidth, особенно на инференсе.

8. Реализации и инструменты

9. Пример архитектуры: Mixtral 8x7B

Mixtral 8x7B — одна из самых известных MoE-моделей, открытая Mistral AI. Её ключевые параметры:

• Transformer декодер с 32 слоями.
• Каждый второй FFN-слой заменён на MoE-слой с 8 экспертами (total 32 MoE-слоя).
• Каждый эксперт — FFN с размерностью скрытого слоя 14336 (как в Mistral 7B).
• Топ-2 активации (expert count = 2).
• Router — небольшой MLP с входом размером 4096 и выходом 8 (веса для экспертов).
• Общее количество параметров: 56B (8×7B), но на один токен активны ~14B (2 эксперта + общие слои).

Expert parallelism для Mixtral позволяет обслуживать его на 8 GPU (например, A100-80GB) с разумной памятью: каждый GPU хранит 7B эксперта, при этом требуется эффективная all-to-all между GPU для каждого MoE-слоя.

10. Комбинация expert parallelism с другими видами параллелизма

В больших кластерах для обучения крупных MoE-моделей используется 4D-параллелизм:

• Data parallelism: батчи распределяются по группам GPU.
• Tensor parallelism: внутри одного трансформер-слоя разбиваются линейные операции (для больших скрытых размерностей).
• Pipeline parallelism: слои модели разрезаются по стадиям.
• Expert parallelism: эксперты распределяются по GPU внутри одной data‑parallel группы.

Грамотное сочетание позволяет достичь хорошего scale-up без узких мест.

11. Влияние на latency и throughput

Latency (задержка): expert parallelism добавляет дополнительные all-to-all коммуникации на каждом MoE-слое. Для Mixtral с 32 MoE-слоями это 64 all-to-all за инференс (туда-обратно). При низкой latency сети (NVLink) overhead невелик, но на Ethernet может быть значительным. Для реального времени latency-sensitive приложений может потребоваться оптимизация (например, fuse коммуникации).

Throughput (пропускная способность): expert parallelism позволяет утилизировать больше GPU, увеличивая throughput за счёт большего батч-сайза. Проблема балансировки нагрузки может снизить throughput, если часть GPU простаивает.

Пет-проект для закрепления

Задача: Реализовать простой expert parallelism для маленькой MoE-модели с двумя экспертами на двух GPU, обучить её на synthetic data и сравнить по скорости с полной репликацией на одном GPU (data parallelism).

Инструменты: PyTorch, torch.distributed, one GPU simulation or actual multi-GPU (можно использовать Colab с TPU/GPU, но проще эмулировать с помощью torch.cuda.device).

Шаги:

1. Создать tiny MoE-модель: embedding 32 -> два эксперта (FFN 32->64->32), router (32->2).
2. Разделить эксперты: на device 0 хранить expert 0, на device 1 — expert 1.
3. Реализовать forward: каждый токен определяет топ-1 эксперта, пересылает скрытое состояние на соответствующий device, получает результат, суммирует.
4. Написать процесс обучения (оптимизатор Adam, loss = MSE) на случайных данных.
5. Замерить время одного шага и сравнить с baseline: обычный data parallelism (оба эксперта на каждом device, градиенты усредняются).
6. Добавить capacity factor и простой вспомогательный load-balancing loss (aux_loss = std(token_count_per_expert)).

Ожидаемый результат: Поймёте практические детали: необходимость all-to-all (torch.distributed.all_to_all_single), балансировку нагрузки, влияние коммуникаций на скорость.

Связь с другими вопросами

Навигация

• Предыдущий: 848
• Следующий: 850
• Индекс: 00. Индекс разборов