Expert-parallel token dispatch: each rank trains on its own part of the batch

docs/source/en/expert_parallelism.md

@@ -50,9 +50,33 @@ Launch your inference script with [torchrun](https://pytorch.org/docs/stable/ela
 torchrun --nproc-per-node 8 your_script.py
 ```
 
+## Token dispatch
+
+By default, every expert parallel rank runs the whole batch, keeps only the experts it owns, and all-reduces expert outputs after every MoE layer. Dense layers then do `tp_size` times the same work, and the all-reduce moves full activations. Set `experts_dispatch="all-to-all"` to send each token to the rank that owns its experts. Each rank trains on its own batch shard, and a lot less data is required to travel between GPUs/nodes during large-scale training.
+
+```py
+from transformers import AutoModelForCausalLM
+from transformers.distributed import DistributedConfig
+
+distributed_config = DistributedConfig(
+    tp_size=8,
+    enable_expert_parallel=True,
+    experts_dispatch="all-to-all",
+)
+```
+
+Each rank trains on its own part of the batch. At every MoE layer it routes its tokens, sends each (token, expert) pair to the rank that owns the expert with an all-to-all, runs its local experts, and gets the results back with a second all-to-all. Only the routed tokens travel.
+
+For the rest of the model:
+
+- The parameters outside the experts are data-parallel across the whole group, so they are sharded with [FSDP2](./fsdp) across every rank (`fsdp` and `tp` together when both are set), and FSDP2 reduces their gradients.
+- The experts stay sharded across `tp`, and across `fsdp` too when `fsdp_size > 1`. With `fsdp_size=1` they are outside FSDP2, so `fsdp_mixed_precision` and `fsdp_cpu_offload` do not apply to them.
+- The [`Trainer`] gives each rank its own training batches and counts tokens across all of them. Evaluation is unchanged from plain expert parallelism: every `tp` rank sees the same batches.
+- Training needs a sized (map-style) dataset, `dispatch_batches=False`, and `train_sampling_strategy="random"`. Iterable datasets and other sampling strategies are rejected.
+
 ## 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](./fsdp) on a second mesh dimension with `fsdp_size`, and keep using `tp_size` for the expert parallel width (`tp_size` is the EP size).
+Without token dispatch, 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](./fsdp) on a second mesh dimension with `fsdp_size`, and keep using `tp_size` for the expert parallel width (`tp_size` is the EP size).
 
 ```py
 from transformers import AutoModelForCausalLM
@@ -66,7 +90,7 @@ distributed_config = DistributedConfig(
 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.
+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. With `experts_dispatch="all-to-all"` the non-expert parameters are data-parallel across the whole mesh rather than replicated across `tp`, so FSDP2 shards them across `fsdp` and `tp` together.
 
 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. [`~Trainer.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`].
 

src/transformers/distributed/configuration_utils.py

@@ -18,6 +18,12 @@
 from typing import Literal
 
 
+# How the expert outputs get back to the tokens that need them, mapped to the parallel style the experts take.
+# `all-reduce` runs the whole batch on every rank and needs no style of its own. A new backend such as DeepEP is a
+# new entry here plus the matching style in `ParallelInterface`, not a new flag.
+EXPERTS_DISPATCH_STRATEGIES = {"all-reduce": None, "all-to-all": "ep_dispatch_experts"}
+
+
 @dataclass
 class DistributedConfig:
     """
@@ -34,6 +40,12 @@ class DistributedConfig:
             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``).
+        experts_dispatch (`str`, *optional*, defaults to `"all-reduce"`):
+            How the expert outputs get back to the tokens that need them. `"all-reduce"` runs the whole batch on
+            every rank and all-reduces the expert outputs. `"all-to-all"` sends each token to the rank that owns its
+            experts instead, so each rank trains on its own part of the batch and the parameters that are not
+            expert-parallel are sharded with FSDP2 across every rank. Anything but `"all-reduce"` requires
+            `enable_expert_parallel`.
         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`):
@@ -48,11 +60,18 @@ class DistributedConfig:
     tp_plan: dict[str, str] | Literal["auto"] | None = None
     enable_sequence_parallel: bool = False
     enable_expert_parallel: bool = False
+    experts_dispatch: str = "all-reduce"
     fsdp_size: int | None = None
     fsdp_cpu_offload: bool = False
     fsdp_mixed_precision: bool = False
     pp_size: int | None = None
 
+    @property
+    def dispatches_tokens(self) -> bool:
+        """Whether each rank routes and trains on its own part of the batch, which is every strategy but the
+        `all-reduce` default."""
+        return self.experts_dispatch != "all-reduce"
+
     def __post_init__(self):
         if self.tp_plan is None and self.tp_size is None and self.fsdp_size is None and self.pp_size is None:
             return
@@ -73,6 +92,20 @@ def __post_init__(self):
         elif self.tp_size is None:
             self.tp_size = 1
 
+        if self.experts_dispatch not in EXPERTS_DISPATCH_STRATEGIES:
+            raise ValueError(
+                f"Unknown `experts_dispatch={self.experts_dispatch!r}`, expected one of "
+                f"{sorted(EXPERTS_DISPATCH_STRATEGIES)}."
+            )
+        if self.dispatches_tokens and not self.enable_expert_parallel:
+            raise ValueError(f"`experts_dispatch={self.experts_dispatch!r}` requires `enable_expert_parallel=True`.")
+
+        if self.dispatches_tokens and self.pp_size > 1:
+            raise ValueError(
+                f"Combining `experts_dispatch={self.experts_dispatch!r}` with pipeline parallelism is not "
+                "supported yet."
+            )
+
         if self.fsdp_size > 1 and self.pp_size > 1:
             raise ValueError(
                 "Combining FSDP with pipeline parallelism is not supported yet. "

src/transformers/distributed/fsdp.py

@@ -20,7 +20,7 @@
 from ..utils import is_torch_available, is_torch_distributed_available, is_torch_greater_or_equal, logging, strtobool
 from ..utils.quantization_config import QuantizationMethod
 from .tensor_parallel import replace_layer_number_by_wildcard
-from .utils import _is_torch_distributed_initialized
+from .utils import _is_torch_distributed_initialized, is_dtensor
 
 
 if TYPE_CHECKING:
@@ -61,12 +61,16 @@ def is_fsdp_managed_module(module: nn.Module) -> bool:
     return isinstance(module, FullyShardedDataParallel)
 
 
-def _get_fsdp_policy_kwargs(distributed_config: DistributedConfig | None) -> dict[str, Any]:
+def _get_fsdp_policy_kwargs(
+    distributed_config: DistributedConfig | None, ignored_params: set[torch.nn.Parameter] | None = None
+) -> dict[str, Any]:
     """Build ``fully_shard`` policy kwargs from ``DistributedConfig`` runtime flags."""
+    fsdp_policy_kwargs = {}
+    if ignored_params:
+        fsdp_policy_kwargs["ignored_params"] = ignored_params
     if distributed_config is None:
-        return {}
+        return fsdp_policy_kwargs
 
-    fsdp_policy_kwargs = {}
     if distributed_config.fsdp_cpu_offload:
         fsdp_policy_kwargs["offload_policy"] = CPUOffloadPolicy()
     if distributed_config.fsdp_mixed_precision:
@@ -185,13 +189,19 @@ def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None)
 
 
 def apply_fully_sharded_data_parallelism(
-    model: nn.Module, fsdp_mesh: torch.distributed.device_mesh.DeviceMesh
+    model: nn.Module,
+    fsdp_mesh: torch.distributed.device_mesh.DeviceMesh,
+    expert_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
 ) -> nn.Module:
     """
     Apply FSDP2 (fully_shard) to a model.
 
     Torch availability, distributed initialization and the version requirement
     are asserted upstream by `initialize_distributed_mesh`.
+
+    With expert-parallel token dispatch `fsdp_mesh` spans the expert-parallel ranks, which the experts are already
+    sharded across: the experts are fully sharded across `expert_mesh` in their own group, and passed to FSDP2 as
+    `ignored_params` when `expert_mesh` is `None`.
     """
     fsdp_plan = dict(getattr(model, "_fsdp_plan", None) or {})
     if not fsdp_plan:
@@ -201,11 +211,26 @@ def apply_fully_sharded_data_parallelism(
         )
 
     distributed_config = getattr(model.config, "distributed_config", None)
-    fsdp_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config)
 
     adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, model)
     reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan)
 
+    ignored_params = None
+    if distributed_config is not None and distributed_config.dispatches_tokens:
+        # The DTensor parameters are the expert-parallel experts: `maybe_distribute_model` rewrote the expert
+        # parallel plan to shard only them.
+        expert_modules = [
+            module for module in model.modules() if any(is_dtensor(p) for p in module.parameters(recurse=False))
+        ]
+        if expert_mesh is not None:
+            expert_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config)
+            for module in expert_modules:
+                fully_shard(module, mesh=expert_mesh, reshard_after_forward=True, **expert_policy_kwargs)
+        else:
+            ignored_params = {p for module in expert_modules for p in module.parameters()}
+
+    fsdp_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config, ignored_params=ignored_params)
+
     for module_name, module in reshard_targets:
         fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=True, **fsdp_policy_kwargs)
         logger.debug(f"Applied fully_shard to {module_name} (reshard=True)")

src/transformers/distributed/mixin.py

@@ -20,7 +20,7 @@
 
 from ..utils import is_torch_greater_or_equal, logging
 from ..utils.hub import create_and_tag_model_card
-from .configuration_utils import DistributedConfig
+from .configuration_utils import EXPERTS_DISPATCH_STRATEGIES, DistributedConfig
 from .fsdp import apply_fully_sharded_data_parallelism, is_fsdp_managed_module
 from .pipeline_parallel import apply_pipeline_parallelism
 from .tensor_parallel import (
@@ -53,6 +53,7 @@ class DistributedMixin:
     _ep_plan: dict[str, str] | None = None
     _tp_size = None
     _fsdp_size = None
+    _expert_parallel_dispatch = False
     _pp_plan: dict[str, tuple[str, str]] | None = None
     _fsdp_plan: dict[str, str] | None = None
 
@@ -156,6 +157,8 @@ def prepare_distribute_model(
             raise ValueError("Tensor parallelism and `device_map` are mutually exclusive.")
         if distributed_config.fsdp_size > 1 and not is_torch_greater_or_equal("2.7"):
             raise OSError("FSDP2 requires `torch>=2.7` (distributed checkpoint save/load).")
+        if distributed_config.dispatches_tokens and not is_torch_greater_or_equal("2.7"):
+            raise OSError("Expert-parallel token dispatch requires `torch>=2.7`.")
 
         device_map, device_mesh = initialize_distributed_mesh(distributed_config)
 
@@ -174,6 +177,7 @@ def maybe_distribute_model(
             model._device_mesh = device_mesh
             model._tp_size = distributed_config.tp_size
             model._fsdp_size = distributed_config.fsdp_size
+            model._expert_parallel_dispatch = distributed_config.dispatches_tokens
 
             if distributed_config.pp_size > 1:
                 pp_mesh = device_mesh["pp"] if device_mesh.ndim > 1 else device_mesh
@@ -185,9 +189,40 @@ def maybe_distribute_model(
                 tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh
                 if isinstance(distributed_config.tp_plan, dict):
                     model.tp_plan = distributed_config.tp_plan
+                if distributed_config.dispatches_tokens:
+                    # Every rank trains on its own part of the batch, so only the experts can be sharded across the
+                    # group: the experts get the dispatch style, the router keeps its global ids and scores, and
+                    # whatever else the plan shards stays replicated, data-parallel like the rest of the trunk.
+                    # Replicated parameters inside the experts module keep their gradient all-reduce: FSDP2 treats
+                    # that module as expert-owned and does not reduce them, and each rank saw different tokens.
+                    kept = ("grouped_gemm", "moe_tp_experts", "replicated_with_grad_allreduce")
+                    replicated = sorted(
+                        name for name, style in model.tp_plan.items() if style not in ("ep_router", *kept)
+                    )
+                    if replicated:
+                        logger.warning(
+                            f"`experts_dispatch={distributed_config.experts_dispatch!r}` shards only the experts, "
+                            "so these expert parallel plan "
+                            f"entries are ignored and their modules stay replicated: {replicated}."
+                        )
+                    # `tp_plan` reads `_ep_plan` under expert parallelism, so that is the plan to rewrite.
+                    dispatch_style = EXPERTS_DISPATCH_STRATEGIES[distributed_config.experts_dispatch]
+                    model._ep_plan = {
+                        name: dispatch_style if style == "moe_tp_experts" else style
+                        for name, style in model.tp_plan.items()
+                        if style in kept
+                    }
                 model = apply_tensor_parallelism(model, tp_mesh)
 
-            if distributed_config.fsdp_size > 1:
+            if distributed_config.dispatches_tokens:
+                # Every expert-parallel rank trains on its own part of the batch, so the parameters outside the
+                # experts are data-parallel across the whole mesh: FSDP2 shards them across all of it and owns
+                # their gradient reduction. The experts stay sharded across `tp` and, if any, across `fsdp`.
+                flattened = "_".join(device_mesh.mesh_dim_names)
+                trunk_mesh = device_mesh[flattened] if device_mesh.ndim > 1 else device_mesh
+                expert_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else None
+                model = apply_fully_sharded_data_parallelism(model, trunk_mesh, expert_mesh=expert_mesh)
+            elif distributed_config.fsdp_size > 1:
                 fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh
                 model = apply_fully_sharded_data_parallelism(model, fsdp_mesh)
         return model
@@ -250,9 +285,9 @@ def gather_sharded_state_dict_for_save(
         if distributed_config is None:
             return state_dict
 
-        if distributed_config.fsdp_size > 1:
-            # Also covers the 2-D (fsdp, tp) mesh: every parameter is FSDP-managed, and the full
-            # state dict is only materialized on rank 0.
+        if distributed_config.fsdp_size > 1 or distributed_config.dispatches_tokens:
+            # Also covers the 2-D (fsdp, tp) mesh and token dispatch: every parameter is FSDP-managed, and
+            # the full state dict is only materialized on rank 0.
             if not _is_torch_distributed_initialized():
                 raise ValueError(
                     "Saving an FSDP-wrapped model requires torch.distributed to be initialized. "