Экспертная параллельность

Экспертная параллельность — это стратегия параллелизма для моделей mixture-of-experts (MoE). Каждый экспертный слой прямой передачи размещается на отдельном аппаратном ускорителе. Роутер направляет токены к соответствующим экспертам и собирает результаты. Этот подход позволяет масштабировать модели до значительно большего количества параметров без увеличения вычислительных затрат, поскольку каждый токен активирует лишь несколько экспертов.

DistributedConfig

Включите экспертную параллельность с помощью класса DistributedConfig и аргумента ep_size. Большинство MoE-моделей по умолчанию используют token dispatch, поэтому ep_size можно задавать независимо от tp_size. Некоторые модели, такие как Llama 4 и Gemma 4, по-прежнему по умолчанию используют маскирование и all-reduce ("ep_router" и "moe_tp_experts" в model.ep_plan). Маскирование также доступно для любой модели с переопределением ep_plan и требует, чтобы ep_size=tp_size, чтобы каждый ранк в экспертной группе получал одни и те же токены.

import os

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.distributed.configuration_utils import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=int(os.environ["WORLD_SIZE"]),
    ep_size=int(os.environ["WORLD_SIZE"]),
)

model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-120b",
    distributed_config=distributed_config,
)

Каждая MoE-модель определяет два плана в своём конфиге: base_model_tp_plan для плотных модулей и base_model_ep_plan для экспертов. Они доступны в загруженной модели как model.tp_plan и model.ep_plan. При tp_size > 1 и ep_size > 1 применяются оба: план тензорного параллелизма разбивает внимание и плотные MLP, а план экспертного параллелизма разбивает экспертов. Правила EP имеют приоритет над правилами TP для одних и тех же модулей, поэтому веса экспертов разбиваются один раз, согласно плану EP. В плане EP стиль GroupedGemmParallel разбивает веса экспертов вдоль измерения экспертов, так что каждый ранг загружает только своих локальных экспертов, а ep_router маскирует экспертов, находящихся на других рангах, до того, как all-reduce объединит выходы экспертов.

tp_plan применяется только при tp_size > 1, а ep_plan — только при ep_size > 1. Если TP включён, а EP отключён, применяется полный план TP, включая правила для экспертов.

Правило экспертного прямого прохода в ep_plan определяет, как токены поступают к экспертам:

Использование enable_expert_parallel=True без ep_size устарело и вызывает FutureWarning. Это разделяет только экспертов, как в случае ep_size=tp_size, tp_plan={}.

Запустите скрипт инференса с помощью torchrun. Количество процессов должно быть равно tp_size * fsdp_size * pp_size, а ep_size должно делить количество экспертов нацело.

torchrun --nproc-per-node 8 your_script.py

Переопределение планов

Передайте tp_plan={...} или ep_plan={...} в DistributedConfig, чтобы переопределить отдельные правила предопределённых планов. Нерасписанные правила сохраняются, а объединённые планы записываются в модель. Каждый ключ должен соответствовать модулю, параметру или существующей записи плана; в противном случае при загрузке возникает ValueError до начала шардирования. Используйте полный путь, как он виден из загруженной модели, то есть model.layers.* для causal LM и layers.* для её базовой модели.

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,
    tp_plan={"model.layers.*.self_attn.q_proj": "colwise_rep"},
    ep_plan={"model.layers.*.mlp.experts.down_proj": "grouped_gemm"},
)

Указание плана не определяет размеры параллелизма: явно задайте tp_size и ep_size.

Большинство MoE-моделей по умолчанию используют "ep_dispatch_experts". Чтобы вместо этого применить маскирование и all-reduce, установите ep_size=tp_size и переопределите правила forward для обоих модулей — роутера и экспертов (имя модуля роутера зависит от модели, например, mlp.router в gpt-oss):

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,
    ep_plan={
        "model.layers.*.mlp.gate": "ep_router",
        "model.layers.*.mlp.experts": "moe_tp_experts",
    },
)

Напротив, переопределите правило forward для экспертов модели, чей план использует маскирование, с помощью "ep_dispatch_experts", чтобы использовать диспетчеризацию токенов. Правило роутера в этом случае игнорируется, поскольку для диспетчеризации необходимы глобальные идентификаторы экспертов, чтобы определить владельца каждого эксперта.

Диспетчеризация токенов

With token dispatch, each rank trains on its own part of the batch. At every MoE layer, a rank routes its tokens, sends each (token, expert) pair to the rank that owns the expert with an all-to-all, runs its local experts on what it receives, gets the results back with a second all-to-all and combines them with the routing weights. Only the routed activations and expert outputs travel, and no rank computes experts for tokens it does not own.

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=1,
    fsdp_size=8,
    ep_size=4,
)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_config=distributed_config)

При tp_size=1 ep_size должно делить fsdp_size и количество экспертов, а внимание не ограничивается num_key_value_heads. Для остальной части модели:

• Параметры вне экспертов распределяются с помощью FSDP2 по fsdp, что уменьшает их градиенты.
• Эксперты разбиты на шарды по ep и, когда efsdp_size = fsdp_size * tp_size // ep_size больше единицы, дополнительно разбиты на шарды FSDP по efsdp. Они всегда обёрнуты в FSDP, поэтому fsdp_mixed_precision и fsdp_cpu_offload применяются и к ним, а save_pretrained() собирает их так же, как и любой другой параметр.
• Группа экспертного параллелизма содержит ep_size / tp_size батчей, поэтому градиент эксперта представляет собой сумму по такому количеству батчей. Редукция efsdp делит на fsdp_size, а не на размер своей группы, что даёт тот же средний по батчу результат, который FSDP2 использует для плотных модулей.
• Каждый локальный эксперт также обрабатывает одну строку нулевых отступов на каждый слой. Ранг, эксперты которого не получили токены, всё равно участвует в обратном all-to-all и в редукции градиентов экспертов.

Как комбинируются размеры

• tp_size * fsdp_size — это количество процессов. ep_size не добавляет новых: он лишь перегруппировывает те же ранги, но только для весов экспертов.
• ep_size разбивает список экспертов на ep_size блоков. Каждый ранг вычисляет num_experts / ep_size экспертов, а ep_size последовательных рангов хранят один полный набор. Этот набор рангов является группой, в которой выполняется all-to-all.
• efsdp_size = fsdp_size * tp_size / ep_size — это количество полных копий набора экспертов. Ранги, находящиеся на одной и той же позиции в разных копиях, разделяют этих экспертов по памяти и усредняют их градиенты, как это делает FSDP для плотных модулей.
• Партия, которую удерживает ранг, зависит только от fsdp. Последовательные ранги образуют группу TP и получают одну и ту же партию; группы fsdp_size получают разные партии.

Восемь процессов, tp_size=2, fsdp_size=4, восемь экспертов:

rank         0      1       2      3       4      5       6      7
batch       [====B0====]   [====B1====]   [====B2====]   [====B3====]     one batch per TP pair
tp           0      1       0      1       0      1       0      1

ep_size=2   E0-3   E4-7    E0-3   E4-7    E0-3   E4-7    E0-3   E4-7     group = a TP pair,   efsdp_size=4
ep_size=4   E0E1   E2E3    E4E5   E6E7    E0E1   E2E3    E4E5   E6E7     group = two pairs,   efsdp_size=2
ep_size=8   E0     E1      E2     E3      E4     E5      E6     E7       group = all ranks,   efsdp_size=1

Из рисунка следуют два числа:

• Внутри группы EP каждый токен существует tp_size раз, по одному разу на каждый ранг пары, которая хранит его батч. Это не зависит от ep_size.
• Группа EP содержит ep_size / tp_size различных батчей. Это количество, по которому суммируются градиенты эксперта.

С параллелизмом тензоров

Установите tp_size > 1, чтобы разбить плотные модули по плану TP, в то время как эксперты используют dispatch. На восьми процессах:

distributed_config = DistributedConfig(
    tp_size=2,
    fsdp_size=4,
    ep_size=4,
)

Каждая пара рангов TP получает один и тот же батч, поскольку тензорное параллелизм дублирует активации внутри пары. Если бы оба ранга передали все свои токены, ранг, владеющий экспертом, получил бы каждый токен дважды, вычислил бы его дважды, и градиент его весов удвоился бы. Эксперты целиком находятся на одном ранге, поэтому дубликат нельзя разделить по весам, а владелец обычно является другим рангом, поэтому проблему нельзя решить через владение, как это делает маскирование. Поэтому пара разделяет строки: каждый ранг TP передает непересекающуюся часть 1 / tp_size токенов, результаты возвращаются на ранг, который их отправил, а all-reduce по паре нулевых дополненных половин восстанавливает дублированный вывод, который ожидает следующий слой. Разделение выполняется по tp_size, а не по ep_size, поскольку только ранги, хранящие батч, могут его отправить. Группы экспертов охватывают четыре ранга, и каждый эксперт шардирован FSDP по efsdp_size = 2 рангам, в то время как FSDP-группа ствольной части охватывает четыре ранга. Обычные ограничения TP модели, такие как делимость заголовков внимания, по-прежнему применяются к плотным модулям. Срезы токенов могут быть неравными или пустыми, в том числе во время декодирования одного токена.

Пока нельзя использовать диспетчеризацию токенов вместе с параллелизмом по конвейеру; используйте pp_size=1 (ещё не протестировано) Каждая из этих конфигураций использует восемь GPU:

Класс Trainer пока не учитывает диспетчеризацию токенов: подсчёт батчей и токенов предполагает компоновку all-reduce, в которой ранги группы TP разделяют один батч, а существует fsdp_size шардов с параллелизмом данных. Поддержка диспетчеризации в Trainer будет добавлена в последующих обновлениях.

Комбинирование с FSDP2

Параллелизм тензоров и экспертов распределяет веса по tp, но состояние оптимизатора и модули без правил по-прежнему дублируются на каждом ранге группы, что ограничивает размер модели, которую можно обучать. Добавьте FSDP2 по второму измерению сетки с fsdp_size. С использованием маскирования и all-reduce, поддерживайте ep_size=tp_size и передавайте `ep_plan={“layers..mlp.gate”: “ep_router”, “layers..mlp.experts”: “moe_tp_experts”}“.

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,  # expert parallel size, must match tp_size with masking and all-reduce
    fsdp_size=2,  # data parallel shards
    ep_plan={
        "model.layers.*.mlp.gate": "ep_router",
        "model.layers.*.mlp.experts": "moe_tp_experts",
    },
)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_config=distributed_config)

Загрузите модель как обычно, затем обучите с помощью Trainer. Он вычисляет норму градиента по обеим сеткам и назначает каждой сетке собственную группу параметров оптимизатора. save_model() собирает разбитые веса в обычный чекпоинт. Для этого требуется accelerate>=1.12, чтобы Trainer мог отразить tp_size и fsdp_size в ~Accelerate.ParallelismConfig.

Таблица ниже сравнивает обучение только с EP с обучением 2D EP+FSDP2 на 8xH100 GPU. Нагрузкой является полное дообучение Qwen3-30B-A3B в bf16 при длине последовательности 2048. Увеличение количества шардов FSDP снижает пиковое потребление памяти, а токены/с немного падают, так как FSDP2 выполняет all-gather и reduce-scatter экспертов через fsdp.

Возобновление обучения из контрольной точки пока не поддерживается для моделей, разбитых на части при загрузке, поэтому Trainer принимает для них только save_only_model=True или save_strategy="no".

Справочник API