🚨 Separate EP logic into its own file

@@ -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):

@@ -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",         ...

@@ -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

@@ -0,0 +1,408 @@+# Copyright 2026 The HuggingFace Team. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.+"""Expert parallelism (EP) for MoE models: the EP styles and the plan application."""++from __future__ import annotations++import contextlib+from collections.abc import Callable+from typing import TYPE_CHECKING++from ..utils import logging+from ..utils.import_utils import is_torch_available, is_torch_distributed_available+from .parallel_utils import ALL_PARALLEL_STYLES, _use_local_dtensor_params+from .tensor_parallel import _get_parameter_plan+++if TYPE_CHECKING:+    from torch import nn+    from torch.distributed.device_mesh import DeviceMesh+++logger = logging.get_logger(__name__)++if is_torch_available():+    import torch++if is_torch_distributed_available():+    from torch.distributed._functional_collectives import all_to_all_single_autograd+    from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor++    from .parallel_utils import _AllReduceBackward, _WaitTensor+++class ExpertParallelLayer:+    """Base class of the expert-parallel styles, with the same hooks as `TensorParallelLayer`."""++    def should_use_local_tensors(self, module):+        """Whether this module's forward requires local inputs and parameters."""+        return False++    def validate_param(self, module, param, mesh, parameter_name=None):+        """Validate a parameter before applying this EP style."""+        pass++    def shard_param(self, module, param, mesh):+        """Wrap ONE parameter as a DTensor placeholder. Default: no-op."""+        pass++    def transform_inputs_pre_forward(self, module, args, kwargs, mesh):+        return args, kwargs++    def context_around_forward(self, module, mesh):+        if self.should_use_local_tensors(module):+            return _use_local_dtensor_params(module)+        return contextlib.nullcontext()++    def transform_output_post_forward(self, module, output, mesh):+        return output++    def install_forward(self, module, mesh):+        """Install pre / around / post transforms by replacing module.forward."""+        original_forward = module.forward++        def ep_forward(*args, **kwargs):+            args, kwargs = self.transform_inputs_pre_forward(module, args, kwargs, mesh)+            with self.context_around_forward(module, mesh):+                output = original_forward(*args, **kwargs)+            return self.transform_output_post_forward(module, output, mesh)++        module.forward = ep_forward+        return module+++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.+    """++    def __init__(self, placement, *, shards_expert_dim: bool = False):+        self.placement = placement+        self.shards_expert_dim = shards_expert_dim++    def shard_param(self, module, param, mesh):+        meta = module._parameters.get(param)+        if meta is None:+            return+        if self.shards_expert_dim and hasattr(module, "num_experts"):+            global_num_experts = meta.shape[0]+            expert_parallel_size = mesh.size()+            if global_num_experts % expert_parallel_size != 0:+                raise ValueError(+                    f"Cannot evenly shard {global_num_experts} experts across "+                    f"{expert_parallel_size} expert-parallel ranks."+                )+            module.num_experts = global_num_experts // expert_parallel_size+            # The experts forward masks sentinel rows only when its experts are actually split.+            module._is_expert_parallel = True+        module._parameters[param] = torch.nn.Parameter(+            distribute_tensor(meta, mesh, [self.placement], src_data_rank=None),+            requires_grad=meta.requires_grad,+        )+++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).+        if torch.is_grad_enabled() and router_scores.requires_grad:+            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 RouterParallelMegaMoe(EpRouterParallel):+    """Router TP plan used with DeepGEMM Mega MoE.++    Mega MoE handles EP dispatch inside the kernel and wants raw global expert ids+    with unmasked routing weights, so the router doesn't pre-shard per EP rank like+    ``EpRouterParallel`` does.+    """++    def transform_output_post_forward(self, module, output, mesh):+        return output+++class EpDispatchExpertsParallel(ExpertParallelLayer):+    """+    Dispatch disjoint TP token slices to the experts' owners, then replicate the combined output on TP.++    Example:+    Let's say we have 8 experts [E0, E7] with DistributedConfig(tp_size=2, fsdp_size=4, ep_size=4). That imply:+        - Since fsdp_size=4, we have 4 batches B denoted [B0, B3]+        - Because we have tp_size=2, that means *both ranks share the same batch*+        - efsdp = (fsdp_size * tp_size) / ep_size = 4 * 2 / 4 = 2++    GPU         0      1       2      3       4      5       6      7+                |      |       |      |       |      |       |      |+    dense view  ---------------------------------------------------------+    batch       [====B0====]   [====B1====]   [====B2====]   [====B3====]+    tp_size     [___________ 0 ___________]   [___________ 1 ___________]+    fsdp_size        0              1              2              3++    expert view ---------------------------------------------------------++    experts     E0E1   E2E3    E4E5   E6E7    E0E1   E2E3    E4E5   E6E7+    ep_size      0      1       2      3       0      1       2      3+    efsdp_size [___________ 0 ___________]   [___________ 1 ___________]++    Assume token 1 in B0 chose E4, which lives on another rank, so it must travel by all-to-all. B0 sits on 2 ranks (cf diagram), so if both sent it, E4 would compute it twice.+    We need to make sure that token 1 is in rank 0 range, so rank 0 sends it and rank 1 does not have it in its slice.+    """++    def should_use_local_tensors(self, module):+        return True++    def _dispatch_tokens(+        self,+        hidden_states: torch.Tensor,+        top_k_index: torch.Tensor,+        num_local_experts: int,+        ep_group,+        ep_size: int,+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[int], list[int]]:+        """Send each selected (token, expert) pair to the rank that owns the expert.++        Also returns the sort order and the per-rank split sizes that `_combine_tokens` needs to reverse the exchange.+        """+        num_top_k = top_k_index.size(-1)++        # Sorting the selected pairs by expert groups them by owner rank, since each rank owns a contiguous range of+        # experts, and the per-expert counts tell every receiver which expert each token it gets is for. The split+        # sizes are the one host sync of the layer.+        expert_ids = top_k_index.reshape(-1)+        order = torch.argsort(expert_ids)+        send_tokens = hidden_states[order // num_top_k]+        send_counts = torch.zeros(num_local_experts * ep_size, dtype=torch.long, device=hidden_states.device)+        send_counts = send_counts.scatter_add_(0, expert_ids, torch.ones_like(expert_ids)).view(+            ep_size, num_local_experts+        )+        recv_counts = torch.empty_like(send_counts)+        torch.distributed.all_to_all_single(recv_counts, send_counts, group=ep_group)+        send_sizes, recv_sizes = torch.stack([send_counts.sum(dim=1), recv_counts.sum(dim=1)]).tolist()+        recv_tokens = all_to_all_single_autograd(send_tokens, recv_sizes, send_sizes, ep_group)+        recv_expert_ids = torch.arange(num_local_experts, device=hidden_states.device).repeat(ep_size)+        recv_expert_ids = recv_expert_ids.repeat_interleave(recv_counts.reshape(-1), output_size=sum(recv_sizes))+        return recv_tokens, recv_expert_ids, order, send_sizes, recv_sizes++    def _run_local_experts(+        self,+        experts_forward: Callable,+        tokens: torch.Tensor,+        expert_ids: torch.Tensor,+    ) -> torch.Tensor:+        """Run local experts with top-1 routing and unit weights; apply routing weights after combine."""+        num_tokens = tokens.shape[0]+        if num_tokens == 0:+            # Keep the empty tokens connected to backward so the reverse all-to-all and FSDP reduction still run.+            dummy = tokens.sum(dim=0, keepdim=True)+            dummy_ids = torch.zeros((1, 1), dtype=expert_ids.dtype, device=tokens.device)+            dummy_weights = torch.ones((1, 1), dtype=tokens.dtype, device=tokens.device)+            return experts_forward(dummy, dummy_ids, dummy_weights)[:0]++        expert_ids = expert_ids.unsqueeze(-1)+        weights = torch.ones_like(expert_ids, dtype=tokens.dtype)+        return experts_forward(tokens, expert_ids, weights)++    def _combine_tokens(+        self,+        expert_output: torch.Tensor,+        top_k_weights: torch.Tensor,+        order: torch.Tensor,+        send_sizes: list[int],+        recv_sizes: list[int],+        ep_group,+    ) -> torch.Tensor:+        """Return expert outputs to the token owners and combine them with routing weights."""+        num_tokens, num_top_k = top_k_weights.shape+        hidden_dim = expert_output.size(-1)+        recv_out = all_to_all_single_autograd(expert_output, send_sizes, recv_sizes, ep_group)+        # Restore the original (token, top-k slot) order, then apply routing weights.+        token_outputs = torch.empty_like(recv_out)+        token_outputs[order] = recv_out+        token_outputs = token_outputs.view(num_tokens, num_top_k, hidden_dim)+        return (token_outputs * top_k_weights.unsqueeze(-1)).sum(dim=1)++    def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, tp_mesh=None):+        hidden_states, top_k_index, top_k_weights = args+        if isinstance(hidden_states, DTensor):+            hidden_states = hidden_states.to_local()+        if isinstance(top_k_weights, DTensor):+            top_k_weights = top_k_weights.to_local()+        if tp_mesh is None or tp_mesh.size() == 1:+            return (hidden_states, top_k_index, top_k_weights), kwargs+        # TP ranks share the same batch, so keep only this rank's rows and each token is dispatched once.+        # Replicate -> Shard(0) is a local chunk (no communication); its backward all-gathers the row gradients.+        hidden_states = DTensor.from_local(hidden_states, tp_mesh, [Replicate()], run_check=False)+        top_k_index = DTensor.from_local(top_k_index, tp_mesh, [Replicate()], run_check=False)+        top_k_weights = DTensor.from_local(top_k_weights, tp_mesh, [Replicate()], run_check=False)++        hidden_states = hidden_states.redistribute(tp_mesh, [Shard(0)]).to_local()+        top_k_index = top_k_index.redistribute(tp_mesh, [Shard(0)]).to_local()+        top_k_weights = top_k_weights.redistribute(tp_mesh, [Shard(0)]).to_local()+        return (hidden_states, top_k_index, top_k_weights), kwargs++    def transform_output_post_forward(self, module, output, mesh, *, tp_mesh=None, num_tokens=None):+        if tp_mesh is None or tp_mesh.size() == 1:+            return output+        # Shard(0) -> Replicate is one all-gather of the row slices, which also handles uneven and empty slices.+        hidden_dim = output.size(-1)+        output = DTensor.from_local(+            output.contiguous(), tp_mesh, [Shard(0)], shape=(num_tokens, hidden_dim), stride=(hidden_dim, 1)+        )+        return output.full_tensor()++    def install_forward(self, module, ep_mesh, *, tp_mesh=None):+        """Experts stay whole (Shard(0) on `ep_mesh`); `tp_mesh` is only the group of ranks holding the same batch."""+        experts_forward = module.forward+        ep_group, ep_size = ep_mesh.get_group(), ep_mesh.size()++        def ep_forward(hidden_states, top_k_index, top_k_weights):+            # Read the full token count before the pre hook slices the inputs across the batch replicas.+            num_tokens = hidden_states.size(0)+            (hidden_states, top_k_index, top_k_weights), _ = self.transform_inputs_pre_forward(+                module, (hidden_states, top_k_index, top_k_weights), {}, ep_mesh, tp_mesh=tp_mesh+            )+            with self.context_around_forward(module, ep_mesh):+                tokens, expert_ids, order, send_sizes, recv_sizes = self._dispatch_tokens(+                    hidden_states, top_k_index, module.num_experts, ep_group, ep_size+                )+                # hand the experts the received rows, not the all-to-all's pending result (whose pointer is null until+                # waited on): https://github.com/huggingface/transformers/pull/49421+                tokens = _WaitTensor.apply(tokens)+                expert_output = self._run_local_experts(experts_forward, tokens, expert_ids)+                output = self._combine_tokens(+                    expert_output, top_k_weights, order, send_sizes, recv_sizes, ep_group+                ).to(hidden_states.dtype)++            return self.transform_output_post_forward(module, output, ep_mesh, tp_mesh=tp_mesh, num_tokens=num_tokens)++        module.forward = ep_forward+        return module+++class MoeTensorParalellMegaMoeExperts(ExpertParallelLayer):+    """TP layer for DeepGEMM Mega MoE experts.++    Mega MoE is inference-only (the kernel has no backward) and handles EP dispatch ++    combine + per-rank token sharding internally — so we skip the gradient-sync hooks+    that ``MoeExpertsParallel`` would apply, and we forward the EP ``process_group``+    into the module so the symm-buffer rendezvous can run on first forward.+    """++    def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_parallel=False):+        hidden_states, top_k_index, top_k_weights = args[0], args[1], args[2]+        return (hidden_states, top_k_index, top_k_weights, mesh.get_group()), kwargs++    def context_around_forward(self, module, mesh):+        return _use_local_dtensor_params(module)++    def transform_output_post_forward(self, module, output, mesh, source=None):+        return output+++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())+    ALL_PARALLEL_STYLES.register("megamoe_router", RouterParallelMegaMoe())+    ALL_PARALLEL_STYLES.register("megamoe_experts", MoeTensorParalellMegaMoeExperts())+++def apply_expert_parallelism(model: nn.Module, ep_mesh: DeviceMesh, tp_mesh: DeviceMesh, plan: dict[str, str]):+    """Shard experts on EP; use TP to split shared tokens and reconstruct outputs around dispatch."""+    for name, module in model.named_modules():+        for p_name, _ in list(module.named_parameters(recurse=False)):+            full = f"{name}.{p_name}" if name else p_name+            style_name = _get_parameter_plan(parameter_name=full, plan=plan, is_weight=True)+            if style_name is not None and style_name in ALL_PARALLEL_STYLES:+                style = ALL_PARALLEL_STYLES[style_name]+                style.validate_param(module, p_name, ep_mesh, parameter_name=full)+                style.shard_param(module, p_name, ep_mesh)++        # Dispatch hooks need both meshes to redistribute tokens between TP and EP ranks.+        style_name = _get_parameter_plan(parameter_name=name, plan=plan, is_weight=False)+        if style_name is not None and style_name in ALL_PARALLEL_STYLES:+            style = ALL_PARALLEL_STYLES[style_name]+            if style_name == "ep_dispatch_experts":+                style.install_forward(module, ep_mesh=ep_mesh, tp_mesh=tp_mesh)+            else:+                style.install_forward(module, ep_mesh)+        module._is_hooked = True++    return model

@@ -31,11 +31,11 @@     consolidate_distributed_checkpoint, ) from .configuration_utils import DistributedConfig+from .expert_parallel import apply_expert_parallelism from .fsdp import apply_fully_sharded_data_parallelism from .pipeline_parallel import apply_pipeline_parallelism from .tensor_parallel import (     _validate_parallel_plan_styles,-    apply_expert_parallelism,     apply_tensor_parallelism,     gather_state_dict_for_save,     resolve_parallel_plans,

@@ -0,0 +1,114 @@+# Copyright 2026 The HuggingFace Team. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.+import contextlib++from ..utils.generic import GeneralInterface+from ..utils.import_utils import is_torch_available, is_torch_distributed_available+++if is_torch_available():+    import torch++if is_torch_distributed_available():+    import torch.distributed as dist+    from torch.distributed._functional_collectives import wait_tensor+    from torch.distributed.tensor import DTensor+++class ParallelInterface(GeneralInterface):+    """Registry of the named parallel styles: `tensor_parallel.py` registers the TP styles, `expert_parallel.py` the EP ones."""++    _global_mapping = {}+++ALL_PARALLEL_STYLES: ParallelInterface = ParallelInterface()+++@contextlib.contextmanager+def _use_local_dtensor_params(module):+    # Kernels as DeepGEMM require local tensors rather than DTensors.+    # We temporarily convert the DTensors to local tensors for the duration of forward() and swap them back after.+    originals = {name: param for name, param in module.named_parameters(recurse=False) if isinstance(param, DTensor)}+    local_params = {name: param.to_local() for name, param in originals.items()}+    module._parameters.update(local_params)+    try:+        yield+    finally:+        for name, original in originals.items():+            current = module._parameters[name]++            if current is local_params[name]:+                module._parameters[name] = original++            # Some kernels such as Megamoe performs changing on FP4 weights and scale factors during its forward pass.+            # That implies creating a new Parameter so we should not restore the original DTensor.+            elif current is not None and not isinstance(current, DTensor):+                replacement = DTensor.from_local(+                    current,+                    original.device_mesh,+                    original.placements,+                    run_check=False,+                )+                module._parameters[name] = torch.nn.Parameter(+                    replacement,+                    requires_grad=current.requires_grad,+                )+++if is_torch_distributed_available():++    class _AllReduceForward(torch.autograd.Function):+        """Allreduce-sum forward, identity backward."""++        @staticmethod+        def forward(ctx, x, process_group):+            if dist.get_world_size(process_group) > 1:+                dist.all_reduce(x, group=process_group)+            return x++        @staticmethod+        def backward(ctx, grad):+            return grad, None++    class _AllReduceBackward(torch.autograd.Function):+        """Identity forward, allreduce-sum backward.++        Used for MoE routing weights: the forward value is replicated (same on all+        ranks), but the backward gradient is partial (each rank has 1/tp_size from+        its expert shard). We need to sum the partial gradients without dividing by+        world_size, which is what DTensor's Replicate backward does incorrectly.+        """++        @staticmethod+        def forward(ctx, x, process_group):+            ctx.process_group = process_group+            return x++        @staticmethod+        def backward(ctx, grad):+            grad = grad.contiguous()+            dist.all_reduce(grad, group=ctx.process_group)+            return grad, None++    class _WaitTensor(torch.autograd.Function):+        """Wait on a functional collective's result, identity backward (`wait_tensor` has no autograd kernel before+        torch 2.11)."""++        @staticmethod+        def forward(ctx, x):+            return wait_tensor(x)++        @staticmethod+        def backward(ctx, grad):+            return grad

@@ -15,19 +15,21 @@  import contextlib import re-from collections.abc import Callable+import warnings from itertools import chain from typing import TYPE_CHECKING  from ..utils import logging-from ..utils.generic import GeneralInterface from ..utils.import_utils import is_torch_available, is_torch_distributed_available+from .parallel_utils import (+    ALL_PARALLEL_STYLES,+    ParallelInterface,  # noqa: F401+    _use_local_dtensor_params,+)+from .utils import _get_torch_distributed_rank   if TYPE_CHECKING:-    from torch import nn-    from torch.distributed.device_mesh import DeviceMesh-     from ..modeling_utils import PreTrainedModel     from .configuration_utils import DistributedConfig @@ -39,10 +41,11 @@  if is_torch_distributed_available():     import torch.distributed as dist-    from torch.distributed._functional_collectives import all_to_all_single_autograd, wait_tensor     from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor     from torch.distributed.tensor.placement_types import _StridedShard +    from .parallel_utils import _AllReduceBackward, _AllReduceForward+  def replace_layer_number_by_wildcard(name: str) -> str:     """@@ -115,37 +118,6 @@ def get_kv_heads_per_rank(tp_plan: dict[str, str], num_heads: int | list[int], t     return num_heads // tp_size if isinstance(num_heads, int) else [h // tp_size for h in layer_heads]  -@contextlib.contextmanager-def _use_local_dtensor_params(module):-    # Kernels as DeepGEMM require local tensors rather than DTensors.-    # We temporarily convert the DTensors to local tensors for the duration of forward() and swap them back after.-    originals = {name: param for name, param in module.named_parameters(recurse=False) if isinstance(param, DTensor)}-    local_params = {name: param.to_local() for name, param in originals.items()}-    module._parameters.update(local_params)-    try:-        yield-    finally:-        for name, original in originals.items():-            current = module._parameters[name]--            if current is local_params[name]:-                module._parameters[name] = original--            # Some kernels such as Megamoe performs changing on FP4 weights and scale factors during its forward pass.-            # That implies creating a new Parameter so we should not restore the original DTensor.-            elif current is not None and not isinstance(current, DTensor):-                replacement = DTensor.from_local(-                    current,-                    original.device_mesh,-                    original.placements,-                    run_check=False,-                )-                module._parameters[name] = torch.nn.Parameter(-                    replacement,-                    requires_grad=current.requires_grad,-                )-- class TensorParallelLayer:     def should_use_local_tensors(self, module):         """Whether this module's forward requires local inputs and parameters."""@@ -603,86 +575,6 @@ def shard_param(self, module, param, mesh):         )  -class MoEParamShard(TensorParallelLayer):-    """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.-    """--    def __init__(self, placement, *, shards_expert_dim: bool = False):-        self.placement = placement-        self.shards_expert_dim = shards_expert_dim--    def shard_param(self, module, param, mesh):-        meta = module._parameters.get(param)-        if meta is None:-            return-        if self.shards_expert_dim and hasattr(module, "num_experts"):-            global_num_experts = meta.shape[0]-            expert_parallel_size = mesh.size()-            if global_num_experts % expert_parallel_size != 0:-                raise ValueError(-                    f"Cannot evenly shard {global_num_experts} experts across "-                    f"{expert_parallel_size} expert-parallel ranks."-                )-            module.num_experts = global_num_experts // expert_parallel_size-            # The experts forward masks sentinel rows only when its experts are actually split.-            module._is_expert_parallel = True-        module._parameters[param] = torch.nn.Parameter(-            distribute_tensor(meta, mesh, [self.placement], src_data_rank=None),-            requires_grad=meta.requires_grad,-        )---if is_torch_distributed_available():--    class _AllReduceForward(torch.autograd.Function):-        """Allreduce-sum forward, identity backward."""--        @staticmethod-        def forward(ctx, x, process_group):-            if dist.get_world_size(process_group) > 1:-                dist.all_reduce(x, group=process_group)-            return x--        @staticmethod-        def backward(ctx, grad):-            return grad, None--    class _AllReduceBackward(torch.autograd.Function):-        """Identity forward, allreduce-sum backward.--        Used for MoE routing weights: the forward value is replicated (same on all-        ranks), but the backward gradient is partial (each rank has 1/tp_size from-        its expert shard). We need to sum the partial gradients without dividing by-        world_size, which is what DTensor's Replicate backward does incorrectly.-        """--        @staticmethod-        def forward(ctx, x, process_group):-            ctx.process_group = process_group-            return x--        @staticmethod-        def backward(ctx, grad):-            grad = grad.contiguous()-            dist.all_reduce(grad, group=ctx.process_group)-            return grad, None--    class _WaitTensor(torch.autograd.Function):-        """Wait on a functional collective's result, identity backward (`wait_tensor` has no autograd kernel before-        torch 2.11)."""--        @staticmethod-        def forward(ctx, x):-            return wait_tensor(x)--        @staticmethod-        def backward(ctx, grad):-            return grad-- class MoeExpertsParallel(TensorParallelLayer):     def should_use_local_tensors(self, module):         return True@@ -753,300 +645,29 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh):         return (input_tensor / mesh.size(), *args[1:]), kwargs  -class EpRouterParallel(TensorParallelLayer):-    """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).-        if torch.is_grad_enabled() and router_scores.requires_grad:-            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 RouterParallelMegaMoe(EpRouterParallel):-    """Router TP plan used with DeepGEMM Mega MoE.--    Mega MoE handles EP dispatch inside the kernel and wants raw global expert ids-    with unmasked routing weights, so the router doesn't pre-shard per EP rank like-    ``EpRouterParallel`` does.-    """--    def transform_output_post_forward(self, module, output, mesh):-        return output---class EpDispatchExpertsParallel(MoeExpertsParallel):-    """-    Dispatch disjoint TP token slices to the experts' owners, then replicate the combined output on TP.--    Example:-    Let's say we have 8 experts [E0, E7] with DistributedConfig(tp_size=2, fsdp_size=4, ep_size=4). That imply:-        - Since fsdp_size=4, we have 4 batches B denoted [B0, B3]-        - Because we have tp_size=2, that means *both ranks share the same batch*-        - efsdp = (fsdp_size * tp_size) / ep_size = 4 * 2 / 4 = 2--    GPU         0      1       2      3       4      5       6      7-                |      |       |      |       |      |       |      |-    dense view  ----------------------------------------------------------    batch       [====B0====]   [====B1====]   [====B2====]   [====B3====]-    tp_size     [___________ 0 ___________]   [___________ 1 ___________]-    fsdp_size        0              1              2              3--    expert view -----------------------------------------------------------    experts     E0E1   E2E3    E4E5   E6E7    E0E1   E2E3    E4E5   E6E7-    ep_size      0      1       2      3       0      1       2      3-    efsdp_size [___________ 0 ___________]   [___________ 1 ___________]--    Assume token 1 in B0 chose E4, which lives on another rank, so it must travel by all-to-all. B0 sits on 2 ranks (cf diagram), so if both sent it, E4 would compute it twice.-    We need to make sure that token 1 is in rank 0 range, so rank 0 sends it and rank 1 does not have it in its slice.-    """--    def _dispatch_tokens(-        self,-        hidden_states: torch.Tensor,-        top_k_index: torch.Tensor,-        num_local_experts: int,-        ep_group,-        ep_size: int,-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[int], list[int]]:-        """Send each selected (token, expert) pair to the rank that owns the expert.--        Also returns the sort order and the per-rank split sizes that `_combine_tokens` needs to reverse the exchange.-        """-        num_top_k = top_k_index.size(-1)--        # Sorting the selected pairs by expert groups them by owner rank, since each rank owns a contiguous range of-        # experts, and the per-expert counts tell every receiver which expert each token it gets is for. The split-        # sizes are the one host sync of the layer.-        expert_ids = top_k_index.reshape(-1)-        order = torch.argsort(expert_ids)-        send_tokens = hidden_states[order // num_top_k]-        send_counts = torch.zeros(num_local_experts * ep_size, dtype=torch.long, device=hidden_states.device)-        send_counts = send_counts.scatter_add_(0, expert_ids, torch.ones_like(expert_ids)).view(-            ep_size, num_local_experts-        )-        recv_counts = torch.empty_like(send_counts)-        torch.distributed.all_to_all_single(recv_counts, send_counts, group=ep_group)-        send_sizes, recv_sizes = torch.stack([send_counts.sum(dim=1), recv_counts.sum(dim=1)]).tolist()-        recv_tokens = all_to_all_single_autograd(send_tokens, recv_sizes, send_sizes, ep_group)-        recv_expert_ids = torch.arange(num_local_experts, device=hidden_states.device).repeat(ep_size)-        recv_expert_ids = recv_expert_ids.repeat_interleave(recv_counts.reshape(-1), output_size=sum(recv_sizes))-        return recv_tokens, recv_expert_ids, order, send_sizes, recv_sizes--    def _run_local_experts(-        self,-        experts_forward: Callable,-        tokens: torch.Tensor,-        expert_ids: torch.Tensor,-    ) -> torch.Tensor:-        """Run local experts with top-1 routing and unit weights; apply routing weights after combine."""-        num_tokens = tokens.shape[0]-        if num_tokens == 0:-            # Keep the empty tokens connected to backward so the reverse all-to-all and FSDP reduction still run.-            dummy = tokens.sum(dim=0, keepdim=True)-            dummy_ids = torch.zeros((1, 1), dtype=expert_ids.dtype, device=tokens.device)-            dummy_weights = torch.ones((1, 1), dtype=tokens.dtype, device=tokens.device)-            return experts_forward(dummy, dummy_ids, dummy_weights)[:0]--        expert_ids = expert_ids.unsqueeze(-1)-        weights = torch.ones_like(expert_ids, dtype=tokens.dtype)-        return experts_forward(tokens, expert_ids, weights)--    def _combine_tokens(-        self,-        expert_output: torch.Tensor,-        top_k_weights: torch.Tensor,-        order: torch.Tensor,-        send_sizes: list[int],-        recv_sizes: list[int],-        ep_group,-    ) -> torch.Tensor:-        """Return expert outputs to the token owners and combine them with routing weights."""-        num_tokens, num_top_k = top_k_weights.shape-        hidden_dim = expert_output.size(-1)-        recv_out = all_to_all_single_autograd(expert_output, send_sizes, recv_sizes, ep_group)-        # Restore the original (token, top-k slot) order, then apply routing weights.-        token_outputs = torch.empty_like(recv_out)-        token_outputs[order] = recv_out-        token_outputs = token_outputs.view(num_tokens, num_top_k, hidden_dim)-        return (token_outputs * top_k_weights.unsqueeze(-1)).sum(dim=1)--    def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, tp_mesh=None):-        hidden_states, top_k_index, top_k_weights = args-        if isinstance(hidden_states, DTensor):-            hidden_states = hidden_states.to_local()-        if isinstance(top_k_weights, DTensor):-            top_k_weights = top_k_weights.to_local()-        if tp_mesh is None or tp_mesh.size() == 1:-            return (hidden_states, top_k_index, top_k_weights), kwargs-        # TP ranks share the same batch, so keep only this rank's rows and each token is dispatched once.-        # Replicate -> Shard(0) is a local chunk (no communication); its backward all-gathers the row gradients.-        hidden_states = DTensor.from_local(hidden_states, tp_mesh, [Replicate()], run_check=False)-        top_k_index = DTensor.from_local(top_k_index, tp_mesh, [Replicate()], run_check=False)-        top_k_weights = DTensor.from_local(top_k_weights, tp_mesh, [Replicate()], run_check=False)--        hidden_states = hidden_states.redistribute(tp_mesh, [Shard(0)]).to_local()-        top_k_index = top_k_index.redistribute(tp_mesh, [Shard(0)]).to_local()-        top_k_weights = top_k_weights.redistribute(tp_mesh, [Shard(0)]).to_local()-        return (hidden_states, top_k_index, top_k_weights), kwargs--    def transform_output_post_forward(self, module, output, mesh, *, tp_mesh=None, num_tokens=None):-        if tp_mesh is None or tp_mesh.size() == 1:-            return output-        # Shard(0) -> Replicate is one all-gather of the row slices, which also handles uneven and empty slices.-        hidden_dim = output.size(-1)-        output = DTensor.from_local(-            output.contiguous(), tp_mesh, [Shard(0)], shape=(num_tokens, hidden_dim), stride=(hidden_dim, 1)-        )-        return output.full_tensor()--    def install_forward(self, module, ep_mesh, *, tp_mesh=None):-        """Experts stay whole (Shard(0) on `ep_mesh`); `tp_mesh` is only the group of ranks holding the same batch."""-        experts_forward = module.forward-        ep_group, ep_size = ep_mesh.get_group(), ep_mesh.size()--        def ep_forward(hidden_states, top_k_index, top_k_weights):-            # Read the full token count before the pre hook slices the inputs across the batch replicas.-            num_tokens = hidden_states.size(0)-            (hidden_states, top_k_index, top_k_weights), _ = self.transform_inputs_pre_forward(-                module, (hidden_states, top_k_index, top_k_weights), {}, ep_mesh, tp_mesh=tp_mesh-            )-            with self.context_around_forward(module, ep_mesh):-                tokens, expert_ids, order, send_sizes, recv_sizes = self._dispatch_tokens(-                    hidden_states, top_k_index, module.num_experts, ep_group, ep_size-                )-                # hand the experts the received rows, not the all-to-all's pending result (whose pointer is null until-                # waited on): https://github.com/huggingface/transformers/pull/49421-                tokens = _WaitTensor.apply(tokens)-                expert_output = self._run_local_experts(experts_forward, tokens, expert_ids)-                output = self._combine_tokens(-                    expert_output, top_k_weights, order, send_sizes, recv_sizes, ep_group-                ).to(hidden_states.dtype)--            return self.transform_output_post_forward(module, output, ep_mesh, tp_mesh=tp_mesh, num_tokens=num_tokens)--        module.forward = ep_forward-        return module---class MoeTensorParalellMegaMoeExperts(MoeExpertsParallel):-    """TP layer for DeepGEMM Mega MoE experts.--    Mega MoE is inference-only (the kernel has no backward) and handles EP dispatch +-    combine + per-rank token sharding internally — so we skip the gradient-sync hooks-    that ``MoeExpertsParallel`` would apply, and we forward the EP ``process_group``-    into the module so the symm-buffer rendezvous can run on first forward.-    """--    def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_parallel=False):-        hidden_states, top_k_index, top_k_weights = args[0], args[1], args[2]-        return (hidden_states, top_k_index, top_k_weights, mesh.get_group()), kwargs--    def context_around_forward(self, module, mesh):-        return _use_local_dtensor_params(module)--    def transform_output_post_forward(self, module, output, mesh, source=None):-        return output---class ParallelInterface(GeneralInterface):-    """Registry of named TP styles for the DTensor backend."""--    _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(),-            "ep_dispatch_experts": EpDispatchExpertsParallel(),-            "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(),-            # Every rank keeps the full param, e.g. a per-tensor quantization scale-            "replicate": TensorParallelLayer(),-        }-        if is_torch_distributed_available()-        else {}+if is_torch_distributed_available():+    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: ParallelInterface = ParallelInterface()+    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 rank keeps the full param, e.g. a per-tensor quantization scale+    ALL_PARALLEL_STYLES.register("replicate", TensorParallelLayer())   def _validate_parallel_plan_styles(plan: dict[str, str] | None) -> None:@@ -1125,30 +746,6 @@ def apply_tensor_parallelism(model, tp_mesh, tp_plan=None):     return model  -def apply_expert_parallelism(model: nn.Module, ep_mesh: DeviceMesh, tp_mesh: DeviceMesh, plan: dict[str, str]):-    """Shard experts on EP; use TP to split shared tokens and reconstruct outputs around dispatch."""-    for name, module in model.named_modules():-        for p_name, _ in list(module.named_parameters(recurse=False)):-            full = f"{name}.{p_name}" if name else p_name-            style_name = _get_parameter_plan(parameter_name=full, plan=plan, is_weight=True)-            if style_name is not None and style_name in ALL_PARALLEL_STYLES:-                style = ALL_PARALLEL_STYLES[style_name]-                style.validate_param(module, p_name, ep_mesh, parameter_name=full)-                style.shard_param(module, p_name, ep_mesh)--        # Dispatch hooks need both meshes to redistribute tokens between TP and EP ranks.-        style_name = _get_parameter_plan(parameter_name=name, plan=plan, is_weight=False)-        if style_name is not None and style_name in ALL_PARALLEL_STYLES:-            style = ALL_PARALLEL_STYLES[style_name]-            if style_name == "ep_dispatch_experts":-                style.install_forward(module, ep_mesh=ep_mesh, tp_mesh=tp_mesh)-            else:-                style.install_forward(module, ep_mesh)-        module._is_hooked = True--    return model-- def gather_state_dict_for_save(     state_dict: dict[str, torch.Tensor],     _tp_plan: dict[str, str],@@ -1168,3 +765,30 @@ def gather_state_dict_for_save(         else:             gathered[key] = tensor     return gathered+++# BC: these moved to expert_parallel.py in v5.21, keep the old import path with a warning until v5.23.+# Lazy import because expert_parallel.py imports from this file (circular import otherwise).+_MOVED_TO_EXPERT_PARALLEL = (+    "MoEParamShard",+    "EpRouterParallel",+    "RouterParallelMegaMoe",+    "EpDispatchExpertsParallel",+    "MoeTensorParalellMegaMoeExperts",+    "apply_expert_parallelism",+)+++def __getattr__(name: str):+    if name in _MOVED_TO_EXPERT_PARALLEL:+        from . import expert_parallel++        if _get_torch_distributed_rank() == 0:+            warnings.warn(+                f"`transformers.distributed.tensor_parallel.{name}` is deprecated and will be removed in v5.23: "+                f"import it from `transformers.distributed.expert_parallel` instead.",+                FutureWarning,+                stacklevel=2,+            )+        return getattr(expert_parallel, name)+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

@@ -13,28 +13,31 @@ # limitations under the License. """Backward-compatibility shim for the tensor parallel API. -The canonical implementation lives in ``transformers.distributed.tensor_parallel``.+The canonical implementation lives in ``transformers.distributed.tensor_parallel`` and, for the expert-parallel+styles, ``transformers.distributed.expert_parallel``. """  from __future__ import annotations  import warnings +from ..distributed.expert_parallel import (+    EpRouterParallel,+    MoEParamShard,+    MoeTensorParalellMegaMoeExperts,+    RouterParallelMegaMoe,+) from ..distributed.tensor_parallel import (     ALL_PARALLEL_STYLES,     AllReduceParallel,     ColwiseParallel,-    EpRouterParallel,     MlaKvAProjParallel,     MoeExpertsParallel,     MoeIdentityParallel,-    MoEParamShard,-    MoeTensorParalellMegaMoeExperts,     PackedColwiseParallel,     PackedRowwiseParallel,     ParallelInterface,     ReplicatedWithGradAllReduce,-    RouterParallelMegaMoe,     RowwiseParallel,     SequenceParallel,     TensorParallelLayer,

@@ -325,6 +325,7 @@ def parse_int_from_env(key, default=None): _run_agent_tests = parse_flag_from_env("RUN_AGENT_TESTS", default=False) _run_training_tests = parse_flag_from_env("RUN_TRAINING_TESTS", default=True) _run_tensor_parallel_tests = parse_flag_from_env("RUN_TENSOR_PARALLEL_TESTS", default=True)+_run_expert_parallel_tests = parse_flag_from_env("RUN_EXPERT_PARALLEL_TESTS", default=True) _run_fsdp_tests = parse_flag_from_env("RUN_FSDP_TESTS", default=True)  @@ -408,6 +409,22 @@ def is_tensor_parallel_test(test_case):             return pytest.mark.is_tensor_parallel_test()(test_case)  +def is_expert_parallel_test(test_case):+    """+    Decorator marking a test as an expert parallel test. If RUN_EXPERT_PARALLEL_TESTS is set to a falsy value, those+    tests will be skipped.+    """+    if not _run_expert_parallel_tests:+        return unittest.skip(reason="test is expert parallel test")(test_case)+    else:+        try:+            import pytest  # We don't need a hard dependency on pytest in the main library+        except ImportError:+            return test_case+        else:+            return pytest.mark.is_expert_parallel_test()(test_case)++ def is_fsdp_test(test_case):     """     Decorator marking a test as an FSDP test. If RUN_FSDP_TESTS is set to a falsy value, those tests will be

@@ -29,6 +29,7 @@ )  from .test_configuration_common import ConfigTester+from .test_expert_parallel_mixin import ExpertParallelTesterMixin from .test_fsdp_mixin import FSDPTesterMixin from .test_modeling_common import (     GenerationTesterMixin,@@ -276,6 +277,7 @@ class CausalLMModelTest(     PipelineTesterMixin,     TrainingTesterMixin,     TensorParallelTesterMixin,+    ExpertParallelTesterMixin,     FSDPTesterMixin, ):     model_tester_class = None

GitHub publishes no diff for this file. View it on GitHub ↗

@@ -0,0 +1,96 @@+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.+from unittest.mock import patch++import torch+from torch.distributed.device_mesh import init_device_mesh++from transformers.distributed import expert_parallel+from transformers.distributed.expert_parallel import EpDispatchExpertsParallel+from transformers.distributed.parallel_utils import ALL_PARALLEL_STYLES+from transformers.testing_utils import TestCasePlus, is_expert_parallel_test++from ..test_tensor_parallel_mixin import _init_distributed+++@is_expert_parallel_test+class TestExpertParallelLayer(TestCasePlus):+    class MockDeviceMesh:+        def __init__(self, world_size, rank):+            self.world_size = world_size+            self.rank = rank+            self.shape = (world_size,)+            self.ndim = 1++        def size(self):+            return self.world_size++        def get_local_rank(self):+            return self.rank++    def _get_parameter_placements(self, module, style, mesh):+        placements = {}+        with patch.object(+            expert_parallel, "distribute_tensor", side_effect=lambda tensor, *args, **kwargs: tensor+        ) as distribute:+            for parameter_name in list(module._parameters):+                style.shard_param(module, parameter_name, mesh)+                placements[parameter_name] = distribute.call_args.args[2][0]++        return placements++    def test_grouped_gemm_updates_local_expert_count(self):+        module = torch.nn.Module()+        module.num_experts = 8+        module.register_parameter("weight", torch.nn.Parameter(torch.empty(8, 16, 32)))+        grouped_gemm = ALL_PARALLEL_STYLES["grouped_gemm"]++        placements = self._get_parameter_placements(module, grouped_gemm, self.MockDeviceMesh(world_size=4, rank=0))++        self.assertEqual(placements["weight"].dim, 0)+        self.assertEqual(module.num_experts, 2)++    def test_sharding_does_not_create_unrelated_module_attributes(self):+        module = torch.nn.Module()+        module.random_attr = 123+        module.register_parameter("weight", torch.nn.Parameter(torch.empty(8, 16, 32)))++        self._get_parameter_placements(+            module, ALL_PARALLEL_STYLES["grouped_gemm"], self.MockDeviceMesh(world_size=4, rank=0)+        )++        self.assertEqual(module.random_attr, 123)+        self.assertFalse(hasattr(module, "num_experts"))++    def test_dispatch_hands_experts_the_received_rows(self):+        """The all-to-all's pending result (`AsyncCollectiveTensor`) holds a null pointer until a torch op waits on+        it, and experts that launch kernels on raw pointers would read it: the dispatch waits before any experts run."""+        _init_distributed(tp=2, backend="gloo")(_dispatch_hands_experts_the_received_rows)()+++class _DummyExperts(torch.nn.Module):+    def __init__(self):+        super().__init__()+        self.num_experts = 2++    def forward(self, tokens, expert_ids, weights):+        assert type(tokens) is torch.Tensor, type(tokens)+        return tokens+++# Worker functions for the expert-parallel layer tests, spawned through `_init_distributed`.+def _dispatch_hands_experts_the_received_rows(rank):+    hidden_states = torch.arange(32.0).view(4, 8) + 100 * rank+    experts = EpDispatchExpertsParallel().install_forward(_DummyExperts(), init_device_mesh("cpu", (2,)))+    experts(hidden_states, torch.tensor([[0, 3], [1, 2], [2, 1], [3, 0]]), torch.ones(4, 2))

@@ -30,6 +30,7 @@  from ...generation.test_utils import GenerationTesterMixin from ...test_configuration_common import ConfigTester+from ...test_expert_parallel_mixin import ExpertParallelTesterMixin from ...test_modeling_common import ModelTesterMixin, ids_tensor from ...test_pipeline_mixin import PipelineTesterMixin from ...test_tensor_parallel_mixin import TensorParallelTesterMixin@@ -214,7 +215,12 @@ def prepare_config_and_inputs_for_common(self):  @require_torch class DeepseekV3ModelTest(-    ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase, TensorParallelTesterMixin+    ModelTesterMixin,+    GenerationTesterMixin,+    PipelineTesterMixin,+    unittest.TestCase,+    TensorParallelTesterMixin,+    ExpertParallelTesterMixin, ):     all_model_classes = (         (

@@ -30,6 +30,7 @@ )  from ...test_configuration_common import ConfigTester+from ...test_expert_parallel_mixin import ExpertParallelTesterMixin from ...test_memory_cleanup_mixin import MemoryCleanupMixin from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor from ...test_tensor_parallel_mixin import TensorParallelTesterMixin@@ -187,7 +188,9 @@ def prepare_config_and_inputs_for_common(self):   @require_torch-class DiffusionGemmaVisionText2TextModelTest(ModelTesterMixin, TensorParallelTesterMixin, unittest.TestCase):+class DiffusionGemmaVisionText2TextModelTest(+    ModelTesterMixin, TensorParallelTesterMixin, ExpertParallelTesterMixin, unittest.TestCase+):     all_model_classes = (DiffusionGemmaModel, DiffusionGemmaForBlockDiffusion) if is_torch_available() else ()     all_generative_model_classes = ()  # No class inherits `GenerationMixin`     additional_model_inputs = ["mm_token_type_ids", "decoder_input_ids", "image_position_ids"]

@@ -36,6 +36,7 @@  from ...generation.test_utils import GenerationTesterMixin from ...test_configuration_common import ConfigTester+from ...test_expert_parallel_mixin import ExpertParallelTesterMixin from ...test_memory_cleanup_mixin import MemoryCleanupMixin from ...test_modeling_common import (     ModelTesterMixin,@@ -190,7 +191,9 @@ def prepare_config_and_inputs_for_common(self):   @require_torch-class Glm4vMoeModelTest(ModelTesterMixin, GenerationTesterMixin, TensorParallelTesterMixin, unittest.TestCase):+class Glm4vMoeModelTest(+    ModelTesterMixin, GenerationTesterMixin, TensorParallelTesterMixin, ExpertParallelTesterMixin, unittest.TestCase+):     all_model_classes = (Glm4vMoeModel, Glm4vMoeForConditionalGeneration) if is_torch_available() else ()      model_split_percents = [0.7, 0.9]  # model too big to split at 0.5

@@ -40,6 +40,7 @@ from ...generation.test_utils import GenerationTesterMixin from ...models.bamba.test_modeling_bamba import BambaModelTester from ...test_configuration_common import ConfigTester+from ...test_expert_parallel_mixin import ExpertParallelTesterMixin from ...test_modeling_common import ModelTesterMixin from ...test_pipeline_mixin import PipelineTesterMixin from ...test_tensor_parallel_mixin import TensorParallelTesterMixin@@ -85,7 +86,12 @@ def get_config(self):  @require_torch class GraniteMoeHybridModelTest(-    ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, TensorParallelTesterMixin, unittest.TestCase+    ModelTesterMixin,+    GenerationTesterMixin,+    PipelineTesterMixin,+    TensorParallelTesterMixin,+    ExpertParallelTesterMixin,+    unittest.TestCase, ):     model_tester_class = GraniteMoeHybridModelTester     all_model_classes = (

@@ -36,6 +36,7 @@  from ...generation.test_utils import GenerationTesterMixin from ...test_configuration_common import ConfigTester+from ...test_expert_parallel_mixin import ExpertParallelTesterMixin from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask from ...test_pipeline_mixin import PipelineTesterMixin from ...test_tensor_parallel_mixin import TensorParallelTesterMixin@@ -302,7 +303,12 @@ def prepare_config_and_inputs_for_common(self):  @require_torch class JambaModelTest(-    ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, TensorParallelTesterMixin, unittest.TestCase+    ModelTesterMixin,+    GenerationTesterMixin,+    PipelineTesterMixin,+    TensorParallelTesterMixin,+    ExpertParallelTesterMixin,+    unittest.TestCase, ):     all_model_classes = (         (

@@ -34,6 +34,7 @@  from ...generation.test_utils import GenerationTesterMixin from ...test_configuration_common import ConfigTester+from ...test_expert_parallel_mixin import ExpertParallelTesterMixin from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask from ...test_pipeline_mixin import PipelineTesterMixin from ...test_tensor_parallel_mixin import TensorParallelTesterMixin@@ -403,7 +404,12 @@ def prepare_config_and_inputs_for_common(self):  @require_torch class NemotronHModelTest(-    ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, TensorParallelTesterMixin, unittest.TestCase+    ModelTesterMixin,+    GenerationTesterMixin,+    PipelineTesterMixin,+    TensorParallelTesterMixin,+    ExpertParallelTesterMixin,+    unittest.TestCase, ):     all_model_classes = (         (

@@ -27,6 +27,7 @@  from ...generation.test_utils import GenerationTesterMixin from ...test_configuration_common import ConfigTester+from ...test_expert_parallel_mixin import ExpertParallelTesterMixin from ...test_modeling_common import ModelTesterMixin, ids_tensor from ...test_pipeline_mixin import PipelineTesterMixin from ...test_tensor_parallel_mixin import TensorParallelTesterMixin@@ -177,7 +178,12 @@ def prepare_config_and_inputs_for_common(self):  @require_torch class OlmoeModelTest(-    ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, TensorParallelTesterMixin, unittest.TestCase+    ModelTesterMixin,+    GenerationTesterMixin,+    PipelineTesterMixin,+    TensorParallelTesterMixin,+    ExpertParallelTesterMixin,+    unittest.TestCase, ):     all_model_classes = (OlmoeModel, OlmoeForCausalLM) if is_torch_available() else ()     pipeline_model_mapping = (

@@ -27,6 +27,7 @@ from transformers.testing_utils import Expectations, require_torch, slow, torch_device  from ...test_configuration_common import ConfigTester+from ...test_expert_parallel_mixin import ExpertParallelTesterMixin from ...test_modeling_common import (     TEST_EAGER_MATCHES_BATCHED_AND_GROUPED_INFERENCE_PARAMETERIZATION,     ModelTesterMixin,@@ -162,7 +163,7 @@ def prepare_config_and_inputs_for_common(self):  @require_torch class OpenAIPrivacyFilterModelTest(-    ModelTesterMixin, PipelineTesterMixin, TensorParallelTesterMixin, unittest.TestCase+    ModelTesterMixin, PipelineTesterMixin, TensorParallelTesterMixin, ExpertParallelTesterMixin, unittest.TestCase ):     all_model_classes = (         (

@@ -49,6 +49,7 @@  from ...generation.test_utils import GenerationTesterMixin from ...test_configuration_common import ConfigTester+from ...test_expert_parallel_mixin import ExpertParallelTesterMixin from ...test_memory_cleanup_mixin import MemoryCleanupMixin from ...test_modeling_common import (     ModelTesterMixin,@@ -264,7 +265,7 @@ def create_and_check_qwenomnithinker_model_fp16_forward(self, config, input_ids,  @require_torch class Qwen3OmniMoeThinkerForConditionalGenerationModelTest(-    ModelTesterMixin, GenerationTesterMixin, TensorParallelTesterMixin, unittest.TestCase+    ModelTesterMixin, GenerationTesterMixin, TensorParallelTesterMixin, ExpertParallelTesterMixin, unittest.TestCase ):     """     Model tester for `Qwen3OmniMoeThinkerForConditionalGeneration`.

@@ -17,6 +17,7 @@ from transformers.testing_utils import _TEXT_MODEL_TESTER_DEFAULTS  from .test_configuration_common import ConfigTester+from .test_expert_parallel_mixin import ExpertParallelTesterMixin from .test_modeling_common import (     GenerationTesterMixin,     ModelTesterMixin,@@ -213,7 +214,9 @@ def create_and_check_model(   @require_torch-class MultiModalModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, TensorParallelTesterMixin):+class MultiModalModelTest(+    ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, TensorParallelTesterMixin, ExpertParallelTesterMixin+):     """Shared test-class base for multimodal model families.      Subclasses must set:

@@ -15,24 +15,20 @@ from unittest.mock import patch  import torch-from torch.distributed.device_mesh import init_device_mesh  from transformers import AutoModelForCausalLM from transformers.distributed import tensor_parallel from transformers.distributed.sharding_utils import DtensorShardOperation from transformers.distributed.tensor_parallel import (     ALL_PARALLEL_STYLES,     ColwiseParallel,-    EpDispatchExpertsParallel,     PackedColwiseParallel,     PackedRowwiseParallel,     RowwiseParallel,     get_kv_heads_per_rank, ) from transformers.testing_utils import TestCasePlus, is_tensor_parallel_test -from ..test_tensor_parallel_mixin import _init_distributed-  @is_tensor_parallel_test class TestTensorParallelProperties(TestCasePlus):@@ -355,19 +351,8 @@ def test_packed_rowwise_packed_and_unpacked_shapes(self):             self.assertEqual(op.shard_tensor(packed).shape, (16, 32))             self.assertEqual(op.shard_tensor(unpacked).shape, (16, 16)) -    def test_grouped_gemm_updates_local_expert_count(self):-        module = torch.nn.Module()-        module.num_experts = 8-        module.register_parameter("weight", torch.nn.Parameter(torch.empty(8, 16, 32)))-        grouped_gemm = ALL_PARALLEL_STYLES["grouped_gemm"]--        placements = self._get_parameter_placements(module, grouped_gemm, self.MockDeviceMesh(world_size=4, rank=0))--        self.assertEqual(placements["weight"].dim, 0)-        self.assertEqual(module.num_experts, 2)-     def test_sharding_does_not_create_unrelated_module_attributes(self):-        styles = (ColwiseParallel(), RowwiseParallel(), ALL_PARALLEL_STYLES["grouped_gemm"])+        styles = (ColwiseParallel(), RowwiseParallel())          for style in styles:             with self.subTest(style=type(style).__name__):@@ -379,25 +364,3 @@ def test_sharding_does_not_create_unrelated_module_attributes(self):                  self.assertEqual(module.random_attr, 123)                 self.assertFalse(hasattr(module, "num_experts"))--    def test_dispatch_hands_experts_the_received_rows(self):-        """The all-to-all's pending result (`AsyncCollectiveTensor`) holds a null pointer until a torch op waits on-        it, and experts that launch kernels on raw pointers would read it: the dispatch waits before any experts run."""-        _init_distributed(tp=2, backend="gloo")(_dispatch_hands_experts_the_received_rows)()---class _DummyExperts(torch.nn.Module):-    def __init__(self):-        super().__init__()-        self.num_experts = 2--    def forward(self, tokens, expert_ids, weights):-        assert type(tokens) is torch.Tensor, type(tokens)-        return tokens---# Worker functions for the expert-parallel layer tests, spawned through `_init_distributed`.-def _dispatch_hands_experts_the_received_rows(rank):-    hidden_states = torch.arange(32.0).view(4, 8) + 100 * rank-    experts = EpDispatchExpertsParallel().install_forward(_DummyExperts(), init_device_mesh("cpu", (2,)))-    experts(hidden_states, torch.tensor([[0, 3], [1, 2], [2, 1], [3, 0]]), torch.ones(4, 2))

@@ -0,0 +1,230 @@+# Copyright 2026 The HuggingFace Team. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.+"""Expert parallel tests for model test classes; the distributed harness lives in `test_tensor_parallel_mixin.py`."""++import tempfile+from abc import ABC+from itertools import product++from parameterized import parameterized++from transformers import set_seed+from transformers.distributed.configuration_utils import DistributedConfig+from transformers.testing_utils import is_expert_parallel_test, is_torch_available++from .test_tensor_parallel_mixin import _get_expert_weight_names, _init_distributed, _reference_grad_shard+++if is_torch_available():+    import torch+    import torch.distributed as dist+    from torch.distributed.tensor import DTensor+++# =============================================================================+# Constants+# =============================================================================++# Set to None to run distributed EP tests for every model with a plan.+# Representative MoE model types covered by distributed EP tests.+EP_DISTRIBUTED_TEST_MODEL_TYPES = {+    "qwen4_exp_text",+    "qwen3_moe",+    "glm_moe_dsa",+    "deepseek_v4",+    "glm4_moe_lite",+    "glm4_moe",+    "olmoe",+    "qwen2_moe",+    "cohere2_moe",+    # VLM+    "glm4v_moe",+}+++# =============================================================================+# Distributed helpers (top-level for pickling by mp.spawn)+# =============================================================================+++def _load_ep_and_reference_models(model_path, model_class):+    """Load EP model and non-EP reference model for comparison."""+    model_ep = model_class.from_pretrained(+        model_path,+        distributed_config=DistributedConfig(tp_size=dist.get_world_size(), ep_size=dist.get_world_size()),+    )+    dist.barrier()++    device = model_ep.device+    model_ref = model_class.from_pretrained(model_path)+    model_ref = model_ref.to(device)++    return model_ep, model_ref, device+++def _test_ep_forward_impl(_rank, model_path, model_class, atol, rtol, experts_implementation):+    """Implementation for comparing EP and non-EP model outputs."""+    set_seed(0)++    model_ep, model_ref, device = _load_ep_and_reference_models(model_path, model_class)++    model_ep.eval()+    model_ref.eval()++    model_ep.set_experts_implementation(experts_implementation)+    model_ref.set_experts_implementation(experts_implementation)++    vocab_size = model_ref.config.get_text_config().vocab_size+    input_ids = torch.randint(0, vocab_size, (2, 64)).to(device)++    with torch.no_grad():+        logits_ref = model_ref(input_ids).logits+        logits_ep = model_ep(input_ids).logits++    diff = (logits_ref - logits_ep).abs()+    assert torch.allclose(logits_ref, logits_ep, atol=atol, rtol=rtol), (+        f"EP and non-EP model outputs differ. Max diff: {diff.max().item()} | Min diff: {diff.min().item()}"+    )++    dist.barrier()+++def _test_ep_backward_impl(rank, model_path, model_class, atol, rtol, experts_implementation):+    """Implementation for comparing EP and non-EP model backward passes."""+    set_seed(0)++    model_ep, model_ref, device = _load_ep_and_reference_models(model_path, model_class)+    model_ep.train()+    model_ref.train()++    model_ep.set_experts_implementation(experts_implementation)+    model_ref.set_experts_implementation(experts_implementation)++    vocab_size = model_ref.config.get_text_config().vocab_size+    input_ids = torch.randint(0, vocab_size, (2, 64)).to(device)+    labels = torch.randint(0, vocab_size, (2, 64)).to(device)++    loss_ref = model_ref(input_ids, labels=labels).loss+    loss_ref.backward()++    loss_ep = model_ep(input_ids, labels=labels).loss+    loss_ep.backward()++    assert torch.allclose(loss_ref, loss_ep, atol=atol, rtol=rtol), (+        f"EP and non-EP model losses differ. "+        f"Non-EP loss: {loss_ref.item()}, EP loss: {loss_ep.item()}, "+        f"Diff: {(loss_ref - loss_ep).abs().item()}"+    )++    world_size = dist.get_world_size()+    expert_weights = set(_get_expert_weight_names(model_ep))+    failed_grads = {}+    for (name, param), (_, param_ep) in zip(model_ref.named_parameters(), model_ep.named_parameters()):+        if param.grad is None and param_ep.grad is None:+            continue+        if param.grad is None or param_ep.grad is None:+            failed_grads[name] = "no gradient on the " + ("EP" if param_ep.grad is None else "non-EP") + " model"+            continue+        local_grad = param_ep.grad.to_local() if isinstance(param_ep.grad, DTensor) else param_ep.grad+        expected = _reference_grad_shard(+            name, param.grad, local_grad, rank, world_size, model_ep._tp_plan, expert_weights+        ).cpu()+        if not torch.allclose(expected, local_grad.cpu(), atol=atol, rtol=rtol):+            failed_grads[name] = f"max diff = {(expected - local_grad.cpu()).abs().max().item()}"++    assert not failed_grads, f"Gradients differ for {len(failed_grads)} parameter(s):\n" + "\n".join(+        f"  {name}: {failure}" for name, failure in failed_grads.items()+    )++    dist.barrier()+++class ExpertParallelTesterMixin(ABC):+    """+    Mixin for expert parallel tests. Add to model test classes alongside `ModelTesterMixin` and+    `TensorParallelTesterMixin`, whose helpers (`_skip_if_not_supported`, `_get_tp_config`, `_get_tp_model_class`) and+    `tensor_parallel_*` settings it reuses: EP runs with `tp_size == ep_size == tensor_parallel_size`.+    """++    # ============================================================+    # Helper methods+    # ============================================================+    def _has_ep_plan(self) -> bool:+        """Check if model has an expert parallel plan defined."""+        config = self.model_tester.get_config().get_text_config()+        return hasattr(config, "base_model_ep_plan") and config.base_model_ep_plan is not None++    def _skip_if_ep_distributed_not_enabled(self):+        """Skip expensive distributed EP tests for models outside the allowlist."""+        config = self.model_tester.get_config()+        if EP_DISTRIBUTED_TEST_MODEL_TYPES is not None and config.model_type not in EP_DISTRIBUTED_TEST_MODEL_TYPES:+            self.skipTest(+                f"Expert parallel distributed tests are not enabled for model_type={config.model_type!r} "+                f"(enabled: {sorted(EP_DISTRIBUTED_TEST_MODEL_TYPES)}). "+                "Set EP_DISTRIBUTED_TEST_MODEL_TYPES = None to run all tests."+            )++    def _skip_if_ep_not_supported(self):+        """Check and skip the test if expert parallel is not supported for this model/environment."""+        self._skip_if_distributed_not_supported(parallelism="Expert")++        if not self._has_ep_plan():+            self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)")++        self._skip_if_ep_distributed_not_enabled()++    @parameterized.expand(+        list(+            product(+                [False, True],  # tie_word_embeddings+                ["eager", "grouped_mm", "batched_mm"],  # experts_implementation+            )+        )+    )+    @is_expert_parallel_test+    def test_ep_forward(self, tie_word_embeddings, experts_implementation):+        self._skip_if_ep_not_supported()++        config = self._get_tp_config(tie_word_embeddings=tie_word_embeddings)+        model_class = self._get_tp_model_class()+        atol = self.tensor_parallel_atol+        rtol = self.tensor_parallel_rtol++        with tempfile.TemporaryDirectory() as tmp_dir:+            set_seed(42)+            model = model_class(config)+            model.save_pretrained(tmp_dir, save_original_format=True)++            _init_distributed(tp=self.tensor_parallel_size)(_test_ep_forward_impl)(+                tmp_dir, model_class, atol, rtol, experts_implementation+            )++    @parameterized.expand([("eager",), ("grouped_mm",), ("batched_mm",)])+    @is_expert_parallel_test+    def test_ep_backward(self, experts_implementation):+        self._skip_if_ep_not_supported()++        config = self._get_tp_config()+        model_class = self._get_tp_model_class()+        atol = self.tensor_parallel_atol+        rtol = self.tensor_parallel_rtol++        with tempfile.TemporaryDirectory() as tmp_dir:+            set_seed(42)+            model = model_class(config)+            model.save_pretrained(tmp_dir, save_original_format=True)++            _init_distributed(tp=self.tensor_parallel_size)(_test_ep_backward_impl)(+                tmp_dir, model_class, atol, rtol, experts_implementation+            )

@@ -1,3 +1,4 @@+# Copyright 2026 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License.@@ -15,7 +16,6 @@ import socket import tempfile from abc import ABC, abstractmethod-from itertools import product  from parameterized import parameterized @@ -420,98 +420,6 @@ def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_t     dist.barrier()  -def _load_ep_and_reference_models(model_path, model_class):-    """Load EP model and non-EP reference model for comparison."""-    model_ep = model_class.from_pretrained(-        model_path,-        distributed_config=DistributedConfig(tp_size=dist.get_world_size(), ep_size=dist.get_world_size()),-    )-    dist.barrier()--    device = model_ep.device-    model_ref = model_class.from_pretrained(model_path)-    model_ref = model_ref.to(device)--    return model_ep, model_ref, device---def _test_ep_forward_impl(_rank, model_path, model_class, atol, rtol, experts_implementation):-    """Implementation for comparing EP and non-EP model outputs."""-    set_seed(0)--    model_ep, model_ref, device = _load_ep_and_reference_models(model_path, model_class)--    model_ep.eval()-    model_ref.eval()--    model_ep.set_experts_implementation(experts_implementation)-    model_ref.set_experts_implementation(experts_implementation)--    vocab_size = model_ref.config.get_text_config().vocab_size-    input_ids = torch.randint(0, vocab_size, (2, 64)).to(device)--    with torch.no_grad():-        logits_ref = model_ref(input_ids).logits-        logits_ep = model_ep(input_ids).logits--    diff = (logits_ref - logits_ep).abs()-    assert torch.allclose(logits_ref, logits_ep, atol=atol, rtol=rtol), (-        f"EP and non-EP model outputs differ. Max diff: {diff.max().item()} | Min diff: {diff.min().item()}"-    )--    dist.barrier()---def _test_ep_backward_impl(rank, model_path, model_class, atol, rtol, experts_implementation):-    """Implementation for comparing EP and non-EP model backward passes."""-    set_seed(0)--    model_ep, model_ref, device = _load_ep_and_reference_models(model_path, model_class)-    model_ep.train()-    model_ref.train()--    model_ep.set_experts_implementation(experts_implementation)-    model_ref.set_experts_implementation(experts_implementation)--    vocab_size = model_ref.config.get_text_config().vocab_size-    input_ids = torch.randint(0, vocab_size, (2, 64)).to(device)-    labels = torch.randint(0, vocab_size, (2, 64)).to(device)--    loss_ref = model_ref(input_ids, labels=labels).loss-    loss_ref.backward()--    loss_ep = model_ep(input_ids, labels=labels).loss-    loss_ep.backward()--    assert torch.allclose(loss_ref, loss_ep, atol=atol, rtol=rtol), (-        f"EP and non-EP model losses differ. "-        f"Non-EP loss: {loss_ref.item()}, EP loss: {loss_ep.item()}, "-        f"Diff: {(loss_ref - loss_ep).abs().item()}"-    )--    world_size = dist.get_world_size()-    expert_weights = set(_get_expert_weight_names(model_ep))-    failed_grads = {}-    for (name, param), (_, param_ep) in zip(model_ref.named_parameters(), model_ep.named_parameters()):-        if param.grad is None and param_ep.grad is None:-            continue-        if param.grad is None or param_ep.grad is None:-            failed_grads[name] = "no gradient on the " + ("EP" if param_ep.grad is None else "non-EP") + " model"-            continue-        local_grad = param_ep.grad.to_local() if isinstance(param_ep.grad, DTensor) else param_ep.grad-        expected = _reference_grad_shard(-            name, param.grad, local_grad, rank, world_size, model_ep._tp_plan, expert_weights-        ).cpu()-        if not torch.allclose(expected, local_grad.cpu(), atol=atol, rtol=rtol):-            failed_grads[name] = f"max diff = {(expected - local_grad.cpu()).abs().max().item()}"--    assert not failed_grads, f"Gradients differ for {len(failed_grads)} parameter(s):\n" + "\n".join(-        f"  {name}: {failure}" for name, failure in failed_grads.items()-    )--    dist.barrier()-- class TensorParallelTesterMixin(ABC):     """     Mixin for tensor parallel tests. Add to model test classes alongside ModelTesterMixin.@@ -539,18 +447,13 @@ def model_tester(self):     # ============================================================     # Helper methods     # ============================================================-    def _has_ep_plan(self) -> bool:-        """Check if model has an expert parallel plan defined."""-        config = self.model_tester.get_config().get_text_config()-        return hasattr(config, "base_model_ep_plan") and config.base_model_ep_plan is not None-     def _has_tp_plan(self) -> bool:         """Check if model has a tensor parallel plan defined."""         config = self.model_tester.get_config().get_text_config()         return hasattr(config, "base_model_tp_plan") and config.base_model_tp_plan is not None      def _skip_if_tp_distributed_not_enabled(self):-        """Skip expensive distributed TP/EP tests for models outside the allowlist."""+        """Skip expensive distributed TP tests for models outside the allowlist."""         config = self.model_tester.get_config()         if TP_DISTRIBUTED_TEST_MODEL_TYPES is not None and config.model_type not in TP_DISTRIBUTED_TEST_MODEL_TYPES:             self.skipTest(@@ -582,9 +485,8 @@ def _get_tp_config(self, tie_word_embeddings: bool | None = None):                 config.tie_word_embeddings = tie_word_embeddings         return config -    def _skip_if_not_supported(self, expert_parallel: bool = False):-        """Check and skip the test if tensor/expert parallel is not supported for this model/environment."""-        parallelism = "Expert" if expert_parallel else "Tensor"+    def _skip_if_distributed_not_supported(self, parallelism: str = "Tensor"):+        """Skip when the environment can't run the multi-process CPU tests (torch version, device, CPUs, model tester)."""         if not is_torch_greater_or_equal("2.9"):             self.skipTest(f"{parallelism} parallel tests require torch >= 2.9") @@ -602,10 +504,11 @@ def _skip_if_not_supported(self, expert_parallel: bool = False):         if self._get_tp_model_class() is None:             self.skipTest("Model tester has neither causal_lm_class nor conditional_generation_class") -        if expert_parallel:-            if not self._has_ep_plan():-                self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)")-        elif not self._has_tp_plan():+    def _skip_if_not_supported(self):+        """Check and skip the test if tensor parallel is not supported for this model/environment."""+        self._skip_if_distributed_not_supported()++        if not self._has_tp_plan():             self.skipTest("Model does not have a tensor parallel plan (base_model_tp_plan)")          self._skip_if_tp_distributed_not_enabled()@@ -726,48 +629,3 @@ def test_tp_generation_quantized(self):             _init_distributed(tp=self.tensor_parallel_size)(_test_tp_generation_quantized_impl)(                 tmp_dir, model_class, max_new_tokens             )--    @parameterized.expand(-        list(-            product(-                [False, True],  # tie_word_embeddings-                ["eager", "grouped_mm", "batched_mm"],  # experts_implementation-            )-        )-    )-    @is_tensor_parallel_test-    def test_ep_forward(self, tie_word_embeddings, experts_implementation):-        self._skip_if_not_supported(expert_parallel=True)--        config = self._get_tp_config(tie_word_embeddings=tie_word_embeddings)-        model_class = self._get_tp_model_class()-        atol = self.tensor_parallel_atol-        rtol = self.tensor_parallel_rtol--        with tempfile.TemporaryDirectory() as tmp_dir:-            set_seed(42)-            model = model_class(config)-            model.save_pretrained(tmp_dir, save_original_format=True)--            _init_distributed(tp=self.tensor_parallel_size)(_test_ep_forward_impl)(-                tmp_dir, model_class, atol, rtol, experts_implementation-            )--    @parameterized.expand([("eager",), ("grouped_mm",), ("batched_mm",)])-    @is_tensor_parallel_test-    def test_ep_backward(self, experts_implementation):-        self._skip_if_not_supported(expert_parallel=True)--        config = self._get_tp_config()-        model_class = self._get_tp_model_class()-        atol = self.tensor_parallel_atol-        rtol = self.tensor_parallel_rtol--        with tempfile.TemporaryDirectory() as tmp_dir:-            set_seed(42)-            model = model_class(config)-            model.save_pretrained(tmp_dir, save_original_format=True)--            _init_distributed(tp=self.tensor_parallel_size)(_test_ep_backward_impl)(-                tmp_dir, model_class, atol, rtol, experts_implementation-            )

@@ -1194,7 +1194,8 @@ def parse_commit_message(commit_message: str) -> dict[str, bool]:     # torch-light.     "tests_non_model": r"tests/(?!peft_integration/|conftest_tests/|repo_utils/)[^/]*?/test_.*\.py",     "tests_training_ci": r"tests/models/.*/test_modeling_.*",-    "tests_tensor_parallel_ci": r"(tests/models/.*/test_modeling_.*|tests/tensor_parallel(?:/test_tensor_parallel\.py)?)",+    "tests_tensor_parallel_ci": r"(tests/models/.*/test_modeling_.*|tests/tensor_parallel(?:/test_.*\.py)?)",+    "tests_expert_parallel_ci": r"(tests/models/.*/test_modeling_.*|tests/expert_parallel(?:/test_.*\.py)?)",     "tests_fsdp_ci": r"(tests/models/.*/test_modeling_.*|tests/test_fsdp_mixin\.py)",     "tests_peft_integration": r"tests/peft_integration/test_.*\.py", }