🚨 Remove the ep_router expert-parallel style

docs/source/en/expert_parallelism.md

@@ -19,8 +19,9 @@ rendered properly in your Markdown viewer.
 
 ## DistributedConfig
 
-Enable expert parallelism with the [`DistributedConfig`] class and the `ep_size` argument. Most MoE models default to [token dispatch](#token-dispatch), so `ep_size` can be set independently of `tp_size`. A few, such as Llama 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.
+Enable expert parallelism with the [`DistributedConfig`] class and the `ep_size` argument. MoE models shard their experts with `"ep_dispatch_experts"`. Its implementations pick how tokens reach the experts. [Token dispatch](#token-dispatch) lets `ep_size` differ from `tp_size`. A few EP implementations require `ep_size=tp_size`, e.g. `"all_reduce"` and `"deepgemm_megamoe"`, so every rank of an expert group holds the same tokens.
 
+Llama 4 has no EP plan: its experts take every token already scaled by its router score, so there is nothing to dispatch. Shard them with tensor parallelism (`tp_size`) instead.
 ```py
 import os
 
@@ -39,16 +40,11 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 ```
 
-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](./perf_infer_gpu_multi) 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.
+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](./perf_infer_gpu_multi) 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 the `"ep_dispatch_experts"` rule on the experts module selects how tokens reach them.
 
 `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:
-
-| rule | mechanism | layout |
-| :--- | :--- | :--- |
-| `"moe_tp_experts"` with `"ep_router"` on the router | masking and all-reduce: every rank runs its local experts on the whole batch, the router masks the others, and an all-reduce combines the outputs | `ep_size=tp_size` |
-| `"ep_dispatch_experts"` | whole experts per rank, reached through the implementation `ep_implementation` names, see [Choosing an implementation](#choosing-an-implementation) | `ep_size` a multiple of `tp_size` that divides `fsdp_size * tp_size` |
+`"ep_dispatch_experts"` shards whole experts per rank (`ep_size` a multiple of `tp_size` that divides `fsdp_size * tp_size`) and reaches them through the implementation `ep_implementation` names, see [Choosing an implementation](#choosing-an-implementation).
 
 ### Choosing an implementation
 
@@ -108,21 +104,12 @@ distributed_config = DistributedConfig(
 
 Providing a plan does not infer parallel sizes: set `tp_size` and `ep_size` explicitly.
 
-Most MoE models default to `"ep_dispatch_experts"`, which already uses `"all_reduce"` at inference when `ep_size=tp_size`. To use masking and all-reduce in training too, 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):
+With `ep_size=tp_size`, the default `"auto"` implementation runs `"all_reduce"` at inference and dispatches tokens in training. To mask and all-reduce in training too, set `ep_implementation="all_reduce"`:
 
 ```py
-distributed_config = DistributedConfig(
-    tp_size=4,
-    ep_size=4,
-    ep_plan={
-        "model.layers.*.mlp.gate": "ep_router",
-        "model.layers.*.mlp.experts": "moe_tp_experts",
-    },
-)
+distributed_config = DistributedConfig(tp_size=4, ep_size=4, ep_implementation="all_reduce")
 ```
 
-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
 
 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.
@@ -189,17 +176,17 @@ These configurations each use eight GPUs:
 
 | Configuration | Result |
 | :--- | :--- |
-| `DistributedConfig(tp_size=4, fsdp_size=2, ep_size=4)` | Each TP group of four shares one batch. Training dispatches tokens, each rank taking a quarter of the batch; inference uses `"all_reduce"`. A masked `ep_plan` uses masking and all-reduce throughout. |
+| `DistributedConfig(tp_size=4, fsdp_size=2, ep_size=4)` | Each TP group of four shares one batch. Training dispatches tokens, each rank taking a quarter of the batch; inference uses `"all_reduce"`; setting `ep_implementation="all_reduce"` uses it throughout. |
 | `DistributedConfig(tp_size=1, fsdp_size=8, ep_size=4)` | Dispatch with an independent batch on each rank, no slicing, and experts FSDP-sharded across pairs of ranks. |
 | `DistributedConfig(tp_size=2, fsdp_size=4, ep_size=4)` | Dispatch with a TP pair per batch, each pair slicing its batch in two; two batches per expert group. |
-| `DistributedConfig(tp_size=8, ep_size=8)` | Every rank shares one batch. Training dispatches tokens, each rank taking an eighth of the batch; inference uses `"all_reduce"`. A masked `ep_plan` uses masking and all-reduce throughout. |
+| `DistributedConfig(tp_size=8, ep_size=8)` | Every rank shares one batch. Training dispatches tokens, each rank taking an eighth of the batch; inference uses `"all_reduce"`; setting `ep_implementation="all_reduce"` uses it throughout. |
 
 > [!WARNING]
 > 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](./fsdp) 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"}``.
+Tensor and expert parallelism shard the weights across their own meshes, `tp` and `ep`, 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](./fsdp) on another mesh dimension with `fsdp_size`. With masking and all-reduce, keep `ep_size=tp_size` and pass `ep_implementation="all_reduce"`.
 
 ```py
 from transformers import AutoModelForCausalLM
@@ -209,10 +196,7 @@ 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",
-    },
+    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)
 ```

docs/source/en/perf_infer_gpu_multi.md

@@ -142,7 +142,6 @@ The table below describes each strategy.
 | `MoeIdentityParallel` | Pre-divides the input of a zero or identity expert by the mesh size, cancelling the all-reduce that `moe_tp_experts` applies downstream. |
 | `LocalParamsParallel` | Runs a module's forward with its own sharded parameters as local tensors. Use it on the owner of a bare `nn.Parameter`, such as attention sinks, that its own plan entry shards. |
 | `MoEParamShard` | Shards MoE expert weights on a given placement. Backs the `grouped_gemm` name, where `shards_expert_dim=True` also rewrites `module.num_experts` to the per-rank expert count. |
-| `EpRouterParallel` | Masks router scores for non-local experts and remaps global expert IDs to local ones, so each rank runs only the experts it owns. Requires `num_experts` to be divisible by the mesh size. |
 
 ### Packed strategies
 

src/transformers/distributed/configuration_utils.py

@@ -200,10 +200,6 @@ def _validate_resolved_ep_plan(self, ep_plan: dict[str, str]):
                 raise ValueError("Combining token dispatch with pipeline parallelism is not supported/tested yet.")
             if not is_torch_greater_or_equal("2.7"):
                 raise OSError("Expert-parallel token dispatch requires `torch>=2.7`.")
-        elif {"ep_router", "moe_tp_experts"}.issubset(ep_plan.values()) and self.ep_size != self.tp_size:
-            raise ValueError(
-                "All-reduce expert parallelism requires `ep_size=tp_size`, so every rank of an expert group sees the same tokens"
-            )
 
     @classmethod
     def from_dict(cls, config_dict: dict, **kwargs) -> "DistributedConfig":

src/transformers/distributed/expert_parallel.py

@@ -85,8 +85,8 @@ def ep_forward(*args, **kwargs):
 class MoEParamShard(ExpertParallelLayer):
     """Param-only EP style for MoE expert weights (``grouped_gemm``).
 
-    Shards dim 0 and updates module.num_experts to the per-rank local count so the
-    experts forward and ep_router sentinel agree.
+    Shards dim 0 and updates module.num_experts to the per-rank local count, so the experts forward skips the
+    expert ids the EP implementations mark as off-rank.
     """
 
     def __init__(self, placement, *, shards_expert_dim: bool = False):
@@ -114,77 +114,6 @@ def shard_param(self, module, param, mesh):
         )
 
 
-class EpRouterParallel(ExpertParallelLayer):
-    """Expert-parallel router: forward-only slicing of router outputs to local experts.
-
-    Expects the router to return `(router_logits, router_scores, router_indices, *extra)`.
-    `router_logits` and any trailing `extra` outputs are passed through unchanged.
-
-    The gate runs replicated on every rank and emits global expert IDs and scores.
-    Under EP each rank owns `num_experts // ep_size` experts, so this post-forward hook:
-
-    - zeroes scores for non-local experts
-    - remaps surviving global indices to local indices (`fmod` after masking non-local slots)
-    - sets dropped slots to sentinel `num_local_experts` (skipped by grouped_gemm experts forward)
-
-    Downstream `moe_tp_experts` allreduce-sums partial per-rank expert outputs.
-
-    Example: 4 tokens, top_k=4, 128 experts, EP=8 → num_local_experts=16.
-
-    Router output (identical on all ranks):
-        router_indices:
-            [ 52,  42, 119,  67],
-            [102,  89,  61,  40],
-            [ 82, 103,   4,  34],
-            [ 93,  23, 109,  11]
-
-    Owning rank (`index // 16`):
-            [ 3,  2,  7,  4],
-            [ 6,  5,  3,  2],
-            [ 5,  6,  0,  2],
-            [ 5,  1,  6,  0],
-
-    After slicing on rank 0 (owns experts 0-15):
-        router_indices:          router_scores (illustrative):
-            [16, 16, 16, 16],       [0.0, 0.0, 0.0, 0.0],
-            [16, 16, 16, 16],       [0.0, 0.0, 0.0, 0.0],
-            [16, 16,  4, 16],       [0.0, 0.0, 0.3, 0.0],
-            [16, 16, 16, 11],       [0.0, 0.0, 0.0, 0.1],
-
-    On rank 1 (owns experts 16-31), global expert 23 remaps to local index 7 via `23 % 16`.
-
-    Scores and indices stay paired element-wise in `(seq, top_k)` shape.
-    """
-
-    def transform_output_post_forward(self, module, output, mesh):
-        ep_rank, ep_size = mesh.get_local_rank(), mesh.size()
-        num_experts = getattr(module, "num_experts", None)
-        if num_experts is None:
-            num_experts = getattr(getattr(module, "config", None), "num_experts", None)
-        if num_experts is None:
-            raise AttributeError(
-                f"Router module {type(module).__name__} is missing `num_experts` and `config.num_experts`"
-            )
-        if num_experts % ep_size != 0:
-            raise ValueError(f"num_experts must be divisible by ep_size: {num_experts} % {ep_size} != 0")
-        num_local_experts = num_experts // ep_size
-
-        router_logits, router_scores, router_indices, *extra_outputs = output
-        # Each rank's score gradient covers only its local experts' slots; sum the per-rank partials
-        # before the mask (each slot has exactly one owning rank, so the sum is exact).
-        process_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp")
-        router_scores = _AllReduceBackward.apply(router_scores, process_group)
-        non_local_mask = (router_indices // num_local_experts) != ep_rank
-        router_scores = router_scores.masked_fill(non_local_mask, 0.0)
-        router_indices = router_indices.masked_fill(non_local_mask, -1)
-        if num_local_experts > 1:
-            router_indices = torch.fmod(router_indices, num_local_experts)
-        else:
-            router_indices = router_indices.masked_fill(router_indices > 0, 0).masked_fill(router_indices < 0, -1)
-        router_indices = router_indices.masked_fill(router_indices == -1, num_local_experts)
-        return router_logits, router_scores, router_indices, *extra_outputs
-
-
 class EPImplementation:
     """
     How the experts of an `ep_dispatch_experts` module get their tokens. One instance is built per experts module when
@@ -485,7 +414,6 @@ def ep_forward(hidden_states, top_k_index, top_k_weights):
 
 if is_torch_distributed_available():
     ALL_PARALLEL_STYLES.register("grouped_gemm", MoEParamShard(Shard(0), shards_expert_dim=True))
-    ALL_PARALLEL_STYLES.register("ep_router", EpRouterParallel())
     ALL_PARALLEL_STYLES.register("ep_dispatch_experts", EpDispatchExpertsParallel())
 
 

src/transformers/distributed/mixin.py

@@ -263,9 +263,15 @@ def _resolve_parallel_plans(
                     f"or set `{kind}_size=1`."
                 )
 
-        if "ep_dispatch_experts" in ep_plan.values() and "ep_router" in ep_plan.values():
+        if "ep_router" in ep_plan.values() or "moe_tp_experts" in ep_plan.values():
             raise ValueError(
-                "`ep_dispatch_experts` routes tokens itself; remove the `ep_router` rules from `ep_plan`."
+                "The `ep_router` style was removed, and `moe_tp_experts` is a tensor-parallel style that belongs in "
+                "`tp_plan`. Drop the router rule and shard the experts by expert with `ep_dispatch_experts` and "
+                "`grouped_gemm`, e.g. `ep_plan={'layers.*.mlp.experts': 'ep_dispatch_experts', "
+                "'layers.*.mlp.experts.gate_up_proj': 'grouped_gemm', 'layers.*.mlp.experts.down_proj': "
+                "'grouped_gemm'}` (most MoE models default to this plan, so `ep_plan` can usually be left out). "
+                '`DistributedConfig(ep_implementation="all_reduce")` gives the masking and all-reduce of the old '
+                "`ep_router` + `moe_tp_experts` plan."
             )
 
         # EP rules take precedence: drop TP rules on EP modules and their children.
@@ -324,10 +330,7 @@ def maybe_distribute_model(
             tp_mesh = mesh_manager.get_mesh("tp")
             ep_mesh = mesh_manager.get_mesh("ep")
 
-            if {"ep_router", "moe_tp_experts"}.issubset(ep_plan.values()):
-                # Legacy masked EP: the EP group is the TP group, every rank keeps every token.
-                model = apply_tensor_parallelism(model, tp_mesh, ep_plan)
-            elif "ep_dispatch_experts" in ep_plan.values():
+            if "ep_dispatch_experts" in ep_plan.values():
                 # EP + DP with tp_size >= 1: the ranks of a TP group share the same batch. If we want a specific token,
                 # we will have to slice the batch here in order to avoid computing tp_size times the same batch.
                 model = apply_expert_parallelism(

src/transformers/distributed/tensor_parallel.py

@@ -553,7 +553,7 @@ class MoeExpertsParallel(TensorParallelLayer):
     def should_use_local_tensors(self, module):
         return True
 
-    def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_parallel=False):
+    def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
         hidden_states, *routing_args = args
         tp_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp")
         if isinstance(hidden_states, DTensor):
@@ -564,46 +564,15 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_
             top_k_index, top_k_weights, *extra_args = routing_args
             if isinstance(top_k_weights, DTensor):
                 top_k_weights = top_k_weights.to_local()
-            if not is_expert_parallel:
-                top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group)
+            top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group)
             routing_args = [top_k_index, top_k_weights, *extra_args]
 
         return (hidden_states, *routing_args), kwargs
 
-    def install_forward(self, module, mesh, *, is_expert_parallel=False):
-        """Install the transforms but pass `is_expert_parallel` in the forward call."""
-        original_forward = module.forward
-        output_source = (
-            Partial()
-            if any(
-                isinstance(param, DTensor) and any(not placement.is_replicate() for placement in param.placements)
-                for param in module.parameters()
-            )
-            else Replicate()
-        )
-
-        def tp_forward(*args, **kwargs):
-            args, kwargs = self.transform_inputs_pre_forward(
-                module, args, kwargs, mesh, is_expert_parallel=is_expert_parallel
-            )
-            with self.context_around_forward(module, mesh):
-                output = original_forward(*args, **kwargs)
-            return self.transform_output_post_forward(module, output, mesh, source=output_source)
-
-        module.forward = tp_forward
-        return module
-
-    def transform_output_post_forward(self, module, output, mesh, source=None):
+    def transform_output_post_forward(self, module, output, mesh):
         if output is None:
             return None
-
-        has_sharded_parameters = any(
-            isinstance(param, DTensor) and any(not placement.is_replicate() for placement in param.placements)
-            for param in module.parameters()
-        )
-        if not has_sharded_parameters:
-            return output
-
+        # each rank ran its shard of every expert: sum the partial outputs
         process_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp")
         return _AllReduceForward.apply(output, process_group)
 
@@ -718,7 +687,6 @@ def gather_state_dict_for_save(
 # Lazy import because expert_parallel.py imports from this file (circular import otherwise).
 _MOVED_TO_EXPERT_PARALLEL = (
     "MoEParamShard",
-    "EpRouterParallel",
     "EpDispatchExpertsParallel",
     "apply_expert_parallelism",
 )