Hugging Face's logo

Transformers documentation

Expert parallelism

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.15.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

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

class transformers.DistributedConfig

< >

( tp_size: int | None = Nonetp_plan: typing.Union[dict[str, str], typing.Literal['auto'], NoneType] = Noneenable_sequence_parallel: bool = Falseenable_expert_parallel: bool = Falsefsdp_size: int | None = Nonefsdp_cpu_offload: bool = Falsefsdp_mixed_precision: bool = Falsepp_size: int | None = None )

Parameters

  • tp_size (int, optional) — Number of devices for tensor parallelism. If None and tp_plan is set, defaults to WORLD_SIZE // (other_parallel_size). If None and no tp_plan is set, defaults to 1.
  • tp_plan (dict[str, str] or “auto”, optional) — Tensor parallel sharding plan. Pass “auto”, or leave as None when tp_size is set, to use the model’s predefined base_model_tp_plan. Pass a dictionary to override the predefined plan.
  • enable_sequence_parallel (bool, optional, defaults to False) — Reserved for sequence parallelism. Not wired up yet.
  • enable_expert_parallel (bool, optional, defaults to False) — Route MoE models through the expert-parallel path (base_model_ep_plan).
  • fsdp_size (int, optional) — Number of devices for FSDP (data parallelism). If None and tp_size is set, defaults to 1.
  • fsdp_cpu_offload (bool, optional, defaults to False) — Whether to enable CPU offloading for FSDP2.
  • fsdp_mixed_precision (bool, optional, defaults to False) — Whether to enable mixed precision for FSDP2.
  • pp_size (int, optional) — Number of devices for pipeline parallelism. If None and another parallel mode is set, defaults to 1.

Configuration for native distributed inference and training with tensor, pipeline, or FSDP2 parallelism.

Update on GitHub

←Tensor parallelism for inference How caching works→