diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md
index f66d993495e6..c62eef296119 100644
--- a/docs/source/en/expert_parallelism.md
+++ b/docs/source/en/expert_parallelism.md
@@ -50,4 +50,30 @@ Launch your inference script with [torchrun](https://pytorch.org/docs/stable/ela
 torchrun --nproc-per-node 8 your_script.py
 ```
 
+## Combining with FSDP2
+
+Expert parallelism only shards the experts. Everything else (attention, embeddings, norms) and its optimizer state is replicated on every expert-parallel rank, which is what limits the model size you can train. Set `fsdp_size` together with `tp_size` to add [FSDP2](./fsdp) on a second mesh dimension.
+
+```py
+distributed_config = DistributedConfig(
+    tp_size=4,  # expert parallel size
+    fsdp_size=2,  # data parallel shards
+    enable_expert_parallel=True,
+)
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_config=distributed_config)
+```
+
+The model is loaded on a 2-D `(fsdp, tp)` device mesh, and `tp_size * fsdp_size` must equal the number of processes. The expert parallel plan shards the experts across `tp`, then FSDP2 shards every parameter, experts included, across `fsdp` and owns their gradient reduction. Each `fsdp` rank trains on its own part of the batch. Nothing else changes: train with the [`Trainer`] as usual (it computes the gradient norm across the two meshes and gives each mesh its own optimizer param group), and [`~Trainer.save_model`] gathers the sharded weights and writes a regular checkpoint.
+
+On 8 GPUs, full fine-tuning of Qwen3-30B-A3B in bf16 at sequence length 2048:
+
+| configuration | tokens/s/GPU | peak memory/GPU |
+|---|---|---|
+| `tp_size=8` | 3485 | 38.6 GB |
+| `tp_size=4, fsdp_size=2` | 2900 | 34.2 GB |
+| `tp_size=2, fsdp_size=4` | 2830 | 32.3 GB |
+
+> [!WARNING]
+> Resuming from a checkpoint is not supported yet for models sharded at load time, so the [`Trainer`] only accepts `save_only_model=True` or `save_strategy="no"` for them.
+
 [[autodoc]] DistributedConfig
diff --git a/docs/source/en/fsdp.md b/docs/source/en/fsdp.md
index e0727b77d897..dbbd629f22d1 100644
--- a/docs/source/en/fsdp.md
+++ b/docs/source/en/fsdp.md
@@ -121,6 +121,9 @@ TrainingArguments(
 </hfoption>
 </hfoptions>
 
+> [!TIP]
+> For mixture-of-experts models, `fsdp_size` can be combined with `tp_size` and `enable_expert_parallel=True` to shard the experts across one mesh dimension and everything else across the other. See [expert parallelism](./expert_parallelism#combining-with-fsdp2).
+
 ## Next steps
 
 - See [DDP](./ddp) for data-parallel training when your model fits on one GPU.
diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 580f46849c3c..ea9e8b958a65 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -976,6 +976,10 @@ def materialize_tensors(self) -> dict[str, list[torch.Tensor]]:
             # Add them to the new dictionary
             collected_tensors[key] = tensors
 
+        if any(len(tensors) == 0 for tensors in collected_tensors.values()):
+            # Uneven FSDP sharding left this rank an empty shard: nothing to load, and its pre-sharded empty local tensor is already correct
+            raise SkipParameters()
+
         return collected_tensors
 
     def was_used(self) -> bool:
diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py
index a18a6f9a1dda..e43e5e479281 100644
--- a/src/transformers/distributed/configuration_utils.py
+++ b/src/transformers/distributed/configuration_utils.py
@@ -73,11 +73,10 @@ def __post_init__(self):
         elif self.tp_size is None:
             self.tp_size = 1
 
-        if self.tp_size > 1 and self.fsdp_size > 1 and self.pp_size > 1:
+        if self.pp_size > 1 and (self.tp_size > 1 or self.fsdp_size > 1):
             raise ValueError(
-                "FSDP+TP+PP is not supported yet. "
-                "Use DistributedConfig(fsdp_size=N) or DistributedConfig(tp_size=N) or DistributedConfig(pp_size=N), not all three. "
-                "Only 1D support is available for now."
+                "Pipeline parallelism cannot be combined with tensor or FSDP parallelism yet. "
+                "Use DistributedConfig(pp_size=N) on its own, or DistributedConfig(tp_size=N, fsdp_size=M)."
             )
 
     @classmethod
diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py
index eb1bafd7355f..f2dd2d4510a3 100644
--- a/src/transformers/distributed/mixin.py
+++ b/src/transformers/distributed/mixin.py
@@ -56,6 +56,7 @@ class DistributedMixin:
     _tp_plan: dict[str, str] | None = None
     _ep_plan: dict[str, str] | None = None
     _tp_size = None
+    _fsdp_size = None
     _pp_plan: dict[str, tuple[str, str]] | None = None
     _fsdp_plan: dict[str, str] | None = None
 
@@ -163,17 +164,23 @@ def prepare_distribute_model(
                     f"is not equal to world_size ({world_size})"
                 )
 
-        if distributed_config.tp_size > 1:
-            if distributed_config.tp_plan is None:
-                distributed_config.tp_plan = "auto"
+        if distributed_config.tp_size > 1 and distributed_config.tp_plan is None:
+            distributed_config.tp_plan = "auto"
+
+        if distributed_config.fsdp_size > 1:
+            # Builds a 2-D (fsdp, tp) mesh when tensor/expert parallelism is also requested.
+            if device_mesh is not None:
+                raise ValueError(
+                    "`device_mesh` cannot be passed together with `fsdp_size > 1`: the mesh is built here."
+                )
+            device_map, device_mesh = initialize_fully_sharded_data_parallelism(distributed_config)
+        elif distributed_config.tp_size > 1:
             device_map, device_mesh = initialize_tensor_parallelism(
                 distributed_config.tp_plan,
                 tp_size=distributed_config.tp_size,
                 device_mesh=device_mesh,
                 device_map=device_map,
             )
-        elif distributed_config.fsdp_size > 1:
-            device_map, device_mesh = initialize_fully_sharded_data_parallelism(distributed_config)
         elif distributed_config.pp_size > 1:
             device_map, device_mesh = initialize_pipeline_parallelism(distributed_config)
 
@@ -190,18 +197,24 @@ def maybe_distribute_model(
         if device_mesh is not None:
             model.config.distributed_config = distributed_config
             model._device_mesh = device_mesh
+            # The Trainer mirrors these into accelerate's `ParallelismConfig`; without them accelerate
+            # sees unaccounted ranks and falls back to DDP, which rejects the DTensor parameters.
+            model._tp_size = distributed_config.tp_size
+            model._fsdp_size = distributed_config.fsdp_size
 
+            # Both may apply: the tensor/expert parallel plan shards across `tp` first, then FSDP2
+            # shards every parameter (the `tp`-sharded ones included) across `fsdp`.
             if distributed_config.tp_size > 1:
                 tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh
                 if isinstance(distributed_config.tp_plan, dict):
                     model.tp_plan = distributed_config.tp_plan
                 model = apply_tensor_parallelism(model, tp_mesh)
 
-            elif distributed_config.fsdp_size > 1:
+            if distributed_config.fsdp_size > 1:
                 fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh
                 model = apply_fully_sharded_data_parallelism(model, fsdp_mesh)
 
-            elif distributed_config.pp_size > 1:
+            if distributed_config.pp_size > 1:
                 pp_mesh = device_mesh["pp"] if device_mesh.ndim > 1 else device_mesh
                 model = apply_pipeline_parallelism(model, pp_mesh)
         return model
@@ -264,15 +277,9 @@ def gather_sharded_state_dict_for_save(
         if distributed_config is None:
             return state_dict
 
-        if distributed_config.tp_size > 1:
-            state_dict = gather_state_dict_for_save(
-                state_dict, self._tp_plan, self._device_mesh, distributed_config.tp_size
-            )
-            if not save_on_this_rank:
-                state_dict = {}
-            return state_dict
-
         if distributed_config.fsdp_size > 1:
+            # Also covers the 2-D (fsdp, tp) mesh: every parameter is FSDP-managed, and the full
+            # state dict is only materialized on rank 0.
             if not _is_torch_distributed_initialized():
                 raise ValueError(
                     "Saving an FSDP-wrapped model requires torch.distributed to be initialized. "
@@ -280,6 +287,14 @@ def gather_sharded_state_dict_for_save(
                 )
             return gather_full_state_dict(model_to_save)
 
+        if distributed_config.tp_size > 1:
+            state_dict = gather_state_dict_for_save(
+                state_dict, self._tp_plan, self._device_mesh, distributed_config.tp_size
+            )
+            if not save_on_this_rank:
+                state_dict = {}
+            return state_dict
+
         return state_dict
 
     def barrier_after_gathered_checkpoint_save(self, distributed_config: DistributedConfig | None) -> None:
diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index 3dc23a07a304..8a1a943939ee 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import contextlib
+import os
 import re
 
 from ..utils import logging
@@ -144,7 +145,7 @@ def context_around_forward(self, module, mesh):
     def transform_output_post_forward(self, module, output, mesh):
         return output
 
-    def install_forward(self, module, mesh):
+    def install_forward(self, module, mesh, *, is_expert_parallel=False):
         """Install pre / around / post transforms by replacing module.forward."""
         original_forward = module.forward
 
@@ -344,7 +345,7 @@ class ReplicatedWithGradAllReduce(TensorParallelLayer):
     summed across the mesh.
     """
 
-    def install_forward(self, module, mesh):
+    def install_forward(self, module, mesh, *, is_expert_parallel=False):
         # A module hook rather than `param.register_hook`: params are replaced during weight
         # loading, which happens after TP is applied, and would drop a param-level hook.
         def _all_reduce_grads(mod, grad_input, grad_output):
@@ -415,7 +416,7 @@ def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = True):
         self.sequence_dim = sequence_dim
         self.use_local_output = use_local_output
 
-    def install_forward(self, module, mesh):
+    def install_forward(self, module, mesh, *, is_expert_parallel=False):
         # Replicate the module's params (LayerNorm/RMSNorm ones-init → from_local is safe).
         for p_name, p in list(module.named_parameters(recurse=False)):
             module.register_parameter(
@@ -606,6 +607,7 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_
 
     def install_forward(self, module, mesh, *, is_expert_parallel=False):
         """Install the transforms but pass `is_expert_parallel` in the forward call."""
+        module.is_expert_parallel = is_expert_parallel
         original_forward = module.forward
         output_source = (
             Partial()
@@ -617,6 +619,23 @@ def install_forward(self, module, mesh, *, is_expert_parallel=False):
         )
 
         def tp_forward(*args, **kwargs):
+            if os.environ.get("HF_EP_DISPATCH") == "1":
+                from ..integrations.moe import dispatch_experts_forward
+
+                hidden_states, top_k_index, top_k_weights, *rest = args
+                if isinstance(hidden_states, DTensor):
+                    hidden_states = hidden_states.to_local()
+                ep_mesh = mesh if mesh.ndim == 1 else mesh["tp"]
+                with self.context_around_forward(module, mesh):
+                    return dispatch_experts_forward(
+                        module,
+                        hidden_states,
+                        top_k_index,
+                        top_k_weights,
+                        ep_mesh.get_group(),
+                        ep_mesh.get_local_rank(),
+                        ep_mesh.size(),
+                    )
             args, kwargs = self.transform_inputs_pre_forward(
                 module, args, kwargs, mesh, is_expert_parallel=is_expert_parallel
             )
@@ -696,6 +715,10 @@ class EpRouterParallel(TensorParallelLayer):
     """
 
     def transform_output_post_forward(self, module, output, mesh):
+        if os.environ.get("HF_EP_DISPATCH") == "1":
+            # Token-dispatch prototype: keep global expert ids and scores; the experts forward
+            # routes tokens to their owners with an all-to-all instead of masking.
+            return output
         ep_rank, ep_size = mesh.get_local_rank(), mesh.size()
         num_experts = getattr(module, "num_experts", None)
         if num_experts is None:
@@ -709,6 +732,11 @@ def transform_output_post_forward(self, module, output, mesh):
         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)
@@ -816,7 +844,9 @@ def apply_tensor_parallelism(model, tp_mesh):
                 # MLA needs to know the qk_rope_head_dim to split the projection output into KV and RoPE parts.
                 # TODO: Store qk_rope_head_dim on MLA projection modules when the models initialize them.
                 module.config = model.config.get_text_config()
-            ALL_PARALLEL_STYLES[style_name].install_forward(module, tp_mesh)
+            ALL_PARALLEL_STYLES[style_name].install_forward(
+                module, tp_mesh, is_expert_parallel=model.config.distributed_config.enable_expert_parallel
+            )
         module._is_hooked = True
 
     return model
diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py
index 49f0f1c41814..e28e62525162 100644
--- a/src/transformers/distributed/utils.py
+++ b/src/transformers/distributed/utils.py
@@ -59,6 +59,37 @@ def _get_torch_distributed_world_size() -> int:
     return torch.distributed.get_world_size()
 
 
+def prefetch_checkpoint_shards(checkpoint_files: list[str]) -> None:
+    """Warm the page cache for the checkpoint shards before the per-tensor loading pass, opt-in via
+    `HF_SHARD_PREFETCH=<read threads per rank>`.
+
+    The per-tensor read pattern of sharded loading reads a network filesystem at well under 1 GiB/s
+    while large sequential reads sustain many times that; warming the page cache first makes the
+    actual load run at memory speed. Local ranks split the shard list between them (every node needs
+    the full checkpoint cached, since every rank slices tensors from all shards).
+    """
+    prefetch_threads = int(os.environ.get("HF_SHARD_PREFETCH", "0"))
+    if not checkpoint_files or not prefetch_threads:
+        return
+    import time
+    from concurrent.futures import ThreadPoolExecutor
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    local_world = int(os.environ.get("LOCAL_WORLD_SIZE", "1"))
+
+    def _warm(path, bufsize=16 * 2**20):
+        with open(path, "rb", buffering=0) as f:
+            while f.read(bufsize):
+                pass
+
+    prefetch_start = time.time()
+    with ThreadPoolExecutor(max_workers=prefetch_threads) as pool:
+        list(pool.map(_warm, checkpoint_files[local_rank::local_world]))
+    if _is_torch_distributed_initialized():
+        torch.distributed.barrier()
+    logger.warning_once(f"Prefetched {len(checkpoint_files)} checkpoint shards in {time.time() - prefetch_start:.0f}s")
+
+
 def is_local_dist_rank_0() -> bool:
     return _is_torch_distributed_initialized() and int(os.environ.get("LOCAL_RANK", "-1")) == 0
 
@@ -193,11 +224,17 @@ def initialize_fully_sharded_data_parallelism(distributed_config: DistributedCon
         device_map = torch.device(device_type)
 
     fsdp_size = distributed_config.fsdp_size
+    tp_size = distributed_config.tp_size
 
+    # `fsdp` is the outer dimension so that the `tp` ranks of a group are contiguous, which is what
+    # the expert all-to-all and the TP collectives want.
     dims, names = [], []
     if fsdp_size > 1:
         dims.append(fsdp_size)
         names.append("fsdp")
+    if tp_size > 1:
+        dims.append(tp_size)
+        names.append("tp")
 
     # Build the N-dimensional device mesh
     mesh = torch.distributed.init_device_mesh(device_type, tuple(dims), mesh_dim_names=tuple(names))
diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py
index ab266873f821..f7334f20b0f2 100644
--- a/src/transformers/integrations/deepgemm.py
+++ b/src/transformers/integrations/deepgemm.py
@@ -504,6 +504,7 @@ def _dispatch_routed_input(
     num_experts: int,
     m_alignment: int,
     use_psum_layout: bool,
+    is_expert_parallel: bool = False,
 ) -> tuple:
     """Sort tokens by expert id and build the M-grouped padded layout.
 
@@ -533,8 +534,10 @@ def _dispatch_routed_input(
     # keeps any per-row gather (e.g. bias) in-bounds — bias added at sentinel positions falls
     # in rows the kernel skips, so harmless. Safe to mutate now: the layout was built from the
     # unclamped tensor and nothing downstream needs the sentinel info from `expert_ids_g` itself.
-    sentinel_mask = (expert_ids_g >= num_experts).unsqueeze(-1)
-    expert_ids_g.clamp_(max=num_experts - 1)
+    sentinel_mask = None
+    if is_expert_parallel:
+        sentinel_mask = (expert_ids_g >= num_experts).unsqueeze(-1)
+        expert_ids_g.clamp_(max=num_experts - 1)
     return (
         sorted_hidden_states_g,
         sample_weights_g,
@@ -550,7 +553,7 @@ def _dispatch_routed_input(
 def _combine_routed_output(
     out_padded: torch.Tensor,
     sorted_weights: torch.Tensor,
-    sentinel_mask: torch.Tensor,
+    sentinel_mask: torch.Tensor | None,
     perm: torch.Tensor,
     sorted_to_padded: torch.Tensor,
     num_tokens: int,
@@ -563,7 +566,8 @@ def _combine_routed_output(
     weighted = out * sorted_weights.to(out.dtype).unsqueeze(-1)
     # Sentinel rows past the valid expert blocks may carry NaN from allocator
     # reuse (`0 * NaN = NaN`); zero them so the top-k reduction stays finite.
-    weighted.masked_fill_(sentinel_mask, 0.0)
+    if sentinel_mask is not None:
+        weighted.masked_fill_(sentinel_mask, 0.0)
     inv_perm = torch.empty_like(perm)
     inv_perm[perm] = torch.arange(perm.size(0), device=out.device)
     # Deterministic reshape+sum (index_add_ with duplicates is non-deterministic on CUDA).
@@ -646,7 +650,13 @@ def deepgemm_bf16_experts_forward(
         grouped_layout,
         total_padded_rows,
     ) = _dispatch_routed_input(
-        hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100()
+        hidden_states,
+        top_k_index,
+        top_k_weights,
+        self.num_experts,
+        deepgemm.m_alignment,
+        is_sm100(),
+        is_expert_parallel=self.is_expert_parallel,
     )
 
     weight_up = self.gate_up_proj if self.has_gate else self.up_proj
@@ -732,7 +742,13 @@ def deepgemm_fp8_fp4_experts_forward(
         grouped_layout,
         total_padded_rows,
     ) = _dispatch_routed_input(
-        hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100()
+        hidden_states,
+        top_k_index,
+        top_k_weights,
+        self.num_experts,
+        deepgemm.m_alignment,
+        is_sm100(),
+        is_expert_parallel=self.is_expert_parallel,
     )
     sf_recipe = (1, 1, cast_kwargs["gran_k"]) if cast_kwargs.get("use_packed_ue8m0") else None
 
diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py
index 412fa42630e0..9089241fdf75 100644
--- a/src/transformers/integrations/finegrained_fp8.py
+++ b/src/transformers/integrations/finegrained_fp8.py
@@ -463,7 +463,7 @@ def fp8_batched_mm_experts_forward(
     # EP sentinel handling: leave `expert_ids` unclamped — the batched kernel early-returns on
     # `expert_id >= NUM_EXPERTS`, leaving sentinel output rows uninitialized. The post-mask below
     # zeroes them before the per-token reduction so `uninit * 0 = NaN` can't poison the sum.
-    sentinel_mask = (expert_ids >= self.num_experts).unsqueeze(-1)
+    sentinel_mask = (expert_ids >= self.num_experts).unsqueeze(-1) if self.is_expert_parallel else None
 
     weight_up = self.gate_up_proj if self.has_gate else self.up_proj
     weight_scale_up = self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv
@@ -501,7 +501,8 @@ def fp8_batched_mm_experts_forward(
 
     # Post-mask sentinel rows: kernel left them uninitialized, so zero them out
     # before the reduction below (uninit may be NaN; NaN * 0 = NaN).
-    weighted_out.masked_fill_(sentinel_mask, 0.0)
+    if self.is_expert_parallel:
+        weighted_out.masked_fill_(sentinel_mask, 0.0)
 
     # Accumulate results using deterministic reshape+sum instead of index_add_
     # (index_add_ with duplicate indices is non-deterministic on CUDA due to atomicAdd)
@@ -551,7 +552,7 @@ def fp8_grouped_mm_experts_forward(
     # valid rows, so sentinel-tail `proj_out` rows are uninit; without the post-mask below,
     # `proj_out[sentinel] * 0 = NaN * 0 = NaN` would poison the per-token reduction. FP8
     # quantized weights are inference-only, so no bwd pre-mask is needed.
-    sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
+    sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1) if self.is_expert_parallel else None
 
     weight_up = self.gate_up_proj if self.has_gate else self.up_proj
     weight_scale_up = self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv
@@ -590,7 +591,8 @@ def fp8_grouped_mm_experts_forward(
     weighted_out = proj_out * sample_weights_g.to(proj_out.dtype).unsqueeze(-1)  # (S, hidden_dim)
 
     # Post-mask (fwd path).
-    weighted_out.masked_fill_(sentinel_mask, 0.0)
+    if self.is_expert_parallel:
+        weighted_out.masked_fill_(sentinel_mask, 0.0)
 
     # Restore original order
     inv_perm = torch.empty_like(perm)
diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py
index b9b288f4a696..3145d21616db 100644
--- a/src/transformers/integrations/moe.py
+++ b/src/transformers/integrations/moe.py
@@ -125,11 +125,12 @@ def batched_mm_experts_forward(
     sample_weights = top_k_weights.reshape(-1)  # (S,)
     expert_ids = top_k_index.reshape(-1)  # (S,)
 
-    # Clamp EP sentinels so `gate_up_proj[expert_ids]` stays in-bounds. Routing weights are already
-    # zero at sentinel slots (RouterParallel masks them at dispatch), so the weighted mul drops
-    # those contributions — we pay the wasted GEMM compute because batched_mm has no offset to skip.
-    # Out-of-place to avoid mutating the caller's routing tensor (a contiguous `reshape(-1)` aliases it).
-    expert_ids = expert_ids.clamp(0, self.num_experts - 1)
+    if self.is_expert_parallel:
+        # Clamp EP sentinels so `gate_up_proj[expert_ids]` stays in-bounds. Routing weights are already
+        # zero at sentinel slots (RouterParallel masks them at dispatch), so the weighted mul drops
+        # those contributions — we pay the wasted GEMM compute because batched_mm has no offset to skip.
+        # Out-of-place to avoid mutating the caller's routing tensor (a contiguous `reshape(-1)` aliases it).
+        expert_ids = expert_ids.clamp(0, self.num_experts - 1)
 
     # Select gate_up or just up projection weights and biases
     if self.has_gate:
@@ -374,6 +375,181 @@ def _grouped_linear(
     return out
 
 
+def dispatch_experts_forward(
+    self,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    ep_group,
+    ep_rank: int,
+    ep_size: int,
+) -> torch.Tensor:
+    """Prototype token-dispatch expert parallelism (HF_EP_DISPATCH=1).
+
+    v3-marker.
+
+    Each rank keeps a 1/ep_size slice of the (replicated) batch, sends every selected
+    token-expert pair to the rank owning that expert with an all-to-all, computes its local
+    experts on what it receives, sends the results back, combines with the local routing
+    weights, and all-gathers the finished slices. Routing arrives unmasked (global expert ids).
+    """
+    import os
+
+    import torch.distributed as dist
+
+    class _AllToAll(torch.autograd.Function):
+        """Autograd all_to_all_single with variable splits (explicit, self-contained)."""
+
+        @staticmethod
+        def forward(ctx, x, out_sizes, in_sizes, group):
+            ctx.group, ctx.out_sizes, ctx.in_sizes = group, out_sizes, in_sizes
+            out = x.new_empty(sum(out_sizes), *x.shape[1:])
+            dist.all_to_all_single(
+                out, x.contiguous(), output_split_sizes=out_sizes, input_split_sizes=in_sizes, group=group
+            )
+            return out
+
+        @staticmethod
+        def backward(ctx, grad):
+            back = grad.new_empty(sum(ctx.in_sizes), *grad.shape[1:])
+            dist.all_to_all_single(
+                back,
+                grad.contiguous(),
+                output_split_sizes=ctx.in_sizes,
+                input_split_sizes=ctx.out_sizes,
+                group=ctx.group,
+            )
+            return back, None, None, None
+
+    class _AllGatherCat(torch.autograd.Function):
+        """Autograd all_gather (concatenated); backward returns this rank's slice of the gradient, summed across ranks."""
+
+        @staticmethod
+        def forward(ctx, x, rank, world, group):
+            ctx.group, ctx.rank, ctx.per = group, rank, x.size(0)
+            out = x.new_empty(x.size(0) * world, *x.shape[1:])
+            dist.all_gather_into_tensor(out, x.contiguous(), group=group)
+            return out
+
+        @staticmethod
+        def backward(ctx, grad):
+            grad = grad.contiguous()
+            dist.all_reduce(grad, group=ctx.group)
+            return grad[ctx.rank * ctx.per : (ctx.rank + 1) * ctx.per], None, None, None
+
+    num_tokens = hidden_states.size(0)
+    hidden_dim = hidden_states.size(-1)
+    num_top_k = top_k_index.size(-1)
+    num_local_experts = self.num_experts  # RouterParallel sets this to experts // ep_size
+
+    # HF_EP_TRUNK_DP=1: every EP rank already carries a *different* microbatch (the dense trunk is
+    # data-parallel), so there is nothing to slice, no replicated gradient to sum at the boundary,
+    # and no gather at the end — dispatch, compute, combine, return. The caller is responsible for
+    # averaging the ep-replicated (non-expert) parameter gradients across the EP group and scaling
+    # expert gradients by 1/ep_size after backward, mirroring data parallelism.
+    trunk_dp = os.environ.get("HF_EP_TRUNK_DP") == "1"
+
+    # Each rank's backward through the dispatch subgraph produces d_hidden only at its own slice's
+    # positions; sum across the group so every rank gets the complete expert-path gradient (the
+    # residual path is replicated and flows outside this op).
+    class _SumGradAcrossEp(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, x, group):
+            ctx.group = group
+            return x
+
+        @staticmethod
+        def backward(ctx, grad):
+            grad = grad.contiguous()
+            import torch.distributed as _dist
+
+            _dist.all_reduce(grad, group=ctx.group)
+            return grad, None
+
+    if trunk_dp:
+        my_tokens = hidden_states
+        my_index = top_k_index.reshape(-1)  # (s*K,) global expert ids
+        my_weights = top_k_weights.reshape(-1)
+    else:
+        hidden_states = _SumGradAcrossEp.apply(hidden_states, ep_group)
+        top_k_weights = _SumGradAcrossEp.apply(top_k_weights, ep_group)
+
+        # This rank's contiguous token slice (pad so every rank has the same slice length).
+        per_rank = (num_tokens + ep_size - 1) // ep_size
+        start, end = ep_rank * per_rank, min((ep_rank + 1) * per_rank, num_tokens)
+        my_tokens = hidden_states[start:end]
+        my_index = top_k_index[start:end].reshape(-1)  # (s*K,) global expert ids
+        my_weights = top_k_weights[start:end].reshape(-1)
+
+    # Group the selected pairs by destination rank.
+    owner = my_index // num_local_experts
+    order = torch.argsort(owner, stable=True)
+    send_tokens = my_tokens.repeat_interleave(num_top_k, dim=0)[order]
+    send_local_ids = (my_index % num_local_experts)[order]
+    send_counts = torch.bincount(owner, minlength=ep_size)
+    recv_counts = torch.empty_like(send_counts)
+    dist.all_to_all_single(recv_counts, send_counts, group=ep_group)
+    send_sizes = send_counts.tolist()
+    recv_sizes = recv_counts.tolist()
+
+    # Exchange the tokens (autograd) and their local expert ids (metadata, no grad).
+    recv_tokens = _AllToAll.apply(send_tokens, recv_sizes, send_sizes, ep_group)
+    recv_local_ids = torch.empty(sum(recv_sizes), dtype=send_local_ids.dtype, device=send_local_ids.device)
+    dist.all_to_all_single(
+        recv_local_ids, send_local_ids, output_split_sizes=recv_sizes, input_split_sizes=send_sizes, group=ep_group
+    )
+
+    # Local expert compute: sort by local expert, grouped GEMM, unsort.
+    ids_sorted, perm = torch.sort(recv_local_ids, stable=True)
+    x = recv_tokens[perm]
+    counts = torch.bincount(ids_sorted, minlength=num_local_experts)
+    offsets = torch.cumsum(counts, dim=0, dtype=torch.int32)
+    if self.has_gate:
+        proj = _grouped_linear(x, self.gate_up_proj, offsets, bias=None, is_transposed=self.is_transposed)
+        proj = self._apply_gate(proj)
+    else:
+        proj = _grouped_linear(x, self.up_proj, offsets, bias=None, is_transposed=self.is_transposed)
+        proj = self.act_fn(proj)
+    proj = _grouped_linear(proj, self.down_proj, offsets, bias=None, is_transposed=self.is_transposed)
+    inv_perm = torch.empty_like(perm)
+    inv_perm[perm] = torch.arange(perm.numel(), device=perm.device)
+    out_unsorted = proj[inv_perm]
+
+    # Send results back to the owning ranks of the tokens and combine.
+    back = _AllToAll.apply(out_unsorted, send_sizes, recv_sizes, ep_group)
+    combined = torch.zeros(my_tokens.size(0) * num_top_k, hidden_dim, device=back.device, dtype=back.dtype)
+    combined[order] = back
+    combined = combined * my_weights.unsqueeze(-1)
+    my_out = combined.view(-1, num_top_k, hidden_dim).sum(dim=1)
+
+    if trunk_dp:
+        return my_out.to(hidden_states.dtype)
+
+    # The dense trunk (and loss) is replicated across the EP group, so the all-gather backward sums
+    # ep_size identical gradient contributions per slice; pre-scale the gradient (forward unchanged).
+    class _ScaleGrad(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, x, scale):
+            ctx.scale = scale
+            return x
+
+        @staticmethod
+        def backward(ctx, grad):
+            return grad * ctx.scale, None
+
+    my_out = _ScaleGrad.apply(my_out, 1.0 / ep_size)
+
+    # Pad to the common slice length and all-gather the slices back to the full batch.
+    if my_out.size(0) < per_rank:
+        my_out = torch.cat(
+            [my_out, torch.zeros(per_rank - my_out.size(0), hidden_dim, device=my_out.device, dtype=my_out.dtype)]
+        )
+    # Gather all slices for the full-batch output. The loss is replicated across the group, so the
+    # gather backward sums ep_size identical gradients per slice; the _ScaleGrad above corrects it.
+    gathered = _AllGatherCat.apply(my_out, ep_rank, ep_size, ep_group)
+    return gathered[:num_tokens].to(hidden_states.dtype)
+
+
 def grouped_mm_experts_forward(
     self: torch.nn.Module,
     hidden_states: torch.Tensor,
@@ -417,8 +593,10 @@ def grouped_mm_experts_forward(
     # In-place clamp on `expert_ids_g` keeps the per-row bias gather in-bounds (bias added at
     # sentinel positions falls in rows the kernel skips, so harmless). Safe to mutate now —
     # nothing downstream needs the sentinel info from `expert_ids_g` itself.
-    sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
-    expert_ids_g.clamp_(max=self.num_experts - 1)
+    sentinel_mask = None
+    if self.is_expert_parallel:
+        sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
+        expert_ids_g.clamp_(max=self.num_experts - 1)
 
     # Select expert weights and biases
     # NOTE: We keep all experts here and rely on offsets to target the active ones.
@@ -434,13 +612,18 @@ def grouped_mm_experts_forward(
         selected_biases = self.up_proj_bias[expert_ids_g] if self.has_bias else None
 
     # Pre-mask (bwd path).
-    selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0)
+    if self.is_expert_parallel:
+        selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0)
 
     # --- Up projection per expert (grouped) ---
     proj_out = _grouped_linear(
         selected_hidden_states_g, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed
     )  # (S, 2 * intermediate_dim) or  (S, intermediate_dim) depending on whether we have gating
 
+    if self.is_expert_parallel:
+        # Zero the sentinel-tail rows the kernel left uninitialized (fwd output and bwd `d_input`).
+        proj_out = proj_out.masked_fill(sentinel_mask, 0.0)
+
     # Apply gating or activation
     if self.has_gate:
         # for gated experts we apply the custom/default gating mechanism
@@ -458,12 +641,13 @@ def grouped_mm_experts_forward(
         proj_out, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed
     )  # (S, hidden_dim)
 
+    if self.is_expert_parallel:
+        # Same: zero the uninitialized sentinel-tail rows.
+        proj_out = proj_out.masked_fill(sentinel_mask, 0.0)
+
     # Apply routing weights
     weighted_out = proj_out * sample_weights_g.unsqueeze(-1)  # (S, hidden_dim)
 
-    # Post-mask (fwd path).
-    weighted_out.masked_fill_(sentinel_mask, 0.0)
-
     # Restore original order
     inv_perm = torch.empty_like(perm)
     inv_perm[perm] = torch.arange(perm.size(0), device=device)
@@ -563,6 +747,7 @@ def __init__(self, config, *args, **kwargs):
             self.has_bias = has_bias
             self.is_transposed = is_transposed
             self.is_concatenated = is_concatenated
+            self.is_expert_parallel = False
 
         @wraps(original_forward)
         def forward(self, *args, **kwargs):
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 3bbd2b3b46d5..8be8dff3edf1 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -43,7 +43,7 @@
 from torch.distributions import constraints
 from torch.utils.checkpoint import checkpoint
 
-from transformers.distributed.utils import is_dtensor
+from transformers.distributed.utils import is_dtensor, prefetch_checkpoint_shards
 
 from . import initialization as init
 from .configuration_utils import PreTrainedConfig
@@ -4360,6 +4360,8 @@ def _load_pretrained_model(
         # Model's definition arriving here is final (TP hooks added, quantized layers replaces)
         expected_keys = list(model.state_dict().keys()) if expected_keys is None else expected_keys
 
+        prefetch_checkpoint_shards(checkpoint_files)
+
         if logger.level >= logging.WARNING:
             verify_tp_plan(expected_keys, getattr(model, "_tp_plan", None))
 
@@ -4600,6 +4602,14 @@ def tp_size(self):
         # if None, the model didn't undergo tensor parallel sharding
         return self._tp_size
 
+    @property
+    def fsdp_size(self):
+        """
+        Returns the model's FSDP sharding degree.
+        """
+        # if None, the model didn't undergo FSDP sharding
+        return self._fsdp_size
+
     @property
     def supports_pp_plan(self):
         # Check if model has a PP plan
@@ -4757,6 +4767,16 @@ def _initialize_missing_keys(self, is_quantized: bool) -> None:
                     pass  # may happen when handling pre-quantized weights
             self._is_hf_initialized = True
 
+        if getattr(self, "_device_mesh", None) is not None:
+            # Empty local shards have nothing to initialize; without the mark, running _init_weights on them issues collectives the other ranks never join (hang)
+            import itertools
+
+            from torch.distributed.tensor import DTensor
+
+            for param_or_buffer in itertools.chain(self.parameters(), self.buffers()):
+                if isinstance(param_or_buffer, DTensor) and param_or_buffer._local_tensor.numel() == 0:
+                    param_or_buffer._is_hf_initialized = True
+
         # This will only initialize submodules that are not marked as initialized by the line above.
         if is_deepspeed_zero3_enabled() and not is_quantized:
             import deepspeed
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 7a5005146858..b08ba795bfda 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -28,6 +28,7 @@
 import tempfile
 import time
 import warnings
+from collections import defaultdict
 from collections.abc import Callable, Iterator, Mapping
 from functools import partial
 from pathlib import Path
@@ -96,6 +97,7 @@
     _OPTIMIZER_HANDLERS,
     OptimizerContext,
     _parse_optim_args,
+    has_mixed_dtensor,
     is_optimizer_factory,
 )
 from .trainer_pt_utils import (
@@ -473,6 +475,9 @@ def __init__(
             or self.is_fsdp_xla_enabled
             or self.is_fsdp_enabled
             or is_sagemaker_mp_enabled()
+            # Sharded at load time (`DistributedConfig`): the model manages its own placement, and
+            # `.to()` on FSDP2-managed (possibly CPU-offloaded) parameters raises in `_apply`.
+            or getattr(model, "_device_mesh", None) is not None
         ):
             self.place_model_on_device = False
         else:
@@ -610,6 +615,17 @@ def __init__(
         self._train_batch_size = args.train_batch_size
         # Guards one-time LR scheduler creation in create_optimizer_and_scheduler
         self._created_lr_scheduler = False
+        # Resolved lazily at the first gradient clip; see `_has_mixed_mesh_grads`.
+        self._mixed_mesh_grads: bool | None = None
+        if (
+            getattr(model, "_device_mesh", None) is not None
+            and args.save_strategy != SaveStrategy.NO
+            and not args.save_only_model
+        ):
+            raise ValueError(
+                "Resuming is not supported for models sharded at load time (`DistributedConfig`), so their "
+                "optimizer state cannot be checkpointed. Pass `save_only_model=True` or `save_strategy='no'`."
+            )
 
         self.control = self.callback_handler.on_init_end(self.args, self.state, self.control)
 
@@ -746,17 +762,21 @@ def _build_accelerator_args(self, **kwargs) -> dict[str, Any]:
                 )
             args["parallelism_config"] = self.args.parallelism_config
 
-        if getattr(self.model, "tp_size", None) is not None and self.model.tp_size > 1:
-            if self.args.parallelism_config is None:
-                if is_accelerate_available("1.12.0"):
-                    if self.args.parallelism_config is None:
-                        from accelerate import ParallelismConfig
-
-                        args["parallelism_config"] = ParallelismConfig(tp_size=self.model.tp_size)
-                else:
-                    raise ValueError("Requires accelerate>1.12.0 to use Tensor Parallelism.")
-            elif args["parallelism_config"].tp_size != self.model.tp_size:
-                args["parallelism_config"].tp_size = self.model.tp_size
+        model_tp_size = getattr(self.model, "tp_size", None) or 1
+        model_fsdp_size = getattr(self.model, "fsdp_size", None) or 1
+        if model_tp_size > 1:
+            # Sharded at load time (tensor/expert parallelism, optionally with FSDP2 on a second mesh
+            # dimension): accelerate has to know both sizes, or it sees unaccounted ranks and wraps
+            # the DTensor model in DDP, which raises.
+            if not is_accelerate_available("1.12.0"):
+                raise ValueError("Requires accelerate>1.12.0 to use Tensor Parallelism.")
+            if args.get("parallelism_config") is None:
+                from accelerate import ParallelismConfig
+
+                args["parallelism_config"] = ParallelismConfig(tp_size=model_tp_size, dp_shard_size=model_fsdp_size)
+            else:
+                args["parallelism_config"].tp_size = model_tp_size
+                args["parallelism_config"].dp_shard_size = model_fsdp_size
 
         if is_accelerate_available("1.2.0"):
             # it we don't have the correct version, we will rely on env var instead that were set in TrainingArguments
@@ -1252,6 +1272,18 @@ def create_optimizer(self, model=None) -> torch.optim.Optimizer:
                     "weight_decay": 0.0,
                 },
             ]
+            if has_mixed_dtensor(p for group in optimizer_grouped_parameters for p in group["params"]):
+                # Parameters on different device meshes (expert parallelism, alone or with FSDP2 on a 2-D
+                # mesh) cannot share one fused/foreach kernel call: give each mesh its own param group.
+                from torch.distributed.tensor import DTensor
+
+                split_groups = []
+                for group in optimizer_grouped_parameters:
+                    by_mesh = defaultdict(list)
+                    for p in group["params"]:
+                        by_mesh[p.device_mesh if isinstance(p, DTensor) else None].append(p)
+                    split_groups.extend({**group, "params": params} for params in by_mesh.values())
+                optimizer_grouped_parameters = split_groups
 
             if self.optimizer_cls_and_kwargs is not None:
                 optimizer_cls, optimizer_kwargs = self.optimizer_cls_and_kwargs
@@ -2613,17 +2645,53 @@ def _track_num_input_tokens(self, inputs):
         input_tokens = torch.as_tensor(input_tokens, device=self.args.device, dtype=torch.int64)
         self.state.num_input_tokens_seen += self.accelerator.gather(input_tokens).sum().item()
 
+    def _mixed_mesh_grad_norm(self, model, max_norm):
+        """
+        Gradient norm (and clip) when the gradients live on different device meshes, which `clip_grad_norm_` cannot
+        span: one norm per mesh, each already reduced over its own mesh.
+        """
+        from torch.distributed.tensor import DTensor
+        from torch.nn.utils import clip_grads_with_norm_, get_total_norm
+
+        params_by_mesh = defaultdict(list)
+        for param in model.parameters():
+            if param.grad is not None:
+                params_by_mesh[param.grad.device_mesh if isinstance(param.grad, DTensor) else None].append(param)
+
+        norms = []
+        for params in params_by_mesh.values():
+            norm = get_total_norm([p.grad for p in params])
+            norms.append(norm.full_tensor() if isinstance(norm, DTensor) else norm)
+        total_norm = torch.linalg.vector_norm(torch.stack(norms))
+
+        if max_norm != float("inf"):
+            for params in params_by_mesh.values():
+                clip_grads_with_norm_(params, max_norm, total_norm)
+        return total_norm
+
+    def _has_mixed_mesh_grads(self, model) -> bool:
+        # Static for the life of the run (sharding never changes after setup), so scan the
+        # parameters only on the first call.
+        if self._mixed_mesh_grads is None:
+            self._mixed_mesh_grads = has_mixed_dtensor(p.grad for p in model.parameters() if p.grad is not None)
+        return self._mixed_mesh_grads
+
     def _clip_grad_norm(self, model):
         """Clip gradients to max_grad_norm. Returns the pre-clip gradient norm."""
         if is_sagemaker_mp_enabled() and self.args.fp16:
             return self.optimizer.clip_master_grads(self.args.max_grad_norm)
+        if self._has_mixed_mesh_grads(model):
+            return self._mixed_mesh_grad_norm(model, self.args.max_grad_norm)
         return self.accelerator.clip_grad_norm_(model.parameters(), self.args.max_grad_norm)
 
     def _get_grad_norm(self, model, grad_norm=None):
         """Return the gradient norm as a Python float."""
         if grad_norm is None:
             # Compute norm without clipping (inf means no actual clipping happens)
-            grad_norm = self.accelerator.clip_grad_norm_(model.parameters(), float("inf"))
+            if self._has_mixed_mesh_grads(model):
+                grad_norm = self._mixed_mesh_grad_norm(model, float("inf"))
+            else:
+                grad_norm = self.accelerator.clip_grad_norm_(model.parameters(), float("inf"))
 
         if self.accelerator.distributed_type == DistributedType.DEEPSPEED:
             if hasattr(grad_norm, "item"):
@@ -3925,6 +3993,12 @@ def save_model(self, output_dir: str | None = None, _internal_call: bool = False
                 remove_dummy_checkpoint(self.args.should_save, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME])
                 self.model_wrapped.save_checkpoint(output_dir)
 
+        elif getattr(self.model, "_device_mesh", None) is not None and not _is_peft_model(self.model):
+            # Sharded at load time (`DistributedConfig`): gathering the weights inside `save_pretrained`
+            # is collective, so every rank saves; only the main process writes, the others leave at the
+            # closing barrier. (PEFT models fall through to the adapter-only save below.)
+            self._save(output_dir)
+
         elif self.args.should_save:
             self._save(output_dir)
 
@@ -3934,10 +4008,10 @@ def save_model(self, output_dir: str | None = None, _internal_call: bool = False
 
     def _save(self, output_dir: str | None = None, state_dict: dict | None = None) -> None:
         """Save model weights, configuration, and processing class to `output_dir`."""
-        # If we are executing this function, we are the process zero, so we don't check for that.
         output_dir = output_dir if output_dir is not None else self.args.output_dir
         os.makedirs(output_dir, exist_ok=True)
-        logger.info(f"Saving model checkpoint to {output_dir}")
+        if self.args.should_save:
+            logger.info(f"Saving model checkpoint to {output_dir}")
 
         supported_classes = (PreTrainedModel,) if not is_peft_available() else (PreTrainedModel, PeftModel)
         # Save a trained model and configuration using `save_pretrained()`.
@@ -3958,6 +4032,10 @@ def _save(self, output_dir: str | None = None, state_dict: dict | None = None) -
         else:
             self.model.save_pretrained(output_dir, state_dict=state_dict)
 
+        # A non-writer rank of a model sharded at load time is only here for the collectives above.
+        if not self.args.should_save:
+            return
+
         if self.processing_class is not None:
             self.processing_class.save_pretrained(output_dir)
         elif (
diff --git a/src/transformers/trainer_optimizer.py b/src/transformers/trainer_optimizer.py
index de63c614f156..13ab9edc18da 100644
--- a/src/transformers/trainer_optimizer.py
+++ b/src/transformers/trainer_optimizer.py
@@ -208,6 +208,19 @@ def _get_adamw_torch(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
     return AdamW, ctx.optimizer_kwargs
 
 
+def has_mixed_dtensor(tensors) -> bool:
+    """
+    Whether `tensors` do not all share one device mesh, so whole-set ops (fused/foreach kernels, `clip_grad_norm_`)
+    cannot span them and the Trainer groups them by mesh. Expert parallelism leaves the experts sharded and everything
+    else as plain tensors; under a 2-D (fsdp, tp) mesh everything is a `DTensor`, but the experts live on the full mesh
+    and the rest on the `fsdp` sub-mesh.
+    """
+    from torch.distributed.tensor import DTensor
+
+    meshes = {t.device_mesh if isinstance(t, DTensor) else None for t in tensors}
+    return len(meshes) > 1
+
+
 def _get_adamw_torch_xla(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
     """Get Torch XLA syncfree AdamW optimizer."""
     try:
diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index 51dd0c7b20e5..e780b74f56c9 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -23,6 +23,7 @@
 import time
 import traceback
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from contextlib import contextmanager
 
 from parameterized import parameterized
@@ -494,6 +495,98 @@ def _test_fsdp2_plan_vs_ddp_impl(rank, config_class, config_dict, tie_word_embed
         logger.debug("DDP and FSDP2 comparison checks passed.")
 
 
+def _grad_norm_across_meshes(model):
+    """Total gradient norm of parameters living on different device meshes (what the Trainer does)."""
+    from torch.distributed.tensor import DTensor
+    from torch.nn.utils import get_total_norm
+
+    grads_by_mesh = defaultdict(list)
+    for param in model.parameters():
+        if param.grad is not None:
+            grads_by_mesh[param.grad.device_mesh if isinstance(param.grad, DTensor) else None].append(param.grad)
+    norms = [get_total_norm(grads) for grads in grads_by_mesh.values()]
+    norms = [n.full_tensor() if isinstance(n, DTensor) else n for n in norms]
+    return torch.linalg.vector_norm(torch.stack(norms))
+
+
+def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, dtype=None):
+    """
+    DDP vs a 2-D (fsdp, tp) mesh with expert parallelism on `tp`. DDP sees the whole batch on every rank; each `fsdp`
+    rank of the 2-D run sees its own slice of it, so FSDP2's reduction over `fsdp` is exercised. Losses, gradient norms
+    and final weights have to match step by step.
+    """
+    init_test_logger()
+
+    if dtype is None:
+        dtype = torch.float32
+
+    device = _get_rank_device(rank)
+    config = config_class.from_dict(config_dict)
+    world_size = dist.get_world_size()
+    dp = world_size // 2
+    generator = torch.Generator(device=device)
+    generator.manual_seed(SEED)
+    input_ids = torch.randint(0, config.vocab_size, (dp * BATCH_SIZE, SEQ_LEN), device=device, generator=generator)
+    batches = [(input_ids, input_ids.clone())] * NUM_STEPS
+
+    with _deterministic_init_model_dir(rank, config, dtype) as init_model_dir:
+        ddp_losses, ddp_grad_norms, ddp_state_dict = train_ddp(rank, batches, LR, device, dtype, init_model_dir)
+
+        _set_determinism(SEED)
+        model = AutoModelForCausalLM.from_pretrained(
+            init_model_dir,
+            torch_dtype=dtype,
+            distributed_config=DistributedConfig(tp_size=2, fsdp_size=dp, enable_expert_parallel=True),
+        )
+        assert model.tp_size == 2 and model.fsdp_size == dp
+        assert model._device_mesh.mesh_dim_names == ("fsdp", "tp")
+        model.train()
+        optimizer = torch.optim.Adam(model.parameters(), lr=LR, foreach=False)
+        dp_rank = model._device_mesh["fsdp"].get_local_rank()
+        dp_group = model._device_mesh["fsdp"].get_group()
+        losses, grad_norms = [], []
+        for ids, labels in batches:
+            rows = slice(dp_rank * BATCH_SIZE, (dp_rank + 1) * BATCH_SIZE)
+            optimizer.zero_grad()
+            loss = model(input_ids=ids[rows], labels=labels[rows], use_cache=False).loss
+            loss.backward()
+            grad_norms.append(_grad_norm_across_meshes(model).item())
+            optimizer.step()
+            loss = loss.detach()
+            dist.all_reduce(loss, group=dp_group)
+            losses.append(loss.item() / dp)
+        state_dict = gather_full_state_dict(model)
+
+    for step in range(len(ddp_losses)):
+        torch.testing.assert_close(
+            torch.tensor(ddp_losses[step]),
+            torch.tensor(losses[step]),
+            rtol=DDP_FSDP_RTOL,
+            atol=DDP_FSDP_ATOL,
+            msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, FSDP2+EP={losses[step]}",
+        )
+        torch.testing.assert_close(
+            torch.tensor(ddp_grad_norms[step]),
+            torch.tensor(grad_norms[step]),
+            rtol=DDP_FSDP_RTOL,
+            atol=DDP_FSDP_ATOL,
+            msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, FSDP2+EP={grad_norms[step]}",
+        )
+
+    for key in ddp_state_dict:
+        assert key in state_dict, f"Key {key} missing from FSDP2+EP state dict"
+        torch.testing.assert_close(
+            ddp_state_dict[key],
+            state_dict[key],
+            rtol=DDP_FSDP_RTOL,
+            atol=DDP_FSDP_ATOL,
+            msg=f"Weight mismatch for {key}: DDP vs FSDP2+EP",
+        )
+
+    if rank == 0:
+        logger.debug("DDP and FSDP2+EP (2-D mesh) comparison checks passed.")
+
+
 # =============================================================================
 # Mixin class
 # =============================================================================
@@ -509,10 +602,10 @@ def model_tester(self):
         """The model tester instance (e.g., CausalLMModelTester)."""
         ...
 
-    def _skip_if_insufficient_devices(self):
+    def _skip_if_insufficient_devices(self, world_size):
         available_workers = _get_available_fsdp_workers()
-        if available_workers < self.fsdp_nproc_per_node:
-            self.skipTest(f"Need at least {self.fsdp_nproc_per_node} FSDP workers, have {available_workers}")
+        if available_workers < world_size:
+            self.skipTest(f"Need at least {world_size} FSDP workers, have {available_workers}")
 
     def _skip_if_mps(self):
         if torch._C._get_accelerator().type == "mps":
@@ -558,9 +651,10 @@ def _get_tiny_config(self):
             config.vocab_size_per_layer_input = config.vocab_size
         return type(config), config.to_diff_dict()
 
-    def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_kwargs):
+    def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, world_size=None, **test_kwargs):
+        world_size = world_size or self.fsdp_nproc_per_node
         self._skip_if_mps()
-        self._skip_if_insufficient_devices()
+        self._skip_if_insufficient_devices(world_size)
         self._skip_if_fsdp_distributed_not_enabled()
 
         config_class, config_dict = self._get_tiny_config()
@@ -575,8 +669,8 @@ def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_k
         try:
             mp.spawn(
                 _fsdp_global_wrapper,
-                args=(test_name, test_impl, func_args, test_kwargs, self.fsdp_nproc_per_node, port, results_file),
-                nprocs=self.fsdp_nproc_per_node,
+                args=(test_name, test_impl, func_args, test_kwargs, world_size, port, results_file),
+                nprocs=world_size,
             )
 
             with open(results_file) as f:
@@ -636,3 +730,16 @@ def test_fsdp2_plan_vs_ddp(self, label):
             _test_fsdp2_plan_vs_ddp_impl,
             label == "tied",
         )
+
+    @is_fsdp_test
+    def test_fsdp2_expert_parallel_2d_vs_ddp(self):
+        """
+        Training on a 2-D (fsdp, tp) mesh with expert parallelism, each fsdp rank on its own slice of the batch,
+        traces DDP on the whole batch step by step.
+        """
+        config = self.model_tester.get_config()
+        if getattr(config, "base_model_ep_plan", None) is None:
+            self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)")
+        self._run_fsdp2_distributed_test(
+            "fsdp2_expert_parallel_2d_vs_ddp", _test_fsdp2_expert_parallel_2d_vs_ddp_impl, world_size=4
+        )