Fix expert parallelism through Trainer

src/transformers/distributed/mixin.py

@@ -190,6 +190,7 @@ def maybe_distribute_model(
         if device_mesh is not None:
             model.config.distributed_config = distributed_config
             model._device_mesh = device_mesh
+            model._tp_size = distributed_config.tp_size
 
             if distributed_config.tp_size > 1:
                 tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh

src/transformers/trainer.py

@@ -610,6 +610,8 @@ 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
 
         self.control = self.callback_handler.on_init_end(self.args, self.state, self.control)
 
@@ -2613,17 +2615,68 @@ 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 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.
+        """
+        from torch.distributed.tensor import DTensor
+
+        sharded, replicated = [], []
+        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
+
+        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())
+
+        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)
+        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
+
     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"):

src/transformers/trainer_optimizer.py

@@ -205,9 +205,26 @@ 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."""
+    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)
+    )
+
+
 def _get_adamw_torch_xla(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
     """Get Torch XLA syncfree AdamW optimizer."""
     try: