Expert Parallelism

Параллелизм экспертов (EP) — это стратегия распределенного обучения и инференса для моделей со смесью экспертов (MoE): разные эксперты размещаются на разных GPU, а каждый токен отправляется на устройства, где находятся выбранные для него эксперты. Слой MoE заменяет один большой блок прямого распространения трансформера множеством небольших экспертных сетей и маршрутизатором, который активирует для каждого токена лишь нескольких экспертов. Поэтому общее число параметров может расти намного быстрее, чем вычислительные затраты на обработку каждого токена. Параллелизм экспертов позволяет разместить эти параметры в памяти: вместо того чтобы хранить всех экспертов на каждом GPU, на каждом устройстве хранится только своя часть. Еще в 2017 году Ном Шазир и его коллеги из Google распределили экспертов между устройствами, разместив по одной общей копии на каждом, в разреженном слое MoE с маршрутизацией. В июне 2020 года GShard масштабировал этот подход, распределив экспертов модели перевода с более чем 600 миллиардами параметров по 2,048 ускорителям TPU v3. Параллелизм экспертов теперь стал стандартом обучения и обслуживания крупных языковых и визуально-языковых моделей MoE.

Как токены перемещаются между GPU#

Каждый слой MoE выполняет четыре шага:

1. Маршрутизация: Маршрутизатор оценивает каждый токен для каждого эксперта и выбирает k лучших экспертов — обычно одного, двух или восьми.
2. Распределение: Коллективная операция обмена «каждый с каждым» отправляет скрытое состояние каждого токена на GPU, где находятся выбранные эксперты.
3. Вычисление: Каждый GPU запускает локальных экспертов для полученных токенов.
4. Объединение: Вторая операция all-to-all возвращает результаты на исходный GPU каждого токена. Там результаты взвешиваются по оценкам маршрутизатора и суммируются.

Приведённый ниже фрагмент имитирует маршрутизацию и распределение для одного слоя с 8 экспертами на 4 GPU — по два эксперта на GPU:

import torch

torch.manual_seed(0)
num_experts, num_gpus, top_k = 8, 4, 2
tokens = torch.randn(4096, 512)  # скрытые состояния для 4,096 токенов
router = torch.nn.Linear(512, num_experts)

# Маршрутизация: каждый токен выбирает 2 лучших экспертов по оценке маршрутизатора
with torch.no_grad():
    _, experts = router(tokens).topk(top_k, dim=-1)

# Параллелизм экспертов: эксперты 0-1 находятся на GPU 0, эксперты 2-3 — на GPU 1 и так далее
gpu = experts // (num_experts // num_gpus)

# Распределение: подсчитываем назначения токенов экспертам, полученные каждым экспертом и GPU
print(torch.bincount(experts.flatten(), minlength=num_experts).tolist())
print(torch.bincount(gpu.flatten(), minlength=num_gpus).tolist())

В первой строке вывода показана нагрузка каждого эксперта: от 948 до 1 087 назначений при равномерной доле 1 024. Во второй строке указано, что четыре GPU получили от 1 968 до 2 139 из 8 192 назначений. Даже случайно инициализированный маршрутизатор распределяет работу неравномерно, а после обучения у маршрутизаторов появляются популярные эксперты, из-за которых нагрузка становится ещё более перекошенной.

Балансировка нагрузки и ёмкость#

Работа слоя завершается только тогда, когда закончит самый загруженный GPU, поэтому дисбаланс маршрутизации приводит к простою оборудования. В обучении MoE с ним борются тремя способами:

• Вспомогательная балансирующая функция потерь: Дополнительное слагаемое в функции потерь поощряет маршрутизатор равномерно распределять токены между экспертами.
• Коэффициент емкости: Каждый эксперт обрабатывает не больше заданного числа токенов в батче, например в 1.25 раза больше равной доли; переполненные токены пропускают этого эксперта.
• Балансировка на основе смещений: DeepSeek-V3 добавляет к оценкам маршрутизатора смещение для каждого эксперта, которое влияет только на выбор экспертов: после каждого шага оно снижается для перегруженных экспертов и повышается для простаивающих. Это балансирует нагрузку без вспомогательной функции потерь, конкурирующей с основной целевой функцией.

При ограничении емкости путь каждого токена зависит от загрузки его экспертов:

Пропущенный токен без изменений проходит через слой по остаточной связи, поэтому частое отбрасывание токенов снижает качество. Во время обучения DeepSeek-V3 не отбрасывает токены, полагаясь вместо этого на балансировку на основе смещений.

Еще одна статья расходов — обмен данными. Два обмена «каждый с каждым» происходят в каждом слое MoE, поэтому параллелизм экспертов эффективнее всего работает внутри узла, где GPU соединены скоростными каналами NVLink, и становится дорогим при использовании более медленных соединений между узлами. DeepSeek-V3 распределяет 256 маршрутизируемых экспертов в каждом слое с параллелизмом экспертов 64-way по 8 узлам, отправляет каждый токен экспертам не более чем на 4 узлах и совмещает обмен данными с вычислениями.

Сочетание стратегий параллелизма#

Параллелизм экспертов — один из нескольких способов распределить обучение между устройствами:

• Параллелизм данных: Реплицирует полную модель на каждом GPU и разделяет каждый батч, синхронизируя градиенты после каждого шага.
• Тензорный параллелизм: Разделяет отдельные матрицы весов между GPU, поэтому каждый GPU участвует в каждом слое.
• Конвейерный параллелизм: Назначает последовательные слои разным GPU и передает через них микробатчи.
• Контекстный параллелизм: Разделяет одну длинную последовательность входных данных между GPU.

Плотным моделям, которые помещаются на одном GPU, достаточно параллелизма по данным. Именно так Ultralytics YOLO26 обучается на нескольких GPU, когда в режиме обучения передаётся список устройств, например device=[0, 1]. Крупные MoE-модели сочетают несколько подходов: DeepSeek-V3 обучила свои 671 миллиард параметров, из которых 37 миллиардов активны для каждого токена, используя параллелизм конвейера с разделением на 16 частей, параллелизм экспертов с разделением на 64 части и шардированный параллелизм по данным, но без параллелизма по тензорам. Механизмы инференса, например vLLM, применяют такое же разделение при обслуживании модели, поэтому MoE-модель, эксперты которой не помещаются в память одного GPU, работает на нескольких GPU, каждый из которых хранит только свою часть экспертов.