[Umbrella] Enable FSDP2 + expert parallelism via a 2-D (fsdp, tp) device mesh

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

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.

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:

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

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,22 +277,24 @@ 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. "
                     "Call save_pretrained from every rank after init_process_group."
                 )
             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:

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

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

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