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 enable_expert_parallel argument.

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"]),
    enable_expert_parallel=True,
)

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

Expert parallelism automatically enables tensor parallelism for attention layers.

This argument switches to the ep_plan (expert parallel plan) defined in each MoE model’s config file. The GroupedGemmParallel class splits expert weights so each device loads only its local experts. The ep_router routes tokens to experts and an all-reduce operation combines their outputs.

Launch your inference script with torchrun and specify how many devices to use. The number of devices must evenly divide the total number of experts.

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

Combining with FSDP2

Expert parallelism only shards the experts. Everything else (attention, embeddings, norms) and its optimizer state is replicated on every expert-parallel rank, which limits how large a model you can train. Add FSDP2 on a second mesh dimension with fsdp_size, and keep using tp_size for the expert parallel width (tp_size is the EP size).

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=4,  # expert parallel size
    fsdp_size=2,  # data parallel shards
    enable_expert_parallel=True,
)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_config=distributed_config)

The model is loaded on a 2D (fsdp, tp) device mesh, and tp_size * fsdp_size must equal the number of processes. The expert parallel plan shards the experts across tp, then FSDP2 shards every parameter, experts included, across fsdp and owns their gradient reduction. Each fsdp rank trains on its own part of the batch.

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.