Expert parallelism

Expert parallelism is a parallelism strategy for mixture-of-experts (MoE) models. Each expert’s feedforward layer lives on a different hardware accelerator. A router dispatches tokens to the appropriate experts and gathers the results. This approach scales models to far larger parameter counts without increasing computation cost because each token activates only a few experts.

DistributedConfig

Enable expert parallelism with the DistributedConfig class and the ep_size argument. Most MoE models default to token dispatch, so ep_size can be set independently of tp_size. A few, such as Llama 4 and Gemma 4, still default to masking and all-reduce ("ep_router" and "moe_tp_experts" in model.ep_plan). Masking is also available on any model with an ep_plan override, and requires ep_size=tp_size so every rank in an expert group receives the same tokens.

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,
)

Each MoE model defines two plans in its config: base_model_tp_plan for the dense modules and base_model_ep_plan for the experts. They are exposed on the loaded model as model.tp_plan and model.ep_plan. With tp_size > 1 and ep_size > 1, both apply: the tensor parallel plan shards attention and the dense MLPs, and the expert parallel plan shards the experts. EP rules take precedence over TP rules for the same modules, so expert weights are sharded once, by the EP plan. In the EP plan, the GroupedGemmParallel style splits the expert weights along the expert dimension so each rank loads only its local experts, and ep_router masks the experts that live on other ranks before an all-reduce combines the expert outputs.

tp_plan is applied only when tp_size > 1, and ep_plan only when ep_size > 1. With TP enabled and EP disabled, the full TP plan applies, expert rules included.

The expert forward rule in ep_plan selects how tokens reach the experts:

enable_expert_parallel=True is a deprecated alias for ep_size=tp_size, used only when ep_size is omitted, and emits a FutureWarning.

Launch your inference script with torchrun. The number of processes must equal tp_size * fsdp_size * pp_size, and ep_size must evenly divide the number of experts.

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

Overriding the plans

Pass tp_plan={...} or ep_plan={...} to DistributedConfig to override individual rules of the predefined plans. Unspecified rules are kept, and the merged plans are stored on the model. Each key must match a module, a parameter, or an existing plan entry; otherwise loading raises a ValueError before anything is sharded. Use the full path as seen from the loaded model, so model.layers.* for a causal LM and layers.* for its base model.

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"},
)

Providing a plan does not infer parallel sizes: set tp_size and ep_size explicitly.

Most MoE models default to "ep_dispatch_experts". To use masking and all-reduce instead, set ep_size=tp_size and override both the router and the expert forward rules (the router module name depends on the model, e.g. mlp.router on 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",
    },
)

Conversely, override the expert forward rule of a model whose plan uses masking with "ep_dispatch_experts" to use token dispatch. The router rule is then ignored, since dispatch needs the global expert ids to find each expert’s owner.

Token dispatch

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)

With tp_size=1, ep_size must divide fsdp_size and the number of experts, and attention is not bound by num_key_value_heads. For the rest of the model:

• The parameters outside the experts are sharded with FSDP2 across fsdp, which reduces their gradients.
• The experts are sharded across ep and, when efsdp_size = fsdp_size * tp_size // ep_size is larger than one, additionally FSDP-sharded across efsdp. They are always FSDP-wrapped, so fsdp_mixed_precision and fsdp_cpu_offload apply to them too and save_pretrained() gathers them like any other parameter.
• An expert parallel group holds ep_size / tp_size batches, so an expert’s gradient is a sum over that many batches. The efsdp reduction divides by fsdp_size instead of its group size, which gives the same per-batch average FSDP2 takes for the dense modules.
• Every local expert also processes one zero pad row per layer. A rank whose experts received no tokens still joins the reverse all-to-all and the expert gradient reduction.

How the sizes combine

• tp_size * fsdp_size is the number of processes. ep_size adds none: it regroups the same ranks for the expert weights only.
• ep_size cuts the expert list into ep_size blocks. Each rank computes num_experts / ep_size experts, and ep_size consecutive ranks hold one complete set. That set of ranks is the group the all-to-all runs in.
• efsdp_size = fsdp_size * tp_size / ep_size is how many complete copies of the expert set exist. Ranks at the same position in different copies shard those experts for memory and average their gradients, like FSDP does for the dense modules.
• The batch a rank holds depends on fsdp only. Consecutive ranks form a TP group and get the same batch; the fsdp_size groups get different batches.

Eight processes, tp_size=2, fsdp_size=4, eight experts:

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

Two numbers follow from the picture:

• Inside an EP group, each token exists tp_size times, once per rank of the pair that holds its batch. This does not depend on ep_size.
• An EP group holds ep_size / tp_size different batches. This is the count an expert’s gradient sums over.

With tensor parallelism

Set tp_size > 1 to shard the dense modules with the TP plan while the experts use dispatch. On eight processes:

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

Each pair of TP ranks receives the same batch, because tensor parallelism replicates the activations inside the pair. If both ranks dispatched all of their tokens, the rank owning an expert would receive every token twice, compute it twice, and its weight gradient would double. Experts are whole on one rank, so the duplicate cannot be split by weights, and the owner is usually another rank, so it cannot be resolved by ownership as masking does. The pair therefore splits the rows: each TP rank dispatches a disjoint 1 / tp_size of the tokens, results come back to the rank that sent them, and an all-reduce over the pair of the zero-padded halves restores the replicated output the next layer expects. The split is by tp_size, not ep_size, since only the ranks that hold a batch can send it. Expert groups span four ranks and each expert is FSDP-sharded across efsdp_size = 2 ranks, while the trunk’s FSDP group spans four ranks. The model’s usual TP constraints, such as attention-head divisibility, still apply to the dense modules. Token slices may be uneven or empty, including during single-token decoding.

Token dispatch cannot be combined yet with pipeline parallelism yet; use pp_size=1 (not tested yet) These configurations each use eight GPUs:

The Trainer does not account for token dispatch yet: batch and token counting assume the all-reduce layout, where the ranks of a TP group share a batch and fsdp_size data-parallel shards exist. Trainer support for dispatch comes in a follow-up.

Combining with FSDP2

Tensor and expert parallelism shard the weights across tp, but the optimizer state and the modules without a rule are still replicated on every rank of the group, which limits how large a model you can train. Add FSDP2 on a second mesh dimension with fsdp_size. With masking and all-reduce, keep ep_size=tp_size and pass `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)

Load the model as usual, then train with Trainer. It takes the gradient norm across both meshes and gives each mesh its own optimizer param group. save_model() gathers sharded weights into a regular checkpoint. This requires accelerate>=1.12 so the Trainer can mirror tp_size and fsdp_size into ~Accelerate.ParallelismConfig.

The table below compares EP-only training with 2D EP+FSDP2 on 8xH100 GPUs. The workload is full fine-tuning of Qwen3-30B-A3B in bf16 at sequence length 2048. More FSDP shards cut peak memory, and tokens/s drop some because FSDP2 all-gathers and reduce-scatters the experts across fsdp.

Resuming from a checkpoint is not supported yet for models sharded at load time, so the Trainer only accepts save_only_model=True or save_strategy="no" for them.

API reference