diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index eb1bafd7355f..2de4ebfcdf07 100644 --- a/src/transformers/distributed/mixin.py +++ b/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 diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 7a5005146858..5b2781b51916 100755 --- a/src/transformers/trainer.py +++ b/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"): diff --git a/src/transformers/trainer_optimizer.py b/src/transformers/trainer_optimizer.py index de63c614f156..bff836a4ca45 100644 --- a/src/transformers/trainer_optimizer.py +++ b/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: