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

docs/source/en/expert_parallelism.md

@@ -50,4 +50,35 @@ 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 limits how large a model you can train. Add [FSDP2](./fsdp) on a second mesh dimension with `fsdp_size`, and keep using `tp_size` for the expert parallel width (`tp_size` is the EP size).
+
+```py
+from transformers import AutoModelForCausalLM
+from transformers.distributed import DistributedConfig
+
+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 2D `(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.
+
+Load the model as usual, then train with [`Trainer`]. It takes the gradient norm across both meshes and gives each mesh its own optimizer param group. [`~Trainer.save_model`] gathers sharded weights into a regular checkpoint. This requires `accelerate>=1.12` so the `Trainer` can mirror `tp_size` and `fsdp_size` into [`~Accelerate.ParallelismConfig`].
+
+The table below compares EP-only training with 2D EP+FSDP2 on 8xH100 GPUs. The workload is full fine-tuning of Qwen3-30B-A3B in bf16 at sequence length 2048. More FSDP shards cut peak memory, and tokens/s drop some because FSDP2 all-gathers and reduce-scatters the experts across `fsdp`.
+
+| 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/distributed/configuration_utils.py

@@ -73,10 +73,10 @@ def __post_init__(self):
         elif self.tp_size is None:
             self.tp_size = 1
 
-        if self.fsdp_size > 1 and (self.tp_size > 1 or self.pp_size > 1):
+        if self.fsdp_size > 1 and self.pp_size > 1:
             raise ValueError(
-                "Combining FSDP with tensor or pipeline parallelism is not supported yet. "
-                "Use FSDP alone, or combine TP and PP."
+                "Combining FSDP with pipeline parallelism is not supported yet. "
+                "Use DistributedConfig(tp_size=N, fsdp_size=M), or combine TP and PP."
             )
 
     @classmethod

src/transformers/distributed/mixin.py

@@ -52,6 +52,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
 
@@ -172,18 +173,21 @@ def maybe_distribute_model(
             model.config.distributed_config = distributed_config
             model._device_mesh = device_mesh
             model._tp_size = distributed_config.tp_size
+            model._fsdp_size = distributed_config.fsdp_size
 
             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)
 
+            # 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:
-                # TODO(3outeille): turn into a if statement when TP + FSDP is supported
+
+            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)
         return model
@@ -246,22 +250,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/utils.py

@@ -173,6 +173,9 @@ def initialize_distributed_mesh(
         tuple(mesh_shape),
         mesh_dim_names=tuple(mesh_dim_names),
     )
+    # A flattened sub-mesh, so an all-reduce over every rank is one collective instead of one per dimension.
+    if len(mesh_dim_names) > 1:
+        device_mesh._flatten("_".join(mesh_dim_names))
     return device_map, device_mesh
 
 

src/transformers/modeling_utils.py

@@ -4612,6 +4612,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

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:
@@ -612,6 +617,15 @@ def __init__(
         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)
 
@@ -748,17 +762,20 @@ 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.
+            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
@@ -1254,6 +1271,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
@@ -2616,48 +2645,33 @@ def _track_num_input_tokens(self, inputs):
         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 only some parameters are sharded, as under expert parallelism.
-
-        Expert parallelism shards the expert weights and leaves everything else replicated, so
-        `model.parameters()` holds a mix of `DTensor` and plain tensors and `_foreach_norm` cannot span
-        both. Take the two groups separately: the sharded gradients contribute their local norms summed
-        across the mesh, the replicated ones are identical on every rank and are counted once.
+        """
+        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
 
-        sharded, replicated = [], []
+        params_by_mesh = defaultdict(list)
         for param in model.parameters():
-            if param.grad is None:
-                continue
-            (sharded if isinstance(param.grad, DTensor) else replicated).append(param.grad)
-
-        device = (sharded or replicated)[0].device
-        replicated_sq = torch.zeros((), device=device, dtype=torch.float32)
-        if replicated:
-            replicated_sq = torch.linalg.vector_norm(torch.stack([g.norm(2) for g in replicated])) ** 2
+            if param.grad is not None:
+                params_by_mesh[param.grad.device_mesh if isinstance(param.grad, DTensor) else None].append(param)
 
-        sharded_sq = torch.zeros((), device=device, dtype=torch.float32)
-        if sharded:
-            local = [g.to_local() for g in sharded]
-            sharded_sq = torch.linalg.vector_norm(torch.stack([g.norm(2) for g in local])) ** 2
-            # Sum the per-shard contributions over the mesh the experts are sharded on.
-            torch.distributed.all_reduce(sharded_sq, group=sharded[0].device_mesh.get_group())
+        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))
 
-        total_norm = (replicated_sq + sharded_sq).sqrt()
         if max_norm != float("inf"):
-            clip = (max_norm / (total_norm + 1e-6)).clamp(max=1.0)
-            for g in replicated:
-                g.mul_(clip)
-            for g in sharded:
-                g.to_local().mul_(clip)
+            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:
-            from .trainer_optimizer import has_mixed_dtensor
-
             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
 
@@ -3978,6 +3992,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)
 
@@ -3987,10 +4007,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()`.
@@ -4011,6 +4031,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 (

src/transformers/trainer_optimizer.py

@@ -205,24 +205,20 @@ def _get_adamw_torch(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
     ctx.optimizer_kwargs.update(ctx.adam_kwargs)
     if ctx.args.optim == OptimizerNames.ADAMW_TORCH_FUSED:
         ctx.optimizer_kwargs.update({"fused": True})
-    # HACK(ep-2d): expert parallelism leaves only the expert weights as DTensor, so the fused and
-    # foreach kernels cannot span the parameter set. Step per parameter instead.
-    if ctx.model is not None and has_mixed_dtensor(p for p in ctx.model.parameters() if p.requires_grad):
-        ctx.optimizer_kwargs.update({"fused": False, "foreach": False})
     return AdamW, ctx.optimizer_kwargs
 
 
 def has_mixed_dtensor(tensors) -> bool:
-    """Whether `tensors` mixes `DTensor` and plain tensors, as expert parallelism leaves a model:
-    experts sharded, everything else replicated."""
+    """
+    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
 
-    tensors = list(tensors)
-    return (
-        bool(tensors)
-        and any(isinstance(t, DTensor) for t in tensors)
-        and not all(isinstance(t, DTensor) for t in tensors)
-    )
+    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]]: