[distributed] Default MoE ep_plans to token dispatch

docs/source/en/expert_parallelism.md

@@ -19,7 +19,7 @@ rendered properly in your Markdown viewer.
 
 ## DistributedConfig
 
-Enable expert parallelism with the [`DistributedConfig`] class and the `ep_size` argument. Most models route with masking and all-reduce, which requires `ep_size=tp_size` so every rank in an expert group receives the same tokens. Models whose plan uses [token dispatch](#token-dispatch), such as Qwen3 MoE, can set `ep_size` independently of `tp_size`.
+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 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.
 
 ```py
 import os
@@ -74,7 +74,7 @@ distributed_config = DistributedConfig(
 
 Providing a plan does not infer parallel sizes: set `tp_size` and `ep_size` explicitly.
 
-Qwen3 MoE defaults 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:
+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):
 
 ```py
 distributed_config = DistributedConfig(
@@ -155,7 +155,7 @@ These configurations each use eight GPUs:
 
 | Configuration | Result |
 | :--- | :--- |
-| `DistributedConfig(tp_size=4, fsdp_size=2, ep_size=4)` | Dispatch with TP groups of four, each slicing its batch in four (Qwen3 MoE default plan); unless you specify a ep_plan to use the legacy masked EP |
+| `DistributedConfig(tp_size=4, fsdp_size=2, ep_size=4)` | Dispatch with TP groups of four, each slicing its batch in four (default plan); unless you specify a ep_plan to use the legacy masked EP |
 | `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)` | Dispatch with every rank sharing one batch, sliced in eight, or masking and all-reduce for a masked plan. |

src/transformers/distributed/mixin.py

@@ -52,6 +52,7 @@ class DistributedMixin:
     """Distributed orchestration and save/load hooks for [`PreTrainedModel`]."""
 
     _device_mesh = None
+    _is_distributed_loading_by_transformers = False
     _mesh_manager: TransformersDeviceMesh | None = None
     _tp_plan: dict[str, str] | None = None
     _ep_plan: dict[str, str] | None = None
@@ -60,6 +61,14 @@ class DistributedMixin:
     _pp_plan: dict[str, tuple[str, str]] | None = None
     _fsdp_plan: dict[str, str] | None = None
 
+    @property
+    def is_distributed_loading_by_transformers(self) -> bool:
+        """
+        Whether `from_pretrained(distributed_config=...)` sharded the model at load time (each rank reads only its
+        shard through `DtensorShardOperation`), rather than a wrapper such as Accelerate sharding it afterwards.
+        """
+        return self._is_distributed_loading_by_transformers
+
     def init_parallel_plans(self) -> None:
         """Copy class-level plans onto the instance and merge config/children contributions."""
         model_cls = type(self)
@@ -192,6 +201,7 @@ def maybe_distribute_model(
         model.config.distributed_config = distributed_config
         model._mesh_manager = mesh_manager
         model._device_mesh = mesh_manager.get_mesh(("pp", "fsdp", "tp"))
+        model._is_distributed_loading_by_transformers = True
         model._tp_size = distributed_config.tp_size
         model._fsdp_size = distributed_config.fsdp_size
 

src/transformers/models/afmoe/configuration_afmoe.py

@@ -62,10 +62,9 @@ class AfmoeConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.router": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     vocab_size: int = 200192

src/transformers/models/axk1/configuration_axk1.py

@@ -69,10 +69,9 @@ class AXK1Config(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
     attribute_map = {
         "num_local_experts": "n_routed_experts",

src/transformers/models/axk2/configuration_axk2.py

@@ -82,10 +82,9 @@ class AXK2Config(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     attribute_map = {"num_local_experts": "n_routed_experts"}

src/transformers/models/cohere2_moe/configuration_cohere2_moe.py

@@ -85,10 +85,9 @@ class Cohere2MoeConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     vocab_size: int = 256000

src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py

@@ -238,10 +238,9 @@ class DeepseekOcr2TextConfig(PreTrainedConfig):
     mlp_bias: bool = False
     head_dim: int | None = None
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
     attribute_map = {
         "num_experts": "n_routed_experts",

src/transformers/models/deepseek_v2/configuration_deepseek_v2.py

@@ -95,10 +95,9 @@ class DeepseekV2Config(PreTrainedConfig):
     mlp_bias: bool = False
     head_dim: int | None = None
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
     attribute_map = {
         "num_experts": "n_routed_experts",

src/transformers/models/deepseek_v2/modular_deepseek_v2.py

@@ -86,10 +86,9 @@ class DeepseekV2Config(LlamaConfig):
         "layers.*.mlp.down_proj": "rowwise",
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     model_type = "deepseek_v2"

src/transformers/models/deepseek_v3/configuration_deepseek_v3.py

@@ -69,10 +69,9 @@ class DeepseekV3Config(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     attribute_map = {

src/transformers/models/deepseek_v32/configuration_deepseek_v32.py

@@ -77,10 +77,9 @@ class DeepseekV32Config(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     attribute_map = {"num_local_experts": "n_routed_experts"}

src/transformers/models/deepseek_v4/configuration_deepseek_v4.py

@@ -111,29 +111,15 @@ class DeepseekV4Config(PreTrainedConfig):
         "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
         "norm": (["hidden_states"], ["hidden_states"]),
     }
+    # Main attention and the shared MLP stay replicated. Shared-KV MQA and the CSA / HCA compressor
+    # broadcast one KV head via `repeat_kv`, so colwise `q_b_proj` would mismatch the rank-local
+    # head count; the shared MLP is too small to be worth sharding.
     base_model_ep_plan = {
-        # V4 ships EP only (no `base_model_tp_plan` — the runtime picks one plan or
-        # the other, never both, and V4 is MoE so EP is the only sensible config).
-        # MoE parallelism: route on the gate, run the routed experts as a grouped-GEMM
-        # kernel sharded along the expert axis, and wrap the experts module with
-        # `moe_tp_experts` so its output gets all-reduced across ranks. Same shape as
-        # gpt-oss. Main attention stays replicated: V4 is shared-KV MQA + a CSA / HCA
-        # compressor branch — both broadcast a single KV head across all attention
-        # heads via `repeat_kv`, so colwise-sharding `q_b_proj` would leave KV
-        # replicated and `repeat_kv` would no longer match the rank-local query head
-        # count. The shared MLP also stays replicated — it's small and not worth
-        # sharding. The Lightning Indexer is the one carve-out: its keys are
-        # replicated (own compressor at index_head_dim fed by replicated
-        # hidden_states), so head-sharding is well-formed; `q_b_proj` and the
-        # `scorer.weights_proj` go colwise, and the `scorer` output is all-reduced
-        # so every rank sees the same `index_scores` and picks the same top-k.
-        "layers.*.mlp.gate": "ep_router",
+        # V4 ships EP only (no `base_model_tp_plan`). Routed experts run as a grouped-GEMM kernel
+        # sharded along the expert axis, with all-to-all token dispatch (same as gpt-oss).
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
-        "layers.*.self_attn.compressor.indexer.q_b_proj": "colwise",
-        "layers.*.self_attn.compressor.indexer.scorer.weights_proj": "colwise",
-        "layers.*.self_attn.compressor.indexer.scorer": "all_reduce",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     vocab_size: int = 129280

src/transformers/models/dots1/configuration_dots1.py

@@ -71,10 +71,9 @@ class Dots1Config(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     attribute_map = {

src/transformers/models/dots1/modular_dots1.py

@@ -88,10 +88,9 @@ class Dots1Config(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     attribute_map = {

src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py

@@ -83,10 +83,9 @@ class Ernie4_5_MoeConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     vocab_size: int = 103424

src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py

@@ -103,10 +103,9 @@ class Ernie4_5_VLMoeTextConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     vocab_size: int = 103424

src/transformers/models/exaone_moe/configuration_exaone_moe.py

@@ -104,10 +104,9 @@ class ExaoneMoeConfig(PreTrainedConfig):
     layer_types: list[str] | None = None
 
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
     mlp_layer_types: list[str] | None = None
     first_k_dense_replace: int = 1

src/transformers/models/exaone_moe/modular_exaone_moe.py

@@ -80,10 +80,9 @@ class ExaoneMoeConfig(Exaone4Config):
     ```"""
 
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     vocab_size: int = 102400

src/transformers/models/flex_olmo/configuration_flex_olmo.py

@@ -64,10 +64,9 @@ class FlexOlmoConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     vocab_size: int = 100352

src/transformers/models/flex_olmo/modular_flex_olmo.py

@@ -74,10 +74,9 @@ class FlexOlmoConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     vocab_size: int = 100352

src/transformers/models/glm4_moe/configuration_glm4_moe.py

@@ -78,10 +78,9 @@ class Glm4MoeConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     attribute_map = {

src/transformers/models/glm4_moe/modular_glm4_moe.py

@@ -91,10 +91,9 @@ class Glm4MoeConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     attribute_map = {

src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py

@@ -69,10 +69,9 @@ class Glm4MoeLiteConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     attribute_map = {

src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py

@@ -77,10 +77,9 @@ class Glm4MoeLiteConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     attribute_map = {

src/transformers/models/glm4v_moe/configuration_glm4v_moe.py

@@ -67,10 +67,9 @@ class Glm4vMoeTextConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
     attribute_map = {
         "num_local_experts": "n_routed_experts",

src/transformers/models/glm5_next/configuration_glm5_next.py

@@ -90,10 +90,9 @@ class Glm5NextTextConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     attribute_map = {"num_local_experts": "n_routed_experts"}

src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py

@@ -82,10 +82,9 @@ class GlmMoeDsaConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     attribute_map = {"num_local_experts": "n_routed_experts"}

src/transformers/models/gpt_oss/configuration_gpt_oss.py

@@ -33,12 +33,11 @@ class GptOssConfig(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        "layers.*.mlp.router": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.gate_up_proj_bias": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj_bias": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.mlp.experts": "ep_dispatch_experts",
     }
 
     num_hidden_layers: int = 36