🚨 Separate EP logic into its own file (#49679)

* the branch is now pointing to main and not to unify-resolve-plan

* Refactor expert parallelism registration and improve documentation

- Consolidated the registration of parallel styles into the `ParallelInterface` class for better organization.
- Updated the `expert_parallel.py` and `tensor_parallel.py` files to streamline the registration process of expert parallel strategies.
- Enhanced documentation for clarity regarding the usage of tensor parallel strategies and their registration.
- Removed redundant code related to `MoeExpertsParallel` and `MoeIdentityParallel` classes, which are now properly registered in the new structure.

* comment

* some models missed the tests EP mixin

* fix copyright

Author: 3outeille

.circleci/create_circleci_config.py

@@ -306,7 +306,7 @@ def job_name(self):
 torch_job = CircleCIJob(
     "torch",
     docker_image=[{"image": "huggingface/transformers-torch-light"}],
-    marker="not (generate or is_training_test or is_tensor_parallel_test or is_fsdp_test)",
+    marker="not (generate or is_training_test or is_tensor_parallel_test or is_expert_parallel_test or is_fsdp_test)",
     parallelism=6,
 )
 
@@ -388,6 +388,15 @@ def job_name(self):
     parallelism=6,
 )
 
+expert_parallel_ci_job = CircleCIJob(
+    "expert_parallel_ci",
+    additional_env={"RUN_EXPERT_PARALLEL_TESTS": True},
+    docker_image=[{"image": "huggingface/transformers-torch-light"}],
+    install_steps=["uv pip install ."],
+    marker="is_expert_parallel_test",
+    parallelism=6,
+)
+
 fsdp_ci_job = CircleCIJob(
     "fsdp_ci",
     additional_env={"RUN_FSDP_TESTS": True},
@@ -434,9 +443,10 @@ def job_name(self):
 DOC_TESTS = [doc_test_job]
 TRAINING_CI_TESTS = [training_ci_job]
 TENSOR_PARALLEL_CI_TESTS = [tensor_parallel_ci_job]
+EXPERT_PARALLEL_CI_TESTS = [expert_parallel_ci_job]
 FSDP_CI_TESTS = [fsdp_ci_job]
 PEFT_INTEGRATION_TESTS = [peft_integration_job]
-ALL_TESTS = REGULAR_TESTS + PIPELINE_TESTS + REPO_UTIL_TESTS + DOC_TESTS + [custom_tokenizers_job] + [exotic_models_job] + TRAINING_CI_TESTS + TENSOR_PARALLEL_CI_TESTS + FSDP_CI_TESTS + PEFT_INTEGRATION_TESTS  # fmt: skip
+ALL_TESTS = REGULAR_TESTS + PIPELINE_TESTS + REPO_UTIL_TESTS + DOC_TESTS + [custom_tokenizers_job] + [exotic_models_job] + TRAINING_CI_TESTS + TENSOR_PARALLEL_CI_TESTS + EXPERT_PARALLEL_CI_TESTS + FSDP_CI_TESTS + PEFT_INTEGRATION_TESTS  # fmt: skip
 
 
 def create_circleci_config(folder=None):

docs/source/en/perf_infer_gpu_multi.md

@@ -101,35 +101,28 @@ print(model.tp_plan)
 
 ## Partitioning strategies
 
-The `ParallelInterface` class maps each strategy name you can use in a `tp_plan` to a configured strategy instance. You don't interact with it directly to shard a model, but it's the authoritative list of available names.
+`ALL_PARALLEL_STYLES`, a `ParallelInterface` defined in [distributed/parallel_utils.py](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/parallel_utils.py), maps each strategy name you can use in a plan to a configured strategy instance. You don't interact with it directly to shard a model, but it's the authoritative list of available names. [distributed/tensor_parallel.py](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/tensor_parallel.py) registers the tensor parallel strategies into it:
 
 ```py
-class ParallelInterface(GeneralInterface):
-    _global_mapping = {
-        "embedding_rowwise": EmbeddingParallel(shard_dim=0),
-        "embedding_colwise": EmbeddingParallel(shard_dim=1),
-        "colwise_gather_output": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
-        "colwise_rep": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
-        "colwise": ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)),
-        "rowwise": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()),
-        "rowwise_split_input": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
-        "rowwise_rep": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
-        "packed_colwise": PackedColwiseParallel(),
-        "packed_rowwise": PackedRowwiseParallel(),
-        "sequence_parallel": SequenceParallel(use_local_output=True),
-        "grouped_gemm": MoEParamShard(Shard(0), shards_expert_dim=True),
-        "ep_router": EpRouterParallel(),
-        "megamoe_router": RouterParallelMegaMoe(),
-        "moe_tp_experts": MoeExpertsParallel(),
-        "megamoe_experts": MoeTensorParalellMegaMoeExperts(),
-        "moe_identity_expert": MoeIdentityParallel(),
-        "replicated_with_grad_allreduce": ReplicatedWithGradAllReduce(),
-        "mla_kv_a_proj": MlaKvAProjParallel(),
-        "all_reduce": AllReduceParallel(),
-    }
+ALL_PARALLEL_STYLES.register("embedding_rowwise", EmbeddingParallel(shard_dim=0))
+ALL_PARALLEL_STYLES.register("embedding_colwise", EmbeddingParallel(shard_dim=1))
+ALL_PARALLEL_STYLES.register("colwise_gather_output", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
+ALL_PARALLEL_STYLES.register("colwise_rep", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
+ALL_PARALLEL_STYLES.register("colwise", ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)))
+ALL_PARALLEL_STYLES.register("rowwise", RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()))
+ALL_PARALLEL_STYLES.register("rowwise_split_input", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
+ALL_PARALLEL_STYLES.register("rowwise_rep", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
+ALL_PARALLEL_STYLES.register("packed_colwise", PackedColwiseParallel())
+ALL_PARALLEL_STYLES.register("packed_rowwise", PackedRowwiseParallel())
+ALL_PARALLEL_STYLES.register("sequence_parallel", SequenceParallel(use_local_output=True))
+ALL_PARALLEL_STYLES.register("replicated_with_grad_allreduce", ReplicatedWithGradAllReduce())
+ALL_PARALLEL_STYLES.register("mla_kv_a_proj", MlaKvAProjParallel())
+ALL_PARALLEL_STYLES.register("all_reduce", AllReduceParallel())
+ALL_PARALLEL_STYLES.register("moe_tp_experts", MoeExpertsParallel())
+ALL_PARALLEL_STYLES.register("moe_identity_expert", MoeIdentityParallel())
 ```
 
-Every strategy is a subclass of `TensorParallelLayer` in [distributed/tensor_parallel.py](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/tensor_parallel.py). The two you'll reach for most, `ColwiseParallel` and `RowwiseParallel`, take `input_layouts` and `output_layouts` placements, which is how one class covers several plan names. `colwise` leaves its output sharded on the last dim for a following `rowwise` layer to consume, while `colwise_gather_output` all-gathers it back to a full tensor.
+Every tensor parallel strategy is a subclass of `TensorParallelLayer` in that file. The expert parallel strategies, such as `ep_dispatch_experts`, live in [distributed/expert_parallel.py](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/expert_parallel.py), subclass `ExpertParallelLayer`, which has the same hooks, and register into the same `ALL_PARALLEL_STYLES`. The two you'll reach for most, `ColwiseParallel` and `RowwiseParallel`, take `input_layouts` and `output_layouts` placements, which is how one class covers several plan names. `colwise` leaves its output sharded on the last dim for a following `rowwise` layer to consume, while `colwise_gather_output` all-gathers it back to a full tensor.
 
 The table below describes each strategy.
 
@@ -144,12 +137,12 @@ The table below describes each strategy.
 | `ReplicatedWithGradAllReduce` | Replicates a parameter but all-reduces its gradient. Needed for norms that sit between a column-wise and a row-wise layer and normalize along a sharded axis, where each rank only sees its own heads. |
 | `AllReduceParallel` | All-reduces a module's `Partial()` forward output to `Replicate()`. Use it as a sync point for a module whose compute ends in a partial sum. |
 | `MlaKvAProjParallel` | Splits the `kv_a_proj_with_mqa` output of DeepSeek-V2 style MLA attention and all-reduces the gradient of the RoPE half, which bypasses `kv_b_proj` and would otherwise keep a partial gradient. Requires `qk_rope_head_dim` in the model config. |
+| `MoeExpertsParallel` | Tensor parallel MoE experts. All-reduces the expert output forward and adds the backward all-reduces for hidden states and routing weights. |
+| `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. |
 | `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. |
 | `RouterParallelMegaMoe` | Router variant for DeepGEMM Mega MoE, which dispatches experts inside the kernel which wants the router output untouched. |
-| `MoeExpertsParallel` | Tensor parallel MoE experts. All-reduces the expert output forward and adds the backward all-reduces for hidden states and routing weights. |
 | `MoeTensorParalellMegaMoeExperts` | Inference-only experts layer for DeepGEMM Mega MoE. Skips the gradient syncs and passes the process group into the module so the kernel can set up its shared buffers on the first forward. |
-| `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. |
 
 ### Packed strategies
 
@@ -242,9 +235,9 @@ The example below walks through a trimmed version of `ColwiseParallel`.
     import torch
 
     from transformers import AutoModelForCausalLM, DistributedConfig
-    from transformers.distributed.tensor_parallel import ParallelInterface
+    from transformers.distributed.parallel_utils import ALL_PARALLEL_STYLES
 
-    ParallelInterface.register("colwise_custom", ColwiseParallel())
+    ALL_PARALLEL_STYLES.register("colwise_custom", ColwiseParallel())
     tp_plan = {
         "model.layers.*.self_attn.q_proj": "colwise_custom",
         ...

pyproject.toml

@@ -93,6 +93,7 @@ markers = [
     "generate: marks tests that use the GenerationTesterMixin",
     "is_training_test: marks tests that use the TrainingTesterMixin (deselect with '-m \"not is_training_test\"')",
     "is_tensor_parallel_test: marks tests that use the TensorParallelTesterMixin (deselect with '-m \"not is_tensor_parallel_test\"')",
+    "is_expert_parallel_test: marks tests that use the ExpertParallelTesterMixin (deselect with '-m \"not is_expert_parallel_test\"')",
     "is_fsdp_test: marks tests that use the FSDPTesterMixin (deselect with '-m \"not is_fsdp_test\"')",
 ]
 log_cli = 1