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

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

DistributedConfig

Включите экспертную параллельность с помощью класса DistributedConfig и аргумента ep_size. Модели MoE разделяют своих экспертов с помощью "ep_dispatch_experts". Их реализации определяют, как токены поступают к экспертам. Token dispatch позволяет ep_size отличаться от tp_size. Некоторые реализации EP требуют ep_size=tp_size, например "all_reduce" и "deepgemm_megamoe", поэтому каждый ранг в группе экспертов хранит одни и те же токены.

У Llama 4 нет плана EP: её эксперты обрабатывают каждый токен, уже масштабированный по его оценке от роутера, поэтому ничего распределять не нужно. Вместо этого используйте тензорное параллелизм (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_dispatch_experts" для модуля экспертов определяет, как токены к ним поступают.

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

"ep_dispatch_experts" распределяет целых экспертов по рангам (ep_size — кратное tp_size число, делящее fsdp_size * tp_size), и обращается к ним через реализацию, указанную в ep_implementation, см. Выбор реализации.

Выбор реализации

"ep_dispatch_experts" разбивает экспертов один раз и выбирает, как токены к ним поступают, из реализаций, зарегистрированных в EPImplementationInterface, которые выбираются с помощью ep_implementation в DistributedConfig:

ep_implementation выбирает только среди размеров параллелизма, заданных в DistributedConfig. Если ранги группы экспертов обрабатывают разные батчи, например tp_size=1 при fsdp_size=ep_size, то единственной применимой реализацией является "all_to_all".

Добавление реализации

Реализация наследует EPImplementation и регистрируется под именем ep_implementation selects. Она создаётся один раз на модуль экспертов, поэтому может хранить состояние, например, буферы коммуникации. Вызов с батчем, который хранит данный ранк, и глобальными идентификаторами экспертов возвращает взвешенный вывод экспертов для каждого из его токенов. Она запускает экспертов данного ранка с помощью experts_forward на строках, локальных идентификаторах экспертов и весах, которые она им передала.

from transformers.distributed.expert_parallel import EPImplementation, EPImplementationInterface

class MyEP(EPImplementation):
    def __call__(self, hidden_states, top_k_index, top_k_weights):
        tokens, expert_ids, weights = ...  # bring this rank's experts the rows routed to them
        expert_output = self.experts_forward(tokens, expert_ids, weights)
        return ...  # the weighted expert output for each token of the batch this rank holds

EPImplementationInterface.register("my_ep", MyEP)
distributed_config = DistributedConfig(tp_size=1, fsdp_size=8, ep_size=8, ep_implementation="my_ep")

Градиенты должны доходить до каждого ранга, на котором хранится токен, включая веса маршрутизации, поэтому коллективные операции, которые autograd не дифференцирует, необходимо оборачивать в torch.autograd.Function.

Использование 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.

При ep_size=tp_size реализация по умолчанию "auto" выполняет "all_reduce" на этапе инференса и распределяет токены при обучении. Чтобы маскировать и выполнять all-reduce также при обучении, установите ep_implementation="all_reduce":

distributed_config = DistributedConfig(tp_size=4, ep_size=4, ep_implementation="all_reduce")

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

Token dispatch is the "all_to_all" implementation: 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 и ep, но состояние оптимизатора и модули без правил по-прежнему дублируются на каждом ранге группы, что ограничивает размер модели, которую можно обучать. Добавьте FSDP2 по другому измерению сетки с помощью fsdp_size. С использованием маскирования и all-reduce, установите ep_size=tp_size и передайте ep_implementation="all_reduce".

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_implementation="all_reduce",  # "auto" would dispatch tokens in training, which the Trainer does not support yet
)
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