From 03d6d499eea2092bf71bc20ecf379f3a101cd9b9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Sat, 22 Aug 2026 02:36:30 +0000 Subject: [PATCH 1/4] Fix expert parallelism through Trainer Loading a model with DistributedConfig(tp_size=N, enable_expert_parallel=True) and handing it to Trainer fails on main in three places: 1. maybe_distribute_model never assigns model._tp_size, so the Trainer builds no ParallelismConfig and accelerate wraps the DTensor model in DDP: ValueError: Your model contains DTensor parameters, which is incompatible with DDP. 2. _get_grad_norm calls clip_grad_norm_ over the full parameter set, and _foreach_norm cannot span a mix of DTensor (experts) and plain parameters. 3. The fused/foreach AdamW kernels cannot span that mix either. Set _tp_size where the mesh is recorded, compute the gradient norm (and clip) per-gradient with replication-aware discounting when parameters live on different meshes, and fall back to per-parameter AdamW stepping for mixed parameter sets. With this and #48205, expert-parallel full fine-tuning through Trainer runs end-to-end and tracks a single-GPU control step by step (OLMoE-1B-7B, tp=4: 12.13 -> 12.45 -> 11.52 -> ... -> 11.04 vs 12.13 -> 12.45 -> 11.51 -> ... -> 11.09). --- src/transformers/distributed/mixin.py | 4 ++ src/transformers/trainer.py | 54 ++++++++++++++++++++++++++- src/transformers/trainer_optimizer.py | 13 +++++++ 3 files changed, 70 insertions(+), 1 deletion(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index eb1bafd7355f..7676b562b065 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -190,6 +190,10 @@ def maybe_distribute_model( if device_mesh is not None: model.config.distributed_config = distributed_config model._device_mesh = device_mesh + # `Trainer` reads this to build accelerate's `ParallelismConfig`. Without it accelerate + # sees `world_size` unaccounted-for ranks and wraps the already-sharded model in DDP, + # which rejects DTensor parameters. + 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..98e7329be8e7 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2613,17 +2613,69 @@ 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: + from torch.distributed.tensor import DTensor + + grads = [p.grad for p in model.parameters() if p.grad is not None] + return ( + bool(grads) + and any(isinstance(g, DTensor) for g in grads) + and not all(isinstance(g, DTensor) for g in 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..0c2374b15437 100644 --- a/src/transformers/trainer_optimizer.py +++ b/src/transformers/trainer_optimizer.py @@ -205,9 +205,22 @@ 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 _has_mixed_dtensor_params(ctx.model): + ctx.optimizer_kwargs.update({"fused": False, "foreach": False}) return AdamW, ctx.optimizer_kwargs +def _has_mixed_dtensor_params(model) -> bool: + from torch.distributed.tensor import DTensor + + if model is None: + return False + ps = [p for p in model.parameters() if p.requires_grad] + return bool(ps) and any(isinstance(p, DTensor) for p in ps) and not all(isinstance(p, DTensor) for p in ps) + + def _get_adamw_torch_xla(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]: """Get Torch XLA syncfree AdamW optimizer.""" try: From a2b284e389ad50ad96735861161afe83d1e13380 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Sat, 22 Aug 2026 02:41:35 +0000 Subject: [PATCH 2/4] Remove comment --- src/transformers/distributed/mixin.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 7676b562b065..2de4ebfcdf07 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -190,9 +190,6 @@ def maybe_distribute_model( if device_mesh is not None: model.config.distributed_config = distributed_config model._device_mesh = device_mesh - # `Trainer` reads this to build accelerate's `ParallelismConfig`. Without it accelerate - # sees `world_size` unaccounted-for ranks and wraps the already-sharded model in DDP, - # which rejects DTensor parameters. model._tp_size = distributed_config.tp_size if distributed_config.tp_size > 1: From a19ebb87e2dbb6cc28ac14c443830128c4f6e5fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 3 Sep 2026 03:03:23 +0000 Subject: [PATCH 3/4] Share the mixed-DTensor predicate and cache the clip-time result on the trainer --- src/transformers/trainer.py | 17 ++++++++++------- src/transformers/trainer_optimizer.py | 16 ++++++++++------ 2 files changed, 20 insertions(+), 13 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 98e7329be8e7..c2c453673bf8 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) @@ -2651,14 +2653,15 @@ def _mixed_mesh_grad_norm(self, model, max_norm): return total_norm def _has_mixed_mesh_grads(self, model) -> bool: - from torch.distributed.tensor import DTensor + # 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 - grads = [p.grad for p in model.parameters() if p.grad is not None] - return ( - bool(grads) - and any(isinstance(g, DTensor) for g in grads) - and not all(isinstance(g, DTensor) for g in grads) - ) + 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.""" diff --git a/src/transformers/trainer_optimizer.py b/src/transformers/trainer_optimizer.py index 0c2374b15437..bff836a4ca45 100644 --- a/src/transformers/trainer_optimizer.py +++ b/src/transformers/trainer_optimizer.py @@ -207,18 +207,22 @@ def _get_adamw_torch(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]: 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 _has_mixed_dtensor_params(ctx.model): + 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_params(model) -> bool: +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 - if model is None: - return False - ps = [p for p in model.parameters() if p.requires_grad] - return bool(ps) and any(isinstance(p, DTensor) for p in ps) and not all(isinstance(p, DTensor) for p in ps) + 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]]: From f464638c9c19ee0d75ed828cae009d82e3deb0a1 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 3 Sep 2026 15:06:09 +0000 Subject: [PATCH 4/4] Apply repo consistency fixes --- src/transformers/trainer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index c2c453673bf8..5b2781b51916 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2658,9 +2658,7 @@ def _has_mixed_mesh_grads(self, model) -> bool: 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 - ) + 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):