diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index 61d646c56f23..beda647e4d5d 100644 --- a/docs/source/en/expert_parallelism.md +++ b/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. | diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 584e2e5a6fb4..5de80438a678 100644 --- a/src/transformers/distributed/mixin.py +++ b/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 diff --git a/src/transformers/models/afmoe/configuration_afmoe.py b/src/transformers/models/afmoe/configuration_afmoe.py index df3980e51917..40083d09e68a 100644 --- a/src/transformers/models/afmoe/configuration_afmoe.py +++ b/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 diff --git a/src/transformers/models/axk1/configuration_axk1.py b/src/transformers/models/axk1/configuration_axk1.py index c654f1331be1..134ee88b01a6 100644 --- a/src/transformers/models/axk1/configuration_axk1.py +++ b/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", diff --git a/src/transformers/models/axk2/configuration_axk2.py b/src/transformers/models/axk2/configuration_axk2.py index 71478558e12d..a5f57ff65048 100644 --- a/src/transformers/models/axk2/configuration_axk2.py +++ b/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"} diff --git a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py index 1cdf869aa2a5..e94bed7ac32c 100644 --- a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py +++ b/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 diff --git a/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py index 2b44918d37ee..eec8d134ea4b 100644 --- a/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py +++ b/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", diff --git a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py index ee4bad80faa9..0b9093626489 100644 --- a/src/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/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", diff --git a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py b/src/transformers/models/deepseek_v2/modular_deepseek_v2.py index 98425bf98506..f31d8d4528ca 100644 --- a/src/transformers/models/deepseek_v2/modular_deepseek_v2.py +++ b/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" diff --git a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py b/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py index a2f2e4e31572..a030c1db5d03 100644 --- a/src/transformers/models/deepseek_v3/configuration_deepseek_v3.py +++ b/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 = { diff --git a/src/transformers/models/deepseek_v32/configuration_deepseek_v32.py b/src/transformers/models/deepseek_v32/configuration_deepseek_v32.py index b225cf5b722f..48eb2047edcb 100644 --- a/src/transformers/models/deepseek_v32/configuration_deepseek_v32.py +++ b/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"} diff --git a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py b/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py index 78e7f11291b1..e361b66ccc44 100644 --- a/src/transformers/models/deepseek_v4/configuration_deepseek_v4.py +++ b/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 diff --git a/src/transformers/models/dots1/configuration_dots1.py b/src/transformers/models/dots1/configuration_dots1.py index d4d14069e1ae..f04022e58df4 100644 --- a/src/transformers/models/dots1/configuration_dots1.py +++ b/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 = { diff --git a/src/transformers/models/dots1/modular_dots1.py b/src/transformers/models/dots1/modular_dots1.py index 3e3dcb8566c6..ea3a83ec51ed 100644 --- a/src/transformers/models/dots1/modular_dots1.py +++ b/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 = { diff --git a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py index 47c7a53dff84..142ff7a2b25c 100644 --- a/src/transformers/models/ernie4_5_moe/configuration_ernie4_5_moe.py +++ b/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 diff --git a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py index 762b910ee666..b920a2f42652 100644 --- a/src/transformers/models/ernie4_5_vl_moe/configuration_ernie4_5_vl_moe.py +++ b/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 diff --git a/src/transformers/models/exaone_moe/configuration_exaone_moe.py b/src/transformers/models/exaone_moe/configuration_exaone_moe.py index 50439746a568..9cdab1097f8e 100644 --- a/src/transformers/models/exaone_moe/configuration_exaone_moe.py +++ b/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 diff --git a/src/transformers/models/exaone_moe/modular_exaone_moe.py b/src/transformers/models/exaone_moe/modular_exaone_moe.py index 72a5c27279dd..92b92c2c34ff 100644 --- a/src/transformers/models/exaone_moe/modular_exaone_moe.py +++ b/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 diff --git a/src/transformers/models/flex_olmo/configuration_flex_olmo.py b/src/transformers/models/flex_olmo/configuration_flex_olmo.py index b4fff038c4a4..921ef2b2597c 100644 --- a/src/transformers/models/flex_olmo/configuration_flex_olmo.py +++ b/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 diff --git a/src/transformers/models/flex_olmo/modular_flex_olmo.py b/src/transformers/models/flex_olmo/modular_flex_olmo.py index 97a63d1a1d5f..956bea08670c 100644 --- a/src/transformers/models/flex_olmo/modular_flex_olmo.py +++ b/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 diff --git a/src/transformers/models/glm4_moe/configuration_glm4_moe.py b/src/transformers/models/glm4_moe/configuration_glm4_moe.py index b0ff8a259826..884066913209 100644 --- a/src/transformers/models/glm4_moe/configuration_glm4_moe.py +++ b/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 = { diff --git a/src/transformers/models/glm4_moe/modular_glm4_moe.py b/src/transformers/models/glm4_moe/modular_glm4_moe.py index 1dc0c7b840f7..7bef7d9ca2da 100644 --- a/src/transformers/models/glm4_moe/modular_glm4_moe.py +++ b/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 = { diff --git a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py index 3a563fd67120..b02df0b43b53 100644 --- a/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py +++ b/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 = { diff --git a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py b/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py index 4959464fd6ca..db05ee5e7b6f 100644 --- a/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py +++ b/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 = { diff --git a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py index 04087eb21334..a47812983c73 100644 --- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py +++ b/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", diff --git a/src/transformers/models/glm5_next/configuration_glm5_next.py b/src/transformers/models/glm5_next/configuration_glm5_next.py index b369b059e0aa..32f37518d808 100644 --- a/src/transformers/models/glm5_next/configuration_glm5_next.py +++ b/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"} diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index 7f640071adfe..7f6f457469ed 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/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"} diff --git a/src/transformers/models/gpt_oss/configuration_gpt_oss.py b/src/transformers/models/gpt_oss/configuration_gpt_oss.py index 47c029a5bca9..5aa25ed9ee1f 100644 --- a/src/transformers/models/gpt_oss/configuration_gpt_oss.py +++ b/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 diff --git a/src/transformers/models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py b/src/transformers/models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py index f9fd210f788d..f20b240507b8 100644 --- a/src/transformers/models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py +++ b/src/transformers/models/hunyuan_v1_moe/configuration_hunyuan_v1_moe.py @@ -34,10 +34,9 @@ class HunYuanMoEV1Config(PreTrainedConfig): model_type = "hunyuan_v1_moe" keys_to_ignore_at_inference = ["past_key_values"] 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_per_tok": "moe_topk", diff --git a/src/transformers/models/hy_v3/configuration_hy_v3.py b/src/transformers/models/hy_v3/configuration_hy_v3.py index 4399a0f4d546..6266f301095c 100644 --- a/src/transformers/models/hy_v3/configuration_hy_v3.py +++ b/src/transformers/models/hy_v3/configuration_hy_v3.py @@ -73,10 +73,9 @@ class HYV3Config(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 = 120832 diff --git a/src/transformers/models/hy_v3/modular_hy_v3.py b/src/transformers/models/hy_v3/modular_hy_v3.py index c24758bb5100..7aa409123b40 100644 --- a/src/transformers/models/hy_v3/modular_hy_v3.py +++ b/src/transformers/models/hy_v3/modular_hy_v3.py @@ -98,10 +98,9 @@ class HYV3Config(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 = 120832 diff --git a/src/transformers/models/kimi_linear/configuration_kimi_linear.py b/src/transformers/models/kimi_linear/configuration_kimi_linear.py index 3fa8c33fdce9..90fbf51356ec 100644 --- a/src/transformers/models/kimi_linear/configuration_kimi_linear.py +++ b/src/transformers/models/kimi_linear/configuration_kimi_linear.py @@ -60,10 +60,9 @@ class KimiLinearConfig(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 = { "max_position_embeddings": "model_max_length", diff --git a/src/transformers/models/laguna/configuration_laguna.py b/src/transformers/models/laguna/configuration_laguna.py index e10b6fbd4d80..024937999e4b 100644 --- a/src/transformers/models/laguna/configuration_laguna.py +++ b/src/transformers/models/laguna/configuration_laguna.py @@ -84,10 +84,9 @@ class LagunaConfig(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 diff --git a/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py b/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py index d4cf1ca8f77c..0e09fdd607cc 100644 --- a/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py +++ b/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py @@ -49,10 +49,9 @@ class Lfm2MoeConfig(PreTrainedConfig): model_type = "lfm2_moe" keys_to_ignore_at_inference = ["past_key_values"] base_model_ep_plan = { - "layers.*.feed_forward.gate": "ep_router", "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm", "layers.*.feed_forward.experts.down_proj": "grouped_gemm", - "layers.*.feed_forward.experts": "moe_tp_experts", + "layers.*.feed_forward.experts": "ep_dispatch_experts", } default_theta = 1000000.0 diff --git a/src/transformers/models/minimax/configuration_minimax.py b/src/transformers/models/minimax/configuration_minimax.py index 709139a584e5..a7fd614a1d62 100644 --- a/src/transformers/models/minimax/configuration_minimax.py +++ b/src/transformers/models/minimax/configuration_minimax.py @@ -76,10 +76,9 @@ class MiniMaxConfig(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_experts": "num_local_experts"} diff --git a/src/transformers/models/minimax/modular_minimax.py b/src/transformers/models/minimax/modular_minimax.py index 58b39e59911e..37c41a63c126 100644 --- a/src/transformers/models/minimax/modular_minimax.py +++ b/src/transformers/models/minimax/modular_minimax.py @@ -103,10 +103,9 @@ class MiniMaxConfig(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_experts": "num_local_experts"} diff --git a/src/transformers/models/minimax_m2/configuration_minimax_m2.py b/src/transformers/models/minimax_m2/configuration_minimax_m2.py index 6e4fcc421b88..3e52eb9f736e 100644 --- a/src/transformers/models/minimax_m2/configuration_minimax_m2.py +++ b/src/transformers/models/minimax_m2/configuration_minimax_m2.py @@ -62,10 +62,9 @@ class MiniMaxM2Config(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 = { diff --git a/src/transformers/models/minimax_m2/modular_minimax_m2.py b/src/transformers/models/minimax_m2/modular_minimax_m2.py index 0722185c1c18..57090e31a927 100644 --- a/src/transformers/models/minimax_m2/modular_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modular_minimax_m2.py @@ -81,10 +81,9 @@ class MiniMaxM2Config(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 = { diff --git a/src/transformers/models/minimax_m3_vl/configuration_minimax_m3_vl.py b/src/transformers/models/minimax_m3_vl/configuration_minimax_m3_vl.py index 536d2c87d0fa..956785f69a06 100644 --- a/src/transformers/models/minimax_m3_vl/configuration_minimax_m3_vl.py +++ b/src/transformers/models/minimax_m3_vl/configuration_minimax_m3_vl.py @@ -70,10 +70,9 @@ class MiniMaxM3VLTextConfig(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 = { diff --git a/src/transformers/models/mistral4/configuration_mistral4.py b/src/transformers/models/mistral4/configuration_mistral4.py index d537cde40436..eb3c2f3f787d 100644 --- a/src/transformers/models/mistral4/configuration_mistral4.py +++ b/src/transformers/models/mistral4/configuration_mistral4.py @@ -63,10 +63,9 @@ class Mistral4Config(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 = { diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py index f0705cbf0b9c..c1158852cf55 100644 --- a/src/transformers/models/mixtral/configuration_mixtral.py +++ b/src/transformers/models/mixtral/configuration_mixtral.py @@ -57,10 +57,9 @@ class MixtralConfig(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_experts": "num_local_experts"} diff --git a/src/transformers/models/olmoe/configuration_olmoe.py b/src/transformers/models/olmoe/configuration_olmoe.py index 1bdb43a14332..699703b0d94d 100644 --- a/src/transformers/models/olmoe/configuration_olmoe.py +++ b/src/transformers/models/olmoe/configuration_olmoe.py @@ -55,10 +55,9 @@ class OlmoeConfig(PreTrainedConfig): "layers.*.mlp.experts": "moe_tp_experts", } 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 = 50304 diff --git a/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py index e7aaefde4bca..bbd1b002895d 100644 --- a/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py +++ b/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py @@ -57,12 +57,11 @@ class OpenAIPrivacyFilterConfig(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 = 8 num_local_experts: int = 128 diff --git a/src/transformers/models/phimoe/configuration_phimoe.py b/src/transformers/models/phimoe/configuration_phimoe.py index 37673ceca5a2..89afec64ce3d 100644 --- a/src/transformers/models/phimoe/configuration_phimoe.py +++ b/src/transformers/models/phimoe/configuration_phimoe.py @@ -47,10 +47,9 @@ class PhimoeConfig(PreTrainedConfig): model_type = "phimoe" keys_to_ignore_at_inference = ["past_key_values"] 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", } default_theta = 1000000.0 diff --git a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py index 318eedbe203d..6f034ef7dde6 100644 --- a/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/configuration_qwen2_moe.py @@ -68,10 +68,9 @@ class Qwen2MoeConfig(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 = 151936 diff --git a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py index 89283522a3dd..0ae57e89f9f5 100644 --- a/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/configuration_qwen3_5_moe.py @@ -81,10 +81,9 @@ class Qwen3_5MoeTextConfig(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 = 248320 diff --git a/src/transformers/models/qwen3_next/configuration_qwen3_next.py b/src/transformers/models/qwen3_next/configuration_qwen3_next.py index 076ccc380ee6..bcba1a251bea 100644 --- a/src/transformers/models/qwen3_next/configuration_qwen3_next.py +++ b/src/transformers/models/qwen3_next/configuration_qwen3_next.py @@ -84,10 +84,9 @@ class Qwen3NextConfig(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 = 151936 diff --git a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py index 41d44aedb7d5..6a3c3b87ee7f 100644 --- a/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py @@ -370,10 +370,9 @@ class Qwen3OmniMoeTalkerTextConfig(PreTrainedConfig): "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py index c8ebaafcc6d4..e10bb164d1ac 100644 --- a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py @@ -403,10 +403,9 @@ def __post_init__(self, **kwargs): @strict class Qwen3OmniMoeTalkerTextConfig(Qwen3MoeConfig): 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 = 3072 diff --git a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py index 26b4793fb73c..ede38489950f 100644 --- a/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/configuration_qwen3_vl_moe.py @@ -65,10 +65,9 @@ class Qwen3VLMoeTextConfig(PreTrainedConfig): "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py index 286439a60521..0902fd903091 100644 --- a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py @@ -92,10 +92,9 @@ class Qwen3VLMoeTextConfig(Qwen3MoeConfig): "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", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py b/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py index 69a37377a2cb..fcaedba75175 100644 --- a/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py +++ b/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py @@ -100,10 +100,9 @@ class Qwen4ExpTextConfig(PreTrainedConfig): } base_model_pp_plan = 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", } vocab_size: int = 248320 diff --git a/src/transformers/models/solar_open/configuration_solar_open.py b/src/transformers/models/solar_open/configuration_solar_open.py index 4a6e355b7b71..5d7b88efc59b 100644 --- a/src/transformers/models/solar_open/configuration_solar_open.py +++ b/src/transformers/models/solar_open/configuration_solar_open.py @@ -52,10 +52,9 @@ class SolarOpenConfig(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", diff --git a/src/transformers/models/step3p7/configuration_step3p7.py b/src/transformers/models/step3p7/configuration_step3p7.py index 001b62676543..680cad527e4e 100644 --- a/src/transformers/models/step3p7/configuration_step3p7.py +++ b/src/transformers/models/step3p7/configuration_step3p7.py @@ -128,10 +128,9 @@ class Step3p7TextConfig(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", diff --git a/src/transformers/models/youtu/configuration_youtu.py b/src/transformers/models/youtu/configuration_youtu.py index ffe0fa4fc355..6d9f2cef1f96 100644 --- a/src/transformers/models/youtu/configuration_youtu.py +++ b/src/transformers/models/youtu/configuration_youtu.py @@ -60,12 +60,6 @@ class YoutuConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "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", - } attribute_map = {} vocab_size: int = 128256 diff --git a/src/transformers/models/youtu/modular_youtu.py b/src/transformers/models/youtu/modular_youtu.py index 77e8913b3e5c..c5ba1f03da53 100644 --- a/src/transformers/models/youtu/modular_youtu.py +++ b/src/transformers/models/youtu/modular_youtu.py @@ -93,6 +93,7 @@ class YoutuConfig(DeepseekV3Config): pretraining_tp = AttributeError() moe_intermediate_size = AttributeError() num_mtp_layers = AttributeError() + base_model_ep_plan = AttributeError() def __post_init__(self, **kwargs): if self.initializer_range is None: diff --git a/src/transformers/models/zaya/configuration_zaya.py b/src/transformers/models/zaya/configuration_zaya.py index 63fe0c5b7ce4..aaa26896564f 100644 --- a/src/transformers/models/zaya/configuration_zaya.py +++ b/src/transformers/models/zaya/configuration_zaya.py @@ -53,10 +53,9 @@ class ZayaConfig(PreTrainedConfig): model_type = "zaya" keys_to_ignore_at_inference = ["past_key_values"] 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 = 262272 diff --git a/src/transformers/models/zaya/modular_zaya.py b/src/transformers/models/zaya/modular_zaya.py index 100b551c27db..c372b77c640f 100644 --- a/src/transformers/models/zaya/modular_zaya.py +++ b/src/transformers/models/zaya/modular_zaya.py @@ -71,6 +71,11 @@ class ZayaConfig(LagunaConfig): """ model_type = "zaya" + base_model_ep_plan = { + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "ep_dispatch_experts", + } vocab_size: int = 262272 moe_intermediate_size: int = 2048 diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index bfc759c5692c..c1420244524f 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -465,6 +465,10 @@ def __init__( elif len(devices) == 1: self.is_model_parallel = self.args.device != torch.device(devices[0]) + # Sharded at load time by `from_pretrained(distributed_config=...)`, whatever the parallelism: the model owns + # its placement and gradient reduction, so Accelerate must not wrap or shard it again. + self.is_distributed_loading_by_transformers = getattr(model, "is_distributed_loading_by_transformers", False) + self.is_fsdp_xla_enabled = args.fsdp and args.fsdp_config.get("xla", False) if args.fsdp: if self.is_deepspeed_enabled: @@ -486,7 +490,7 @@ def __init__( or is_sagemaker_mp_enabled() # Sharded at load time (`DistributedConfig`): the model manages its own placement, and # `.to()` on FSDP2-managed (possibly CPU-offloaded) parameters raises in `_apply`. - or getattr(model, "_device_mesh", None) is not None + or self.is_distributed_loading_by_transformers ): self.place_model_on_device = False else: @@ -627,7 +631,7 @@ def __init__( # Resolved lazily at the first gradient clip; see `_has_mixed_mesh_grads`. self._mixed_mesh_grads: bool | None = None if ( - getattr(model, "_device_mesh", None) is not None + self.is_distributed_loading_by_transformers and args.save_strategy != SaveStrategy.NO and not args.save_only_model ): @@ -1723,7 +1727,13 @@ def _prepare_for_training(self, max_steps, train_dataloader, resume_from_checkpo use_accelerator_prepare = model is self.model # prepare using `accelerator` prepare - if use_accelerator_prepare: + if self.is_distributed_loading_by_transformers: + # The model already owns placement and gradient reduction. Prepare autocast and compilation only, + # without asking Accelerate to wrap the DTensor parameters in DDP or to shard them again. + model = self.accelerator.prepare_model(model, device_placement=False, evaluation_mode=True) + self.optimizer = self.accelerator.prepare(self.optimizer) + self._sync_replicated_trainable_parameters(model) + elif use_accelerator_prepare: if delay_optimizer_creation: # TODO: check if we can move this somewhere else if self.is_fsdp_enabled and _is_peft_model(self.model): @@ -2174,11 +2184,9 @@ def compute_loss( and (self.model_accepts_loss_kwargs or self.compute_loss_func) and num_items_in_batch is not None ): - # TP and EP-as-TP ranks see replicated batches; `num_processes` over-counts - # them by `tp_size`. Mirror the divisor used in `_get_num_items_in_batch`. - loss_scale = self.accelerator.num_processes - if (pc := getattr(self.accelerator, "parallelism_config", None)) is not None: - loss_scale //= pc.tp_size + # TP ranks (and the expert-parallel ranks sharing their batch) see replicated batches; `num_processes` + # over-counts them by `tp_size`. Mirror the divisor used in `_get_num_items_in_batch`. + loss_scale = self.accelerator.num_processes // self.get_tp_size() loss *= loss_scale if self.args.n_gpu <= 1 else self.args.n_gpu return (loss, outputs) if return_outputs else loss @@ -2327,8 +2335,9 @@ def _get_num_items_in_batch(self, batch_samples: list, device: torch.device) -> # In the DataParallel case, convert the scalar tensor into a 2-dim tensor with the same value repeated num_items_in_batch = num_items_in_batch.unsqueeze(0).expand(self.args.n_gpu, -1) # Divide by number of devices with the same batch - if pc := getattr(self.accelerator, "parallelism_config", None): - num_items_in_batch = num_items_in_batch // pc.non_data_parallel_size + num_items_in_batch = num_items_in_batch // ( + self.get_tp_size() * self.get_cp_size() * self.get_sp_size() + ) return num_items_in_batch @@ -2573,9 +2582,43 @@ def get_tp_size(self) -> int: if self.is_deepspeed_enabled and (deepspeed_config := getattr(self.args, "hf_deepspeed_config", None)): return deepspeed_config.config.get("tensor_parallel", {}).get("autotp_size", 1) - # 3. Default fallback + # 3. Fall back to accelerate, for tensor parallelism configured outside `DistributedConfig` + if (pc := getattr(self.accelerator, "parallelism_config", None)) is not None: + return pc.tp_size + + # 4. Default fallback return 1 + def _sync_replicated_trainable_parameters(self, model: nn.Module) -> None: + """ + Keep the trainable parameters that FSDP2 does not manage identical across ranks. + + A parameter added after `fully_shard` (a PEFT adapter attached to a model sharded at load time) stays a plain + replicated tensor next to the DTensor base weights. FSDP2 reduce-scatters only what it sharded and the DDP + wrap is skipped for such a model, so each rank would init its own copy and train it on its own batch. + Broadcast these parameters from rank 0, then average their gradient over all ranks at the end of each + accumulation window. + """ + if not dist.is_available() or not dist.is_initialized() or dist.get_world_size() == 1: + return + + from torch.distributed.tensor import DTensor + + def average_gradient(param): + # Averaging the accumulated micro-batch gradients once gives the same result as after every backward. + if self.accelerator.sync_gradients: + dist.all_reduce(param.grad, op=dist.ReduceOp.AVG) + + for param in model.parameters(): + if not param.requires_grad or isinstance(param.data, DTensor): + continue + with torch.no_grad(): + dist.broadcast(param.data, src=0) + # `train()` can run more than once on the same model: register the hook only once. + if not getattr(param, "_replicated_grad_hook_registered", False): + param.register_post_accumulate_grad_hook(average_gradient) + param._replicated_grad_hook_registered = True + def _wrap_model(self, model: nn.Module, training: bool = True, dataloader: DataLoader | None = None) -> nn.Module: """Wrap `model` for distributed training if needed (DDP, FSDP, SageMaker, etc.).""" # train/eval could be run multiple-times - if already wrapped, don't re-wrap it again @@ -4044,7 +4087,7 @@ def save_model(self, output_dir: str | None = None, _internal_call: bool = False remove_dummy_checkpoint(self.args.should_save, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME]) self.model_wrapped.save_checkpoint(output_dir) - elif getattr(self.model, "_device_mesh", None) is not None and not _is_peft_model(self.model): + elif self.is_distributed_loading_by_transformers and not _is_peft_model(self.model): # Sharded at load time (`DistributedConfig`): gathering the weights inside `save_pretrained` # is collective, so every rank saves; only the main process writes, the others leave at the # closing barrier. (PEFT models fall through to the adapter-only save below.)