From 8af1a002aaabce1a058dc53a2f00abfb249ef1e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 21 Aug 2026 23:10:29 +0000 Subject: [PATCH 01/52] Fix NaN gradients in expert-parallel training: mask uninitialized grouped_mm rows Under EP, sentinel token-expert slots sit beyond offsets[-1] and torch._grouped_mm leaves those output rows (fwd output and bwd d_input) uninitialized. The forward relied on a single post-mask plus a single pre-mask, letting NaN/Inf from uninitialized memory transit the activation and down-projection backward. The gate product's backward (act_fn(gate) * up) turns 0 x Inf into NaN (torch.autograd anomaly mode names this exact Mul), and it escapes into finite gradients: full fine-tuning of any EP-sharded MoE produced nan grad_norm on the second step (the first step survives only because freshly-allocated CUDA memory happens to be zeroed) and the loss collapsed to 0. Mask the sentinel-tail rows after each grouped GEMM instead. Full fine-tuning of OLMoE-1B-7B under ep=4 now matches the single-GPU loss trajectory. --- src/transformers/integrations/moe.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index b9b288f4a696..bd0adeaaaf89 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -441,6 +441,11 @@ def grouped_mm_experts_forward( selected_hidden_states_g, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed ) # (S, 2 * intermediate_dim) or (S, intermediate_dim) depending on whether we have gating + # Mask the sentinel-tail rows the kernel left uninitialized, in forward AND backward: relying on + # the single post-mask lets NaN/Inf from the uninitialized rows transit through the activation and + # down-projection backward, where it can escape into finite gradients. + proj_out = proj_out.masked_fill(sentinel_mask, 0.0) + # Apply gating or activation if self.has_gate: # for gated experts we apply the custom/default gating mechanism @@ -458,6 +463,9 @@ def grouped_mm_experts_forward( proj_out, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed ) # (S, hidden_dim) + # Same as above: zero the uninitialized sentinel-tail rows before anything consumes them. + proj_out = proj_out.masked_fill(sentinel_mask, 0.0) + # Apply routing weights weighted_out = proj_out * sample_weights_g.unsqueeze(-1) # (S, hidden_dim) From 1f0b7bf661b05ae788b51bec43cc95a4cdbb853f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 21 Aug 2026 23:56:29 +0000 Subject: [PATCH 02/52] Fix wrong gradients for all non-expert parameters in expert-parallel training Under EP the router hook zeroes the routing scores of non-local experts, so in backward each rank's score gradient covers only the slots of its local experts, and nothing sums the per-rank partial gradients: the gate weights and, through the gate's input, every parameter upstream of each MoE block receive gradients missing the contributions that flow through remote experts. The existing _AllReduceBackward on the experts' hidden input covers the dispatch branch, and the top_k_weights branch is explicitly skipped when is_expert_parallel -- but under EP it is exactly as partial as under TP-MoE. Measured against a single-GPU reference (OLMoE-1B-7B, one batch, fp32 so rounding noise vanishes): before the fix, 3/179 parameters agree (relative max-abs errors 0.3-2.5 on attention, norms, embeddings and router gates, 10-100x above the run-to-run noise floor; only the last layer's experts and the final norm -- the parameters backward reaches before crossing an expert block -- are correct). After the fix: 179/179 agree, max relative error 2.7e-5. Fix: allreduce-sum the score gradient in the EP router hook, before the non-local mask (each slot has exactly one owning rank, so the sum is exact). --- src/transformers/distributed/tensor_parallel.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 3dc23a07a304..2bb1e7038260 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -709,6 +709,12 @@ 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 + # The gradient reaching each rank's scores covers only the slots of its local experts (the + # non-local slots are zeroed by the mask below in backward). Sum the per-rank partial + # gradients BEFORE the mask so the gate weights and the hidden states upstream of the router + # receive the full gradient — each slot has exactly one owning rank, so the sum is exact. + 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) From 4636d4253ca784d3e279e72d682b71ebee6e5622 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Sat, 22 Aug 2026 01:08:56 +0000 Subject: [PATCH 03/52] Trim comments --- src/transformers/distributed/tensor_parallel.py | 6 ++---- src/transformers/integrations/moe.py | 6 ++---- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 2bb1e7038260..bc7364502c6b 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -709,10 +709,8 @@ 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 - # The gradient reaching each rank's scores covers only the slots of its local experts (the - # non-local slots are zeroed by the mask below in backward). Sum the per-rank partial - # gradients BEFORE the mask so the gate weights and the hidden states upstream of the router - # receive the full gradient — each slot has exactly one owning rank, so the sum is exact. + # 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). 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 diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index bd0adeaaaf89..49312469d43a 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -441,9 +441,7 @@ def grouped_mm_experts_forward( selected_hidden_states_g, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed ) # (S, 2 * intermediate_dim) or (S, intermediate_dim) depending on whether we have gating - # Mask the sentinel-tail rows the kernel left uninitialized, in forward AND backward: relying on - # the single post-mask lets NaN/Inf from the uninitialized rows transit through the activation and - # down-projection backward, where it can escape into finite gradients. + # Zero the sentinel-tail rows the kernel left uninitialized (fwd output and bwd `d_input`). proj_out = proj_out.masked_fill(sentinel_mask, 0.0) # Apply gating or activation @@ -463,7 +461,7 @@ def grouped_mm_experts_forward( proj_out, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed ) # (S, hidden_dim) - # Same as above: zero the uninitialized sentinel-tail rows before anything consumes them. + # Same: zero the uninitialized sentinel-tail rows. proj_out = proj_out.masked_fill(sentinel_mask, 0.0) # Apply routing weights From 89fa55ca91ec52c19a17c74d75cdae7b80c587ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 3 Sep 2026 02:17:41 +0000 Subject: [PATCH 04/52] Gate the router-score backward all-reduce on grad mode; drop the post-mask superseded by the per-mm masks --- src/transformers/distributed/tensor_parallel.py | 5 +++-- src/transformers/integrations/moe.py | 3 --- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index bc7364502c6b..6a9455637a33 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -711,8 +711,9 @@ def transform_output_post_forward(self, module, output, mesh): 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). - process_group = mesh.get_group() if mesh.ndim == 1 else mesh.get_group("tp") - router_scores = _AllReduceBackward.apply(router_scores, process_group) + 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) diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 49312469d43a..fb3b736c0301 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -467,9 +467,6 @@ def grouped_mm_experts_forward( # Apply routing weights weighted_out = proj_out * sample_weights_g.unsqueeze(-1) # (S, hidden_dim) - # Post-mask (fwd path). - weighted_out.masked_fill_(sentinel_mask, 0.0) - # Restore original order inv_perm = torch.empty_like(perm) inv_perm[perm] = torch.arange(perm.size(0), device=device) From 96769d9972449eb4874553aa7cb3cb968e02f1b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Thu, 3 Sep 2026 14:22:19 -0400 Subject: [PATCH 05/52] Fix expert parallelism through Trainer (#48208) --- src/transformers/distributed/mixin.py | 1 + src/transformers/trainer.py | 55 ++++++++++++++++++++++++++- src/transformers/trainer_optimizer.py | 17 +++++++++ 3 files changed, 72 insertions(+), 1 deletion(-) 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: From c74c1228fb9e0ae515cd433dab2ca24e8aba14e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 3 Sep 2026 19:45:04 +0000 Subject: [PATCH 06/52] Enable FSDP2 + expert parallelism via a 2-D (fsdp, tp) device mesh DistributedConfig(tp_size=E, fsdp_size=D, enable_expert_parallel=True) builds a 2-D mesh: experts are sharded across tp, everything else is fully sharded across fsdp. The Trainer mirrors both dimensions into accelerate's ParallelismConfig, averages the expert gradients over fsdp (FSDP2 only reduces what it shards), computes the gradient norm across parameters on different meshes, and gathers the DTensor state dict on save. --- src/transformers/distributed/mixin.py | 18 +++-- src/transformers/distributed/utils.py | 6 ++ src/transformers/modeling_utils.py | 8 ++ src/transformers/trainer.py | 110 ++++++++++++++++++++------ src/transformers/trainer_optimizer.py | 18 ++--- tests/test_fsdp_mixin.py | 80 +++++++++++++++++-- 6 files changed, 192 insertions(+), 48 deletions(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 2de4ebfcdf07..ff7834cf62ef 100644 --- a/src/transformers/distributed/mixin.py +++ b/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,7 +164,12 @@ def prepare_distribute_model( f"is not equal to world_size ({world_size})" ) - if distributed_config.tp_size > 1: + if distributed_config.fsdp_size > 1: + # Builds a 2-D (fsdp, tp) mesh when tensor/expert parallelism is also requested. + device_map, device_mesh = initialize_fully_sharded_data_parallelism(distributed_config) + if distributed_config.tp_size > 1 and distributed_config.tp_plan is None: + distributed_config.tp_plan = "auto" + elif distributed_config.tp_size > 1: if distributed_config.tp_plan is None: distributed_config.tp_plan = "auto" device_map, device_mesh = initialize_tensor_parallelism( @@ -172,8 +178,6 @@ def prepare_distribute_model( 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,19 +194,23 @@ 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: experts/attention are sharded across `tp`, the rest 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 diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index 49f0f1c41814..50cf27bda817 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -193,11 +193,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 or 1 + # `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)) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index a7da91270ff4..334ec8955476 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4611,6 +4611,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 diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 5b2781b51916..49c370beeaeb 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -473,6 +473,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: @@ -748,7 +751,16 @@ 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: + model_fsdp_size = getattr(self.model, "fsdp_size", None) or 1 + if model_fsdp_size > 1: + # 2-D mesh: accelerate has to know about both dimensions or it will try to wrap the + # already-sharded model in DDP. + from accelerate import ParallelismConfig + + args["parallelism_config"] = ParallelismConfig( + dp_shard_size=model_fsdp_size, tp_size=getattr(self.model, "tp_size", None) or 1 + ) + elif 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: @@ -1859,6 +1871,7 @@ def _run_epoch( self._track_num_input_tokens(inputs) if do_sync_step: + self._sync_replicated_grads(model) grad_norm = None if self.args.max_grad_norm > 0: grad_norm = self._clip_grad_norm(model) @@ -2615,41 +2628,64 @@ 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. - """ + def _sync_replicated_grads(self, model): + """Average the gradients of parameters replicated along the data-parallel dimension (e.g. the + EP-sharded experts under a 2-D mesh) -- FSDP2 only reduces what it shards.""" from torch.distributed.tensor import DTensor - sharded, replicated = [], [] + mesh = getattr(model, "_device_mesh", None) or getattr(getattr(model, "module", None), "_device_mesh", None) + if mesh is None or mesh.ndim < 2 or "fsdp" not in (mesh.mesh_dim_names or ()): + return + dp_group = mesh["fsdp"].get_group() + dp_size = mesh["fsdp"].size() + if dp_size < 2: + return + for param in model.parameters(): if param.grad is None: continue - (sharded if isinstance(param.grad, DTensor) else replicated).append(param.grad) + grad = param.grad + if isinstance(grad, DTensor): + # Already reduced over `fsdp` if that dimension is part of the gradient's own mesh. + if "fsdp" in (grad.device_mesh.mesh_dim_names or ()): + continue + grad = grad.to_local() + torch.distributed.all_reduce(grad, group=dp_group) + grad.div_(dp_size) - 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 + def _mixed_mesh_grad_norm(self, model, max_norm): + """Gradient norm (and clip) when parameters live on different device meshes: sum each gradient's + local squared norm, discounting the mesh dimensions it is replicated over, reduce once.""" + from torch.distributed.tensor import DTensor + + grads = [p.grad for p in model.parameters() if p.grad is not None] + if not grads: + return torch.zeros((), device=self.args.device) + + world = torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1 + total_sq = torch.zeros((), device=grads[0].device, dtype=torch.float32) + for g in grads: + if isinstance(g, DTensor): + local = g.to_local() + # Ranks along a replicated mesh dimension all hold the same values. + replicas = 1 + for dim, placement in enumerate(g.placements): + if placement.is_replicate(): + replicas *= g.device_mesh.size(dim) + # Dimensions of the world this tensor's mesh does not span also hold copies. + replicas *= world // g.device_mesh.size() + else: + local, replicas = g, world # a plain tensor is the same on every rank + total_sq += local.detach().float().pow(2).sum() / replicas - 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()) + if world > 1: + torch.distributed.all_reduce(total_sq) + total_norm = total_sq.sqrt() - 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 g in grads: + (g.to_local() if isinstance(g, DTensor) else g).mul_(clip) return total_norm def _has_mixed_mesh_grads(self, model) -> bool: @@ -3978,6 +4014,28 @@ 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.accelerator.unwrap_model(self.model, keep_torch_compile=False), "_device_mesh", None + ) is not None and not _is_peft_model(self.model): + # The model was sharded at load time (`DistributedConfig`). Gathering its DTensor state + # dict is collective, so every rank must participate; only the main process writes. + # (PEFT models fall through to the adapter-only save below.) + from .distributed.tensor_parallel import gather_state_dict_for_save + + unwrapped = self.accelerator.unwrap_model(self.model, keep_torch_compile=False) + gather_start = time.time() + state_dict = gather_state_dict_for_save(unwrapped.state_dict(), None, None, 0) + gather_s = time.time() - gather_start + if self.args.should_save: + write_start = time.time() + # `save_pretrained` ends with `barrier_after_gathered_checkpoint_save`. + self._save(output_dir, state_dict=state_dict) + logger.info(f"Sharded-model save: gather {gather_s:.1f}s, write {time.time() - write_start:.1f}s") + else: + # Match the writer's end-of-save barrier so no rank runs ahead into the next + # step's collectives while the main process is still writing. + unwrapped.barrier_after_gathered_checkpoint_save(unwrapped.config.distributed_config) + elif self.args.should_save: self._save(output_dir) diff --git a/src/transformers/trainer_optimizer.py b/src/transformers/trainer_optimizer.py index bff836a4ca45..bc78622cdc29 100644 --- a/src/transformers/trainer_optimizer.py +++ b/src/transformers/trainer_optimizer.py @@ -205,24 +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. + # Expert parallelism leaves the parameters on different meshes, 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. 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 `tp` mesh and the rest on the `fsdp` 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]]: diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index 51dd0c7b20e5..ea658341ce40 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -494,6 +494,60 @@ def _test_fsdp2_plan_vs_ddp_impl(rank, config_class, config_dict, tie_word_embed logger.debug("DDP and FSDP2 comparison checks passed.") +def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, dtype=None): + """DDP vs a 2-D (fsdp, tp) mesh with expert parallelism on `tp`: same batch on every rank, so the + sharded run has to trace the DDP run step by step and end on the same weights.""" + init_test_logger() + + if dtype is None: + dtype = torch.float32 + + device = _get_rank_device(rank) + config = config_class.from_dict(config_dict) + batches = _build_repeated_training_batches(config, device, NUM_STEPS) + + with _deterministic_init_model_dir(rank, config, dtype) as init_model_dir: + ddp_losses, _, ddp_state_dict = train_ddp(rank, batches, LR, device, dtype, init_model_dir) + + _set_determinism(SEED) + world_size = dist.get_world_size() + model = AutoModelForCausalLM.from_pretrained( + init_model_dir, + torch_dtype=dtype, + distributed_config=DistributedConfig(tp_size=2, fsdp_size=world_size // 2, enable_expert_parallel=True), + ) + assert model.tp_size == 2 and model.fsdp_size == world_size // 2 + assert model._device_mesh.mesh_dim_names == ("fsdp", "tp") + model.train() + optimizer = torch.optim.Adam(model.parameters(), lr=LR, foreach=False) + # `clip_grad_norm_` cannot span parameters living on different meshes; the Trainer has its own + # norm for that case, so only the losses and the final weights are compared here. + losses, _ = _run_training_steps(model, optimizer, batches, track_grad_norms=False) + state_dict = gather_full_state_dict(model) + + for step in range(len(ddp_losses)): + torch.testing.assert_close( + torch.tensor(ddp_losses[step]), + torch.tensor(losses[step]), + rtol=DDP_FSDP_RTOL, + atol=DDP_FSDP_ATOL, + msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, FSDP2+EP={losses[step]}", + ) + + for key in ddp_state_dict: + assert key in state_dict, f"Key {key} missing from FSDP2+EP state dict" + torch.testing.assert_close( + ddp_state_dict[key], + state_dict[key], + rtol=DDP_FSDP_RTOL, + atol=DDP_FSDP_ATOL, + msg=f"Weight mismatch for {key}: DDP vs FSDP2+EP", + ) + + if rank == 0: + logger.debug("DDP and FSDP2+EP (2-D mesh) comparison checks passed.") + + # ============================================================================= # Mixin class # ============================================================================= @@ -509,10 +563,11 @@ def model_tester(self): """The model tester instance (e.g., CausalLMModelTester).""" ... - def _skip_if_insufficient_devices(self): + def _skip_if_insufficient_devices(self, world_size=None): + world_size = world_size or self.fsdp_nproc_per_node available_workers = _get_available_fsdp_workers() - if available_workers < self.fsdp_nproc_per_node: - self.skipTest(f"Need at least {self.fsdp_nproc_per_node} FSDP workers, have {available_workers}") + if available_workers < world_size: + self.skipTest(f"Need at least {world_size} FSDP workers, have {available_workers}") def _skip_if_mps(self): if torch._C._get_accelerator().type == "mps": @@ -558,9 +613,10 @@ def _get_tiny_config(self): config.vocab_size_per_layer_input = config.vocab_size return type(config), config.to_diff_dict() - def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_kwargs): + def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, world_size=None, **test_kwargs): + world_size = world_size or self.fsdp_nproc_per_node self._skip_if_mps() - self._skip_if_insufficient_devices() + self._skip_if_insufficient_devices(world_size) self._skip_if_fsdp_distributed_not_enabled() config_class, config_dict = self._get_tiny_config() @@ -575,8 +631,8 @@ def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_k try: mp.spawn( _fsdp_global_wrapper, - args=(test_name, test_impl, func_args, test_kwargs, self.fsdp_nproc_per_node, port, results_file), - nprocs=self.fsdp_nproc_per_node, + args=(test_name, test_impl, func_args, test_kwargs, world_size, port, results_file), + nprocs=world_size, ) with open(results_file) as f: @@ -636,3 +692,13 @@ def test_fsdp2_plan_vs_ddp(self, label): _test_fsdp2_plan_vs_ddp_impl, label == "tied", ) + + @is_fsdp_test + def test_fsdp2_expert_parallel_2d_vs_ddp(self): + """Training on a 2-D (fsdp, tp) mesh with expert parallelism traces DDP step by step.""" + config = self.model_tester.get_config() + if getattr(config, "base_model_ep_plan", None) is None: + self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)") + self._run_fsdp2_distributed_test( + "fsdp2_expert_parallel_2d_vs_ddp", _test_fsdp2_expert_parallel_2d_vs_ddp_impl, world_size=4 + ) From 1c66acbc5adf0133ecebf025886658b21fc84c99 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 3 Sep 2026 19:54:59 +0000 Subject: [PATCH 07/52] Gather the sharded state dict through the model's own save helpers --- src/transformers/trainer.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 49c370beeaeb..f4d0baa3f714 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -4018,23 +4018,20 @@ def save_model(self, output_dir: str | None = None, _internal_call: bool = False self.accelerator.unwrap_model(self.model, keep_torch_compile=False), "_device_mesh", None ) is not None and not _is_peft_model(self.model): # The model was sharded at load time (`DistributedConfig`). Gathering its DTensor state - # dict is collective, so every rank must participate; only the main process writes. + # dict is collective, so every rank must take part; only the main process writes. # (PEFT models fall through to the adapter-only save below.) - from .distributed.tensor_parallel import gather_state_dict_for_save - unwrapped = self.accelerator.unwrap_model(self.model, keep_torch_compile=False) - gather_start = time.time() - state_dict = gather_state_dict_for_save(unwrapped.state_dict(), None, None, 0) - gather_s = time.time() - gather_start + distributed_config = unwrapped.config.distributed_config + state_dict = unwrapped.gather_sharded_state_dict_for_save( + unwrapped, unwrapped.state_dict(), distributed_config, save_on_this_rank=self.args.should_save + ) if self.args.should_save: - write_start = time.time() # `save_pretrained` ends with `barrier_after_gathered_checkpoint_save`. self._save(output_dir, state_dict=state_dict) - logger.info(f"Sharded-model save: gather {gather_s:.1f}s, write {time.time() - write_start:.1f}s") else: # Match the writer's end-of-save barrier so no rank runs ahead into the next # step's collectives while the main process is still writing. - unwrapped.barrier_after_gathered_checkpoint_save(unwrapped.config.distributed_config) + unwrapped.barrier_after_gathered_checkpoint_save(distributed_config) elif self.args.should_save: self._save(output_dir) From 38aefbf2214d18dc60b0d7bf717afc90c8b65551 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 3 Sep 2026 19:55:18 +0000 Subject: [PATCH 08/52] Document combining expert parallelism with FSDP2 --- docs/source/en/expert_parallelism.md | 26 ++++++++++++++++++++++++++ docs/source/en/fsdp.md | 3 +++ 2 files changed, 29 insertions(+) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index f66d993495e6..7bab29ddf965 100644 --- a/docs/source/en/expert_parallelism.md +++ b/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 shard the rest of the model with [FSDP2](./fsdp) on a second mesh dimension. + +```py +distributed_config = DistributedConfig( + tp_size=8, # 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: the experts are sharded across the `tp` dimension, all other parameters are fully sharded across the `fsdp` dimension, and `tp_size * fsdp_size` must equal the number of processes. The experts see the same batch on every `fsdp` rank, so their gradients are averaged over that dimension by the [`Trainer`] (FSDP2 only reduces what it shards). Nothing else changes: train with the [`Trainer`] as usual, and `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` | 3023 | 40.9 GB | +| `tp_size=4, fsdp_size=2` | 2326 | 34.2 GB | +| `tp_size=2, fsdp_size=4` | 2244 | 32.3 GB | + +> [!WARNING] +> Resuming from a checkpoint is not supported yet for models sharded at load time. + [[autodoc]] DistributedConfig diff --git a/docs/source/en/fsdp.md b/docs/source/en/fsdp.md index e0727b77d897..dbbd629f22d1 100644 --- a/docs/source/en/fsdp.md +++ b/docs/source/en/fsdp.md @@ -121,6 +121,9 @@ TrainingArguments( +> [!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. From e72b117d836a815b6adfbd8f476e6cdd8a11ec54 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 3 Sep 2026 20:10:32 +0000 Subject: [PATCH 09/52] Address review: drop the dead expert-gradient sync, torch-native mixed-mesh norm, single collective save FSDP2 composes over the tp-sharded experts and shards them across fsdp as well, so nothing is replicated over fsdp and the Trainer-side gradient averaging never ran; remove it and describe the actual layout. The gradient norm is now one get_total_norm per mesh, each reduced over its own mesh. save_model runs save_pretrained on every rank so its gather is collective and only rank 0 writes; the FSDP branch of the gather (full state dict on rank 0 only) now also covers the 2-D mesh. ParallelismConfig keeps a user-supplied config and only claims what the model was loaded with; pipeline parallelism is rejected together with tp/fsdp; optimizer checkpoints are refused for models sharded at load time since they cannot be resumed. --- docs/source/en/expert_parallelism.md | 8 +- .../distributed/configuration_utils.py | 7 +- src/transformers/distributed/mixin.py | 32 ++-- src/transformers/distributed/utils.py | 2 +- src/transformers/trainer.py | 139 ++++++------------ src/transformers/trainer_optimizer.py | 2 +- tests/test_fsdp_mixin.py | 3 +- 7 files changed, 76 insertions(+), 117 deletions(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index 7bab29ddf965..e79e8c2082f4 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -52,18 +52,18 @@ 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 shard the rest of the model with [FSDP2](./fsdp) on a second mesh dimension. +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=8, # expert parallel size + 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: the experts are sharded across the `tp` dimension, all other parameters are fully sharded across the `fsdp` dimension, and `tp_size * fsdp_size` must equal the number of processes. The experts see the same batch on every `fsdp` rank, so their gradients are averaged over that dimension by the [`Trainer`] (FSDP2 only reduces what it shards). Nothing else changes: train with the [`Trainer`] as usual, and `save_model` gathers the sharded weights and writes a regular checkpoint. +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 steps the optimizer per parameter), and `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: @@ -74,6 +74,6 @@ On 8 GPUs, full fine-tuning of Qwen3-30B-A3B in bf16 at sequence length 2048: | `tp_size=2, fsdp_size=4` | 2244 | 32.3 GB | > [!WARNING] -> Resuming from a checkpoint is not supported yet for models sharded at load time. +> 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 diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index a18a6f9a1dda..e43e5e479281 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/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 diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index ff7834cf62ef..f2dd2d4510a3 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -164,14 +164,17 @@ def prepare_distribute_model( f"is not equal to world_size ({world_size})" ) + 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) - if distributed_config.tp_size > 1 and distributed_config.tp_plan is None: - distributed_config.tp_plan = "auto" elif distributed_config.tp_size > 1: - if distributed_config.tp_plan is None: - distributed_config.tp_plan = "auto" device_map, device_mesh = initialize_tensor_parallelism( distributed_config.tp_plan, tp_size=distributed_config.tp_size, @@ -199,7 +202,8 @@ def maybe_distribute_model( model._tp_size = distributed_config.tp_size model._fsdp_size = distributed_config.fsdp_size - # Both may apply: experts/attention are sharded across `tp`, the rest across `fsdp`. + # 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): @@ -273,15 +277,9 @@ 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. " @@ -289,6 +287,14 @@ def gather_sharded_state_dict_for_save( ) 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: diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py index 50cf27bda817..6feb97c5e730 100644 --- a/src/transformers/distributed/utils.py +++ b/src/transformers/distributed/utils.py @@ -193,7 +193,7 @@ 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 or 1 + 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. diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index f4d0baa3f714..c19f8b400c74 100755 --- a/src/transformers/trainer.py +++ b/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 @@ -615,6 +616,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) @@ -751,26 +761,21 @@ def _build_accelerator_args(self, **kwargs) -> dict[str, Any]: ) args["parallelism_config"] = self.args.parallelism_config + model_tp_size = getattr(self.model, "tp_size", None) or 1 model_fsdp_size = getattr(self.model, "fsdp_size", None) or 1 - if model_fsdp_size > 1: - # 2-D mesh: accelerate has to know about both dimensions or it will try to wrap the - # already-sharded model in DDP. - from accelerate import ParallelismConfig - - args["parallelism_config"] = ParallelismConfig( - dp_shard_size=model_fsdp_size, tp_size=getattr(self.model, "tp_size", None) or 1 - ) - elif 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 + 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, or it sees unaccounted ranks and wraps + # the DTensor model in DDP, which raises. + 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 @@ -1871,7 +1876,6 @@ def _run_epoch( self._track_num_input_tokens(inputs) if do_sync_step: - self._sync_replicated_grads(model) grad_norm = None if self.args.max_grad_norm > 0: grad_norm = self._clip_grad_norm(model) @@ -2628,64 +2632,26 @@ 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 _sync_replicated_grads(self, model): - """Average the gradients of parameters replicated along the data-parallel dimension (e.g. the - EP-sharded experts under a 2-D mesh) -- FSDP2 only reduces what it shards.""" - from torch.distributed.tensor import DTensor - - mesh = getattr(model, "_device_mesh", None) or getattr(getattr(model, "module", None), "_device_mesh", None) - if mesh is None or mesh.ndim < 2 or "fsdp" not in (mesh.mesh_dim_names or ()): - return - dp_group = mesh["fsdp"].get_group() - dp_size = mesh["fsdp"].size() - if dp_size < 2: - return - - for param in model.parameters(): - if param.grad is None: - continue - grad = param.grad - if isinstance(grad, DTensor): - # Already reduced over `fsdp` if that dimension is part of the gradient's own mesh. - if "fsdp" in (grad.device_mesh.mesh_dim_names or ()): - continue - grad = grad.to_local() - torch.distributed.all_reduce(grad, group=dp_group) - grad.div_(dp_size) - def _mixed_mesh_grad_norm(self, model, max_norm): - """Gradient norm (and clip) when parameters live on different device meshes: sum each gradient's - local squared norm, discounting the mesh dimensions it is replicated over, reduce 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 - grads = [p.grad for p in model.parameters() if p.grad is not None] - if not grads: - return torch.zeros((), device=self.args.device) - - world = torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1 - total_sq = torch.zeros((), device=grads[0].device, dtype=torch.float32) - for g in grads: - if isinstance(g, DTensor): - local = g.to_local() - # Ranks along a replicated mesh dimension all hold the same values. - replicas = 1 - for dim, placement in enumerate(g.placements): - if placement.is_replicate(): - replicas *= g.device_mesh.size(dim) - # Dimensions of the world this tensor's mesh does not span also hold copies. - replicas *= world // g.device_mesh.size() - else: - local, replicas = g, world # a plain tensor is the same on every rank - total_sq += local.detach().float().pow(2).sum() / replicas + params_by_mesh = defaultdict(list) + for param in model.parameters(): + if param.grad is not None: + params_by_mesh[param.grad.device_mesh if isinstance(param.grad, DTensor) else None].append(param) - if world > 1: - torch.distributed.all_reduce(total_sq) - total_norm = total_sq.sqrt() + 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)) if max_norm != float("inf"): - clip = (max_norm / (total_norm + 1e-6)).clamp(max=1.0) - for g in grads: - (g.to_local() if isinstance(g, DTensor) else g).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: @@ -4014,24 +3980,11 @@ 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.accelerator.unwrap_model(self.model, keep_torch_compile=False), "_device_mesh", None - ) is not None and not _is_peft_model(self.model): - # The model was sharded at load time (`DistributedConfig`). Gathering its DTensor state - # dict is collective, so every rank must take part; only the main process writes. - # (PEFT models fall through to the adapter-only save below.) - unwrapped = self.accelerator.unwrap_model(self.model, keep_torch_compile=False) - distributed_config = unwrapped.config.distributed_config - state_dict = unwrapped.gather_sharded_state_dict_for_save( - unwrapped, unwrapped.state_dict(), distributed_config, save_on_this_rank=self.args.should_save - ) - if self.args.should_save: - # `save_pretrained` ends with `barrier_after_gathered_checkpoint_save`. - self._save(output_dir, state_dict=state_dict) - else: - # Match the writer's end-of-save barrier so no rank runs ahead into the next - # step's collectives while the main process is still writing. - unwrapped.barrier_after_gathered_checkpoint_save(distributed_config) + 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) @@ -4042,10 +3995,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()`. @@ -4066,6 +4019,8 @@ def _save(self, output_dir: str | None = None, state_dict: dict | None = None) - else: self.model.save_pretrained(output_dir, state_dict=state_dict) + if not self.args.should_save: + return # a non-writer rank of a model sharded at load time, only here for the collectives above if self.processing_class is not None: self.processing_class.save_pretrained(output_dir) elif ( diff --git a/src/transformers/trainer_optimizer.py b/src/transformers/trainer_optimizer.py index bc78622cdc29..d5acb6fd2dd6 100644 --- a/src/transformers/trainer_optimizer.py +++ b/src/transformers/trainer_optimizer.py @@ -216,7 +216,7 @@ def has_mixed_dtensor(tensors) -> bool: """Whether `tensors` do not all share one device mesh, so whole-set ops (fused/foreach kernels, `clip_grad_norm_`) cannot span them. 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 `tp` mesh and the rest on the `fsdp` mesh.""" + on the full mesh and the rest on the `fsdp` sub-mesh.""" from torch.distributed.tensor import DTensor meshes = {t.device_mesh if isinstance(t, DTensor) else None for t in tensors} diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index ea658341ce40..6f31f2d7429e 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -563,8 +563,7 @@ def model_tester(self): """The model tester instance (e.g., CausalLMModelTester).""" ... - def _skip_if_insufficient_devices(self, world_size=None): - world_size = world_size or self.fsdp_nproc_per_node + def _skip_if_insufficient_devices(self, world_size): available_workers = _get_available_fsdp_workers() if available_workers < world_size: self.skipTest(f"Need at least {world_size} FSDP workers, have {available_workers}") From 449acd211d512dbb3aed02ae4cb86f491e778e47 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 3 Sep 2026 20:16:20 +0000 Subject: [PATCH 10/52] Give each device mesh its own optimizer param group instead of stepping per parameter Fused/foreach AdamW cannot span parameters on different meshes, but it can run per mesh: 189 ms/step per-parameter vs 41 ms fused per mesh group on Qwen3-30B-A3B at tp=4 x fsdp=2. --- src/transformers/trainer.py | 15 +++++++++++++-- src/transformers/trainer_optimizer.py | 6 +----- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index c19f8b400c74..16bea712cbd5 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -97,6 +97,7 @@ _OPTIMIZER_HANDLERS, OptimizerContext, _parse_optim_args, + has_mixed_dtensor, is_optimizer_factory, ) from .trainer_pt_utils import ( @@ -1271,6 +1272,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 @@ -2658,8 +2671,6 @@ 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 diff --git a/src/transformers/trainer_optimizer.py b/src/transformers/trainer_optimizer.py index d5acb6fd2dd6..c4347a5c0d36 100644 --- a/src/transformers/trainer_optimizer.py +++ b/src/transformers/trainer_optimizer.py @@ -205,16 +205,12 @@ 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}) - # Expert parallelism leaves the parameters on different meshes, 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` do not all share one device mesh, so whole-set ops (fused/foreach kernels, - `clip_grad_norm_`) cannot span them. Expert parallelism leaves the experts sharded and everything + `clip_grad_norm_`) cannot span them; the Trainer then groups 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 From 9d17bd7aacf922fae1bb634fdd9be2cd051fa4ae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 3 Sep 2026 20:19:13 +0000 Subject: [PATCH 11/52] Test the 2-D mesh with a batch split across the fsdp ranks and compare gradient norms --- tests/test_fsdp_mixin.py | 61 ++++++++++++++++++++++++++++++++-------- 1 file changed, 50 insertions(+), 11 deletions(-) diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index 6f31f2d7429e..5863140e4f3b 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -494,9 +494,26 @@ def _test_fsdp2_plan_vs_ddp_impl(rank, config_class, config_dict, tie_word_embed logger.debug("DDP and FSDP2 comparison checks passed.") +def _grad_norm_across_meshes(model): + """Total gradient norm of parameters living on different device meshes (what the Trainer does).""" + from torch.distributed.tensor import DTensor + from torch.nn.utils import get_total_norm + + grads_by_mesh = {} + for param in model.parameters(): + if param.grad is not None: + grads_by_mesh.setdefault(param.grad.device_mesh if isinstance(param.grad, DTensor) else None, []).append( + param.grad + ) + norms = [get_total_norm(grads) for grads in grads_by_mesh.values()] + norms = [n.full_tensor() if isinstance(n, DTensor) else n for n in norms] + return torch.linalg.vector_norm(torch.stack(norms)) + + def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, dtype=None): - """DDP vs a 2-D (fsdp, tp) mesh with expert parallelism on `tp`: same batch on every rank, so the - sharded run has to trace the DDP run step by step and end on the same weights.""" + """DDP vs a 2-D (fsdp, tp) mesh with expert parallelism on `tp`. DDP sees the whole batch on every + rank; each `fsdp` rank of the 2-D run sees its own slice of it, so FSDP2's reduction over `fsdp` + is exercised. Losses, gradient norms and final weights have to match step by step.""" init_test_logger() if dtype is None: @@ -504,25 +521,39 @@ def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, device = _get_rank_device(rank) config = config_class.from_dict(config_dict) - batches = _build_repeated_training_batches(config, device, NUM_STEPS) + world_size = dist.get_world_size() + dp = world_size // 2 + generator = torch.Generator(device=device) + generator.manual_seed(SEED) + input_ids = torch.randint(0, config.vocab_size, (dp * BATCH_SIZE, SEQ_LEN), device=device, generator=generator) + batches = [(input_ids, input_ids.clone())] * NUM_STEPS with _deterministic_init_model_dir(rank, config, dtype) as init_model_dir: - ddp_losses, _, ddp_state_dict = train_ddp(rank, batches, LR, device, dtype, init_model_dir) + ddp_losses, ddp_grad_norms, ddp_state_dict = train_ddp(rank, batches, LR, device, dtype, init_model_dir) _set_determinism(SEED) - world_size = dist.get_world_size() model = AutoModelForCausalLM.from_pretrained( init_model_dir, torch_dtype=dtype, - distributed_config=DistributedConfig(tp_size=2, fsdp_size=world_size // 2, enable_expert_parallel=True), + distributed_config=DistributedConfig(tp_size=2, fsdp_size=dp, enable_expert_parallel=True), ) - assert model.tp_size == 2 and model.fsdp_size == world_size // 2 + assert model.tp_size == 2 and model.fsdp_size == dp assert model._device_mesh.mesh_dim_names == ("fsdp", "tp") model.train() optimizer = torch.optim.Adam(model.parameters(), lr=LR, foreach=False) - # `clip_grad_norm_` cannot span parameters living on different meshes; the Trainer has its own - # norm for that case, so only the losses and the final weights are compared here. - losses, _ = _run_training_steps(model, optimizer, batches, track_grad_norms=False) + dp_rank = model._device_mesh["fsdp"].get_local_rank() + dp_group = model._device_mesh["fsdp"].get_group() + losses, grad_norms = [], [] + for ids, labels in batches: + rows = slice(dp_rank * BATCH_SIZE, (dp_rank + 1) * BATCH_SIZE) + optimizer.zero_grad() + loss = model(input_ids=ids[rows], labels=labels[rows], use_cache=False).loss + loss.backward() + grad_norms.append(_grad_norm_across_meshes(model).item()) + optimizer.step() + loss = loss.detach() + dist.all_reduce(loss, group=dp_group) + losses.append(loss.item() / dp) state_dict = gather_full_state_dict(model) for step in range(len(ddp_losses)): @@ -533,6 +564,13 @@ def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, atol=DDP_FSDP_ATOL, msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, FSDP2+EP={losses[step]}", ) + torch.testing.assert_close( + torch.tensor(ddp_grad_norms[step]), + torch.tensor(grad_norms[step]), + rtol=DDP_FSDP_RTOL, + atol=DDP_FSDP_ATOL, + msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, FSDP2+EP={grad_norms[step]}", + ) for key in ddp_state_dict: assert key in state_dict, f"Key {key} missing from FSDP2+EP state dict" @@ -694,7 +732,8 @@ def test_fsdp2_plan_vs_ddp(self, label): @is_fsdp_test def test_fsdp2_expert_parallel_2d_vs_ddp(self): - """Training on a 2-D (fsdp, tp) mesh with expert parallelism traces DDP step by step.""" + """Training on a 2-D (fsdp, tp) mesh with expert parallelism, each fsdp rank on its own slice of + the batch, traces DDP on the whole batch step by step.""" config = self.model_tester.get_config() if getattr(config, "base_model_ep_plan", None) is None: self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)") From 6f2d9c4bf87ba5f0a1b4e715bc281217a3a8f73b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 3 Sep 2026 20:19:59 +0000 Subject: [PATCH 12/52] Refresh the throughput table --- docs/source/en/expert_parallelism.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index e79e8c2082f4..6e45e28bd27d 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -69,9 +69,9 @@ 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` | 3023 | 40.9 GB | -| `tp_size=4, fsdp_size=2` | 2326 | 34.2 GB | -| `tp_size=2, fsdp_size=4` | 2244 | 32.3 GB | +| `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. From 3a8da11954d9c58ce09a5b3b000620ff9143f016 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 3 Sep 2026 23:33:32 +0000 Subject: [PATCH 13/52] Cosmetic: docstring layout, comment placement, doc link --- docs/source/en/expert_parallelism.md | 2 +- src/transformers/trainer.py | 10 +++++++--- src/transformers/trainer_optimizer.py | 10 ++++++---- tests/test_fsdp_mixin.py | 21 ++++++++++++--------- 4 files changed, 26 insertions(+), 17 deletions(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index 6e45e28bd27d..c62eef296119 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -63,7 +63,7 @@ distributed_config = DistributedConfig( 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 steps the optimizer per parameter), and `save_model` gathers the sharded weights and writes a regular checkpoint. +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: diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 16bea712cbd5..b08ba795bfda 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2646,8 +2646,10 @@ 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 the gradients live on different device meshes, which - `clip_grad_norm_` cannot span: one norm per mesh, each already reduced over its own mesh.""" + """ + 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 @@ -4030,8 +4032,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 # a non-writer rank of a model sharded at load time, only here for the collectives above + return + if self.processing_class is not None: self.processing_class.save_pretrained(output_dir) elif ( diff --git a/src/transformers/trainer_optimizer.py b/src/transformers/trainer_optimizer.py index c4347a5c0d36..13ab9edc18da 100644 --- a/src/transformers/trainer_optimizer.py +++ b/src/transformers/trainer_optimizer.py @@ -209,10 +209,12 @@ def _get_adamw_torch(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]: def has_mixed_dtensor(tensors) -> bool: - """Whether `tensors` do not all share one device mesh, so whole-set ops (fused/foreach kernels, - `clip_grad_norm_`) cannot span them; the Trainer then groups 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.""" + """ + 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 meshes = {t.device_mesh if isinstance(t, DTensor) else None for t in tensors} diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index 5863140e4f3b..e780b74f56c9 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -23,6 +23,7 @@ import time import traceback from abc import ABC, abstractmethod +from collections import defaultdict from contextlib import contextmanager from parameterized import parameterized @@ -499,21 +500,21 @@ def _grad_norm_across_meshes(model): from torch.distributed.tensor import DTensor from torch.nn.utils import get_total_norm - grads_by_mesh = {} + grads_by_mesh = defaultdict(list) for param in model.parameters(): if param.grad is not None: - grads_by_mesh.setdefault(param.grad.device_mesh if isinstance(param.grad, DTensor) else None, []).append( - param.grad - ) + grads_by_mesh[param.grad.device_mesh if isinstance(param.grad, DTensor) else None].append(param.grad) norms = [get_total_norm(grads) for grads in grads_by_mesh.values()] norms = [n.full_tensor() if isinstance(n, DTensor) else n for n in norms] return torch.linalg.vector_norm(torch.stack(norms)) def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, dtype=None): - """DDP vs a 2-D (fsdp, tp) mesh with expert parallelism on `tp`. DDP sees the whole batch on every - rank; each `fsdp` rank of the 2-D run sees its own slice of it, so FSDP2's reduction over `fsdp` - is exercised. Losses, gradient norms and final weights have to match step by step.""" + """ + DDP vs a 2-D (fsdp, tp) mesh with expert parallelism on `tp`. DDP sees the whole batch on every rank; each `fsdp` + rank of the 2-D run sees its own slice of it, so FSDP2's reduction over `fsdp` is exercised. Losses, gradient norms + and final weights have to match step by step. + """ init_test_logger() if dtype is None: @@ -732,8 +733,10 @@ def test_fsdp2_plan_vs_ddp(self, label): @is_fsdp_test def test_fsdp2_expert_parallel_2d_vs_ddp(self): - """Training on a 2-D (fsdp, tp) mesh with expert parallelism, each fsdp rank on its own slice of - the batch, traces DDP on the whole batch step by step.""" + """ + Training on a 2-D (fsdp, tp) mesh with expert parallelism, each fsdp rank on its own slice of the batch, + traces DDP on the whole batch step by step. + """ config = self.model_tester.get_config() if getattr(config, "base_model_ep_plan", None) is None: self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)") From 1b071993173e52ffbd5005b95f2f195ca7669d60 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 00:48:23 +0000 Subject: [PATCH 14/52] Expert-parallel token dispatch: each rank trains on its own part of the batch DistributedConfig(expert_parallel_dispatch=True) sends every selected (token, expert) pair to the rank that owns the expert with an all-to-all, runs the local experts, and sends the results back, instead of running the whole batch on every expert-parallel rank and all-reducing the outputs. Each rank trains on its own part of the batch: the parameters outside the experts are sharded with FSDP2 across every rank (the flattened (fsdp, tp) mesh) so FSDP2 owns their gradient reduction, the experts stay sharded across tp and, if set, fsdp, and the all-to-all backward accumulates their gradients across the group (scaled by 1/ep_size to match the data-parallel average). The Trainer gives each rank its own batches and counts tokens across all of them. --- docs/source/en/expert_parallelism.md | 14 +++ .../distributed/configuration_utils.py | 8 ++ src/transformers/distributed/fsdp.py | 22 +++- src/transformers/distributed/mixin.py | 15 ++- .../distributed/tensor_parallel.py | 30 ++++- src/transformers/integrations/moe.py | 107 ++++++++++++++++++ src/transformers/trainer.py | 15 ++- tests/test_fsdp_mixin.py | 37 +++--- tests/test_tensor_parallel_mixin.py | 29 +++-- 9 files changed, 241 insertions(+), 36 deletions(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index c62eef296119..d62332ae4c08 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -50,6 +50,20 @@ Launch your inference script with [torchrun](https://pytorch.org/docs/stable/ela torchrun --nproc-per-node 8 your_script.py ``` +## Token dispatch + +By default every expert-parallel rank runs the whole batch, keeps only the experts it owns, and the group all-reduces the expert outputs at the end of every MoE layer. The parameters outside the experts do `tp_size` times the same work, and the all-reduce moves the full activations. `expert_parallel_dispatch=True` sends each token to the rank that owns its experts instead. + +```py +distributed_config = DistributedConfig( + tp_size=8, + enable_expert_parallel=True, + expert_parallel_dispatch=True, +) +``` + +Every rank trains on its own part of the batch: at each MoE layer it routes its tokens, sends the selected (token, expert) pairs to their owners with an all-to-all, runs its local experts, and gets the results back with a second all-to-all. Only the routed tokens travel. The parameters outside the experts are data-parallel across the expert-parallel group, so they are sharded with [FSDP2](./fsdp) across every rank (the `fsdp` and `tp` dimensions together when both are set) and FSDP2 reduces their gradients; the experts stay sharded across `tp` and, with `fsdp_size > 1`, across `fsdp`. The [`Trainer`] gives each rank its own batches and counts tokens across all of them. + ## 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. diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index e43e5e479281..b81611a171c6 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -34,6 +34,10 @@ class DistributedConfig: Reserved for sequence parallelism. Not wired up yet. enable_expert_parallel (`bool`, *optional*, defaults to `False`): Route MoE models through the expert-parallel path (``base_model_ep_plan``). + expert_parallel_dispatch (`bool`, *optional*, defaults to `False`): + Send each token to the rank that owns its experts with an all-to-all instead of running the whole batch + on every expert-parallel rank. Each rank then trains on its own part of the batch, and the parameters + that are not expert-parallel are sharded with FSDP2 across every rank. Requires `enable_expert_parallel`. fsdp_size (`int`, *optional*): Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1. fsdp_cpu_offload (`bool`, *optional*, defaults to `False`): @@ -48,6 +52,7 @@ class DistributedConfig: tp_plan: dict[str, str] | Literal["auto"] | None = None enable_sequence_parallel: bool = False enable_expert_parallel: bool = False + expert_parallel_dispatch: bool = False fsdp_size: int | None = None fsdp_cpu_offload: bool = False fsdp_mixed_precision: bool = False @@ -73,6 +78,9 @@ def __post_init__(self): elif self.tp_size is None: self.tp_size = 1 + if self.expert_parallel_dispatch and not self.enable_expert_parallel: + raise ValueError("`expert_parallel_dispatch=True` requires `enable_expert_parallel=True`.") + if self.pp_size > 1 and (self.tp_size > 1 or self.fsdp_size > 1): raise ValueError( "Pipeline parallelism cannot be combined with tensor or FSDP parallelism yet. " diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 3ea0e294bf04..458adfec1662 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -34,6 +34,7 @@ if is_torch_distributed_available() and is_torch_greater_or_equal("2.6"): from torch.distributed._composable.fsdp import fully_shard from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy + from torch.distributed.tensor import DTensor logger = logging.get_logger(__name__) @@ -185,13 +186,18 @@ def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None) def apply_fully_sharded_data_parallelism( - model: nn.Module, fsdp_mesh: torch.distributed.device_mesh.DeviceMesh + model: nn.Module, + fsdp_mesh: torch.distributed.device_mesh.DeviceMesh, + expert_mesh: torch.distributed.device_mesh.DeviceMesh | None = None, ) -> nn.Module: """ Apply FSDP2 (fully_shard) to a model. Torch availability, distributed initialization and the version requirement are asserted upstream by `initialize_fully_sharded_data_parallelism`. + + With expert-parallel token dispatch the experts are already sharded across the expert-parallel ranks, which + `fsdp_mesh` spans: pass the mesh they should be sharded across as `expert_mesh` (`None` leaves them as they are). """ fsdp_plan = dict(getattr(model, "_fsdp_plan", None) or {}) if not fsdp_plan: @@ -206,6 +212,20 @@ def apply_fully_sharded_data_parallelism( adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, model) reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) + if expert_mesh is not None or fsdp_mesh.ndim == 1 and any(isinstance(p, DTensor) for p in model.parameters()): + # The DTensor parameters are the expert-parallel experts: shard them across `expert_mesh` in their own + # group, or keep them out of the groups below when there is no such mesh. + expert_modules = [ + module + for module in model.modules() + if any(isinstance(p, DTensor) for p in module.parameters(recurse=False)) + ] + if expert_mesh is not None: + for module in expert_modules: + fully_shard(module, mesh=expert_mesh, reshard_after_forward=True, **fsdp_policy_kwargs) + else: + fsdp_policy_kwargs["ignored_params"] = {p for module in expert_modules for p in module.parameters()} + for module_name, module in reshard_targets: fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=True, **fsdp_policy_kwargs) logger.debug(f"Applied fully_shard to {module_name} (reshard=True)") diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index f2dd2d4510a3..32241239ead1 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -210,7 +210,14 @@ def maybe_distribute_model( model.tp_plan = distributed_config.tp_plan model = apply_tensor_parallelism(model, tp_mesh) - if distributed_config.fsdp_size > 1: + if distributed_config.expert_parallel_dispatch: + # Every expert-parallel rank trains on its own part of the batch, so the parameters outside the + # experts are data-parallel across the whole mesh: FSDP2 shards them across all of it and owns + # their gradient reduction. The experts stay sharded across `tp` and, if any, across `fsdp`. + trunk_mesh = device_mesh._flatten("fsdp_tp") if device_mesh.ndim > 1 else device_mesh + expert_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else None + model = apply_fully_sharded_data_parallelism(model, trunk_mesh, expert_mesh=expert_mesh) + elif 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) @@ -277,9 +284,9 @@ def gather_sharded_state_dict_for_save( if distributed_config is None: 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 is_fsdp_managed_module(model_to_save): + # Also covers the 2-D (fsdp, tp) mesh and token dispatch: 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. " diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 6a9455637a33..bb65490e9a3d 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -144,7 +144,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, *, expert_parallel_dispatch=False): """Install pre / around / post transforms by replacing module.forward.""" original_forward = module.forward @@ -344,7 +344,7 @@ class ReplicatedWithGradAllReduce(TensorParallelLayer): summed across the mesh. """ - def install_forward(self, module, mesh): + def install_forward(self, module, mesh, *, expert_parallel_dispatch=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 +415,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, *, expert_parallel_dispatch=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( @@ -604,7 +604,7 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_ return (hidden_states, *routing_args), kwargs - def install_forward(self, module, mesh, *, is_expert_parallel=False): + def install_forward(self, module, mesh, *, is_expert_parallel=False, expert_parallel_dispatch=False): """Install the transforms but pass `is_expert_parallel` in the forward call.""" original_forward = module.forward output_source = ( @@ -617,6 +617,17 @@ def install_forward(self, module, mesh, *, is_expert_parallel=False): ) def tp_forward(*args, **kwargs): + if expert_parallel_dispatch: + from ..integrations.moe import dispatch_experts_forward + + # Every rank routes its own tokens to the experts' owners; nothing is replicated across + # the group, so there is neither an input gradient sum nor an output all-reduce. + hidden_states, top_k_index, top_k_weights = args[:3] + 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) args, kwargs = self.transform_inputs_pre_forward( module, args, kwargs, mesh, is_expert_parallel=is_expert_parallel ) @@ -695,7 +706,14 @@ class EpRouterParallel(TensorParallelLayer): Scores and indices stay paired element-wise in `(seq, top_k)` shape. """ + def install_forward(self, module, mesh, *, expert_parallel_dispatch=False): + module.expert_parallel_dispatch = expert_parallel_dispatch + return super().install_forward(module, mesh) + def transform_output_post_forward(self, module, output, mesh): + if module.expert_parallel_dispatch: + # The experts forward sends every token to its experts' owners, so it needs the global ids. + return output ep_rank, ep_size = mesh.get_local_rank(), mesh.size() num_experts = getattr(module, "num_experts", None) if num_experts is None: @@ -821,7 +839,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, expert_parallel_dispatch=model.config.distributed_config.expert_parallel_dispatch + ) module._is_hooked = True return model diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index fb3b736c0301..5374401a6717 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -374,6 +374,113 @@ def _grouped_linear( return out +class _AllToAll(torch.autograd.Function): + """`all_to_all_single` with uneven splits whose backward is the reverse exchange.""" + + @staticmethod + def forward(ctx, tensor, output_split_sizes, input_split_sizes, group): + ctx.group, ctx.output_split_sizes, ctx.input_split_sizes = group, output_split_sizes, input_split_sizes + output = tensor.new_empty(sum(output_split_sizes), *tensor.shape[1:]) + torch.distributed.all_to_all_single( + output, + tensor.contiguous(), + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group, + ) + return output + + @staticmethod + def backward(ctx, grad_output): + grad_input = grad_output.new_empty(sum(ctx.input_split_sizes), *grad_output.shape[1:]) + torch.distributed.all_to_all_single( + grad_input, + grad_output.contiguous(), + output_split_sizes=ctx.input_split_sizes, + input_split_sizes=ctx.output_split_sizes, + group=ctx.group, + ) + return grad_input, None, None, None + + +class _ScaleGrad(torch.autograd.Function): + """Identity whose backward scales the gradient.""" + + @staticmethod + def forward(ctx, tensor, scale): + ctx.scale = scale + return tensor + + @staticmethod + def backward(ctx, grad_output): + return grad_output * ctx.scale, None + + +def dispatch_experts_forward( + self: torch.nn.Module, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, + ep_mesh, +) -> torch.Tensor: + """ + Expert-parallel forward by token dispatch. Every rank routes its own tokens, sends each selected (token, expert) + pair to the rank that owns the expert with an all-to-all, runs its local experts on what it receives, sends the + results back and combines them with the routing weights. Each rank trains on its own part of the batch, so the + parameters outside the experts are data-parallel across the group and averaged by FSDP2; the all-to-all backward + already sums every rank's contribution to the expert gradients, which are scaled by `1 / ep_size` to match. + """ + ep_group, ep_size = ep_mesh.get_group(), ep_mesh.size() + num_tokens, hidden_dim = hidden_states.shape + num_top_k = top_k_index.size(-1) + num_local_experts = self.num_experts # the sharding leaves the module with its local expert count + + # Group the selected pairs by the rank that owns their expert. + expert_ids = top_k_index.reshape(-1) + owner = expert_ids // num_local_experts + order = torch.argsort(owner, stable=True) + send_tokens = hidden_states.repeat_interleave(num_top_k, dim=0)[order] + send_expert_ids = (expert_ids % num_local_experts)[order] + send_counts = torch.bincount(owner, minlength=ep_size) + recv_counts = torch.empty_like(send_counts) + torch.distributed.all_to_all_single(recv_counts, send_counts, group=ep_group) + send_sizes, recv_sizes = send_counts.tolist(), recv_counts.tolist() + + # Exchange the tokens (with gradient) and their local expert ids. + recv_tokens = _AllToAll.apply(send_tokens, recv_sizes, send_sizes, ep_group) + recv_expert_ids = send_expert_ids.new_empty(sum(recv_sizes)) + torch.distributed.all_to_all_single( + recv_expert_ids, send_expert_ids, output_split_sizes=recv_sizes, input_split_sizes=send_sizes, group=ep_group + ) + + # Local experts: sort by expert, grouped GEMMs, unsort. + expert_ids_g, perm = torch.sort(recv_expert_ids, stable=True) + tokens_g = recv_tokens[perm] + offsets = torch.cumsum(torch.bincount(expert_ids_g, minlength=num_local_experts), dim=0, dtype=torch.int32) + grad_scale = 1.0 / ep_size + if self.has_gate: + weight = _ScaleGrad.apply(self.gate_up_proj, grad_scale) + bias = _ScaleGrad.apply(self.gate_up_proj_bias, grad_scale)[expert_ids_g] if self.has_bias else None + else: + weight = _ScaleGrad.apply(self.up_proj, grad_scale) + bias = _ScaleGrad.apply(self.up_proj_bias, grad_scale)[expert_ids_g] if self.has_bias else None + proj_out = _grouped_linear(tokens_g, weight, offsets, bias=bias, is_transposed=self.is_transposed) + proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) + weight = _ScaleGrad.apply(self.down_proj, grad_scale) + bias = _ScaleGrad.apply(self.down_proj_bias, grad_scale)[expert_ids_g] if self.has_bias else None + proj_out = _grouped_linear(proj_out, weight, offsets, bias=bias, is_transposed=self.is_transposed) + inverse_perm = torch.empty_like(perm) + inverse_perm[perm] = torch.arange(perm.numel(), device=perm.device) + expert_out = proj_out[inverse_perm] + + # Send the results back to the owners of the tokens and combine them with the routing weights. + recv_out = _AllToAll.apply(expert_out, send_sizes, recv_sizes, ep_group) + combined = recv_out.new_zeros(num_tokens * num_top_k, hidden_dim) + combined[order] = recv_out + combined = combined * top_k_weights.reshape(-1, 1).to(combined.dtype) + return combined.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype) + + def grouped_mm_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index b08ba795bfda..b2f06aee8ce1 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -617,6 +617,10 @@ 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 + # With expert-parallel token dispatch every rank trains on its own part of the batch, while accelerate + # treats the `tp` ranks as one data-parallel rank: the batches and the token counts are handled here. + distributed_config = getattr(getattr(model, "config", None), "distributed_config", None) + self._expert_parallel_dispatch = distributed_config is not None and distributed_config.expert_parallel_dispatch if ( getattr(model, "_device_mesh", None) is not None and args.save_strategy != SaveStrategy.NO @@ -1046,6 +1050,10 @@ def _get_dataloader( if prepared_bs is not None and getattr(prepared_bs, "batch_sampler", None) is sampler: prepared_bs.num_processes = 1 prepared_bs.process_index = 0 + elif self._expert_parallel_dispatch: + # accelerate hands every `tp` rank the same batch; under token dispatch each rank gets its own. + dataloader.batch_sampler.num_processes = self.args.world_size + dataloader.batch_sampler.process_index = self.args.process_index # Store the prepared dataloader for subsequent evaluations if using persistent workers. if dataloader_key is not None and self.args.dataloader_persistent_workers: @@ -2158,7 +2166,8 @@ def compute_loss( # TP and EP-as-TP ranks see replicated batches; `num_processes` over-counts # them by `tp_size`. Mirror the divisor used in `_get_num_items_in_batch`. loss_scale = self.accelerator.num_processes - if (pc := getattr(self.accelerator, "parallelism_config", None)) is not None: + pc = getattr(self.accelerator, "parallelism_config", None) + if pc is not None and not self._expert_parallel_dispatch: loss_scale //= pc.tp_size loss *= loss_scale if self.args.n_gpu <= 1 else self.args.n_gpu @@ -2308,7 +2317,9 @@ def _get_num_items_in_batch(self, batch_samples: list, device: torch.device) -> # In the DataParallel case, convert the scalar tensor into a 2-dim tensor with the same value repeated num_items_in_batch = num_items_in_batch.unsqueeze(0).expand(self.args.n_gpu, -1) # Divide by number of devices with the same batch - if pc := getattr(self.accelerator, "parallelism_config", None): + if ( + pc := getattr(self.accelerator, "parallelism_config", None) + ) and not self._expert_parallel_dispatch: num_items_in_batch = num_items_in_batch // pc.non_data_parallel_size return num_items_in_batch diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index e780b74f56c9..413ce8bff6c0 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -509,11 +509,11 @@ def _grad_norm_across_meshes(model): return torch.linalg.vector_norm(torch.stack(norms)) -def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, dtype=None): +def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, dtype=None, dispatch=False): """ DDP vs a 2-D (fsdp, tp) mesh with expert parallelism on `tp`. DDP sees the whole batch on every rank; each `fsdp` - rank of the 2-D run sees its own slice of it, so FSDP2's reduction over `fsdp` is exercised. Losses, gradient norms - and final weights have to match step by step. + rank of the 2-D run sees its own slice of it (every rank with token dispatch), so FSDP2's reduction is exercised. + Losses, gradient norms and final weights have to match step by step. """ init_test_logger() @@ -524,9 +524,12 @@ def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, config = config_class.from_dict(config_dict) world_size = dist.get_world_size() dp = world_size // 2 + num_slices = world_size if dispatch else dp generator = torch.Generator(device=device) generator.manual_seed(SEED) - input_ids = torch.randint(0, config.vocab_size, (dp * BATCH_SIZE, SEQ_LEN), device=device, generator=generator) + input_ids = torch.randint( + 0, config.vocab_size, (num_slices * BATCH_SIZE, SEQ_LEN), device=device, generator=generator + ) batches = [(input_ids, input_ids.clone())] * NUM_STEPS with _deterministic_init_model_dir(rank, config, dtype) as init_model_dir: @@ -536,17 +539,21 @@ def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, model = AutoModelForCausalLM.from_pretrained( init_model_dir, torch_dtype=dtype, - distributed_config=DistributedConfig(tp_size=2, fsdp_size=dp, enable_expert_parallel=True), + distributed_config=DistributedConfig( + tp_size=2, fsdp_size=dp, enable_expert_parallel=True, expert_parallel_dispatch=dispatch + ), ) assert model.tp_size == 2 and model.fsdp_size == dp assert model._device_mesh.mesh_dim_names == ("fsdp", "tp") model.train() optimizer = torch.optim.Adam(model.parameters(), lr=LR, foreach=False) - dp_rank = model._device_mesh["fsdp"].get_local_rank() - dp_group = model._device_mesh["fsdp"].get_group() + if dispatch: + slice_index, dp_group = rank, dist.group.WORLD + else: + slice_index, dp_group = model._device_mesh["fsdp"].get_local_rank(), model._device_mesh["fsdp"].get_group() losses, grad_norms = [], [] for ids, labels in batches: - rows = slice(dp_rank * BATCH_SIZE, (dp_rank + 1) * BATCH_SIZE) + rows = slice(slice_index * BATCH_SIZE, (slice_index + 1) * BATCH_SIZE) optimizer.zero_grad() loss = model(input_ids=ids[rows], labels=labels[rows], use_cache=False).loss loss.backward() @@ -554,7 +561,7 @@ def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, optimizer.step() loss = loss.detach() dist.all_reduce(loss, group=dp_group) - losses.append(loss.item() / dp) + losses.append(loss.item() / num_slices) state_dict = gather_full_state_dict(model) for step in range(len(ddp_losses)): @@ -731,15 +738,19 @@ def test_fsdp2_plan_vs_ddp(self, label): label == "tied", ) + @parameterized.expand([("masked", False), ("dispatch", True)]) @is_fsdp_test - def test_fsdp2_expert_parallel_2d_vs_ddp(self): + def test_fsdp2_expert_parallel_2d_vs_ddp(self, label, dispatch): """ - Training on a 2-D (fsdp, tp) mesh with expert parallelism, each fsdp rank on its own slice of the batch, - traces DDP on the whole batch step by step. + Training on a 2-D (fsdp, tp) mesh with expert parallelism, each fsdp rank (each rank with token dispatch) + on its own slice of the batch, traces DDP on the whole batch step by step. """ config = self.model_tester.get_config() if getattr(config, "base_model_ep_plan", None) is None: self.skipTest("Model does not have an expert parallel plan (base_model_ep_plan)") self._run_fsdp2_distributed_test( - "fsdp2_expert_parallel_2d_vs_ddp", _test_fsdp2_expert_parallel_2d_vs_ddp_impl, world_size=4 + "fsdp2_expert_parallel_2d_vs_ddp", + _test_fsdp2_expert_parallel_2d_vs_ddp_impl, + world_size=4, + dispatch=dispatch, ) diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 9b7f2e53387e..ea4b8f816027 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -389,11 +389,13 @@ def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_t dist.barrier() -def _load_ep_and_reference_models(model_path, model_class): +def _load_ep_and_reference_models(model_path, model_class, dispatch=False): """Load EP model and non-EP reference model for comparison.""" model_ep = model_class.from_pretrained( model_path, - distributed_config=DistributedConfig(tp_size=dist.get_world_size(), enable_expert_parallel=True), + distributed_config=DistributedConfig( + tp_size=dist.get_world_size(), enable_expert_parallel=True, expert_parallel_dispatch=dispatch + ), ) dist.barrier() @@ -404,11 +406,11 @@ def _load_ep_and_reference_models(model_path, model_class): return model_ep, model_ref, device -def _test_ep_forward_impl(_rank, model_path, model_class, atol, rtol): +def _test_ep_forward_impl(_rank, model_path, model_class, atol, rtol, dispatch=False): """Implementation for comparing EP and non-EP model outputs.""" set_seed(0) - model_ep, model_ref, device = _load_ep_and_reference_models(model_path, model_class) + model_ep, model_ref, device = _load_ep_and_reference_models(model_path, model_class, dispatch=dispatch) model_ep.eval() model_ref.eval() @@ -428,11 +430,11 @@ def _test_ep_forward_impl(_rank, model_path, model_class, atol, rtol): dist.barrier() -def _test_ep_backward_impl(_rank, model_path, model_class, atol, rtol): +def _test_ep_backward_impl(_rank, model_path, model_class, atol, rtol, dispatch=False): """Implementation for comparing EP and non-EP model backward passes.""" set_seed(0) - model_ep, model_ref, device = _load_ep_and_reference_models(model_path, model_class) + model_ep, model_ref, device = _load_ep_and_reference_models(model_path, model_class, dispatch=dispatch) model_ep.train() model_ref.train() @@ -639,9 +641,9 @@ def test_tp_generation_quantized(self): tmp_dir, model_class, max_new_tokens ) - @parameterized.expand([(False,), (True,)]) + @parameterized.expand([(False, False), (True, False), (False, True)]) @is_tensor_parallel_test - def test_ep_forward(self, tie_word_embeddings): + def test_ep_forward(self, tie_word_embeddings, dispatch): self._skip_if_not_supported(expert_parallel=True) config = self._get_tp_config(tie_word_embeddings=tie_word_embeddings) @@ -654,10 +656,13 @@ def test_ep_forward(self, tie_word_embeddings): model = model_class(config) model.save_pretrained(tmp_dir, save_original_format=True) - _init_distributed(tp=self.tensor_parallel_size)(_test_ep_forward_impl)(tmp_dir, model_class, atol, rtol) + _init_distributed(tp=self.tensor_parallel_size)(_test_ep_forward_impl)( + tmp_dir, model_class, atol, rtol, dispatch=dispatch + ) + @parameterized.expand([(False,), (True,)]) @is_tensor_parallel_test - def test_ep_backward(self): + def test_ep_backward(self, dispatch): self._skip_if_not_supported(expert_parallel=True) config = self._get_tp_config() @@ -670,4 +675,6 @@ def test_ep_backward(self): model = model_class(config) model.save_pretrained(tmp_dir, save_original_format=True) - _init_distributed(tp=self.tensor_parallel_size)(_test_ep_backward_impl)(tmp_dir, model_class, atol, rtol) + _init_distributed(tp=self.tensor_parallel_size)(_test_ep_backward_impl)( + tmp_dir, model_class, atol, rtol, dispatch=dispatch + ) From 7063c456210add33aa2203f18970b56e81d3d6f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 00:51:12 +0000 Subject: [PATCH 15/52] Make the dispatch handling of the experts explicit in apply_fully_sharded_data_parallelism --- src/transformers/distributed/fsdp.py | 11 ++++++----- src/transformers/distributed/mixin.py | 4 +++- src/transformers/distributed/tensor_parallel.py | 5 +++-- 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 458adfec1662..1445964f2a71 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -188,6 +188,7 @@ def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None) def apply_fully_sharded_data_parallelism( model: nn.Module, fsdp_mesh: torch.distributed.device_mesh.DeviceMesh, + expert_parallel_dispatch: bool = False, expert_mesh: torch.distributed.device_mesh.DeviceMesh | None = None, ) -> nn.Module: """ @@ -196,8 +197,9 @@ def apply_fully_sharded_data_parallelism( Torch availability, distributed initialization and the version requirement are asserted upstream by `initialize_fully_sharded_data_parallelism`. - With expert-parallel token dispatch the experts are already sharded across the expert-parallel ranks, which - `fsdp_mesh` spans: pass the mesh they should be sharded across as `expert_mesh` (`None` leaves them as they are). + With expert-parallel token dispatch (`expert_parallel_dispatch=True`) `fsdp_mesh` spans the expert-parallel + ranks, which the experts are already sharded across: they are sharded across `expert_mesh` in their own group + instead, or left as they are when there is no such mesh. """ fsdp_plan = dict(getattr(model, "_fsdp_plan", None) or {}) if not fsdp_plan: @@ -212,9 +214,8 @@ def apply_fully_sharded_data_parallelism( adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, model) reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) - if expert_mesh is not None or fsdp_mesh.ndim == 1 and any(isinstance(p, DTensor) for p in model.parameters()): - # The DTensor parameters are the expert-parallel experts: shard them across `expert_mesh` in their own - # group, or keep them out of the groups below when there is no such mesh. + if expert_parallel_dispatch: + # The DTensor parameters are the expert-parallel experts. expert_modules = [ module for module in model.modules() diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 32241239ead1..af0b3ad4c150 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -216,7 +216,9 @@ def maybe_distribute_model( # their gradient reduction. The experts stay sharded across `tp` and, if any, across `fsdp`. trunk_mesh = device_mesh._flatten("fsdp_tp") if device_mesh.ndim > 1 else device_mesh expert_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else None - model = apply_fully_sharded_data_parallelism(model, trunk_mesh, expert_mesh=expert_mesh) + model = apply_fully_sharded_data_parallelism( + model, trunk_mesh, expert_parallel_dispatch=True, expert_mesh=expert_mesh + ) elif 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) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index bb65490e9a3d..8326b8fab978 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -616,10 +616,11 @@ def install_forward(self, module, mesh, *, is_expert_parallel=False, expert_para else Replicate() ) + if expert_parallel_dispatch: + from ..integrations.moe import dispatch_experts_forward + def tp_forward(*args, **kwargs): if expert_parallel_dispatch: - from ..integrations.moe import dispatch_experts_forward - # Every rank routes its own tokens to the experts' owners; nothing is replicated across # the group, so there is neither an input gradient sum nor an output all-reduce. hidden_states, top_k_index, top_k_weights = args[:3] From dc1fdeee8d28eb0fb476f134befbfcdbe278792a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 00:51:43 +0000 Subject: [PATCH 16/52] Combine the dispatched expert outputs in the routing weights' precision, as the masked path does --- src/transformers/integrations/moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 5374401a6717..4725e182c9ce 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -477,7 +477,7 @@ def dispatch_experts_forward( recv_out = _AllToAll.apply(expert_out, send_sizes, recv_sizes, ep_group) combined = recv_out.new_zeros(num_tokens * num_top_k, hidden_dim) combined[order] = recv_out - combined = combined * top_k_weights.reshape(-1, 1).to(combined.dtype) + combined = combined * top_k_weights.reshape(-1, 1) return combined.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype) From f2a3c26d8afc5a4f2dc12614c4eca617464735f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 00:57:43 +0000 Subject: [PATCH 17/52] Split the batches across every rank with a DistributedSampler under token dispatch accelerate only shards and seeds the sampler when it sees more than one data-parallel rank, so with tp_size alone every rank drew its own random full batch. --- src/transformers/trainer.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index b2f06aee8ce1..63b0ec80aad3 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -50,7 +50,7 @@ from huggingface_hub import CommitInfo, ModelCard from packaging import version from torch import nn -from torch.utils.data import DataLoader, Dataset, IterableDataset, RandomSampler, SequentialSampler +from torch.utils.data import DataLoader, Dataset, DistributedSampler, IterableDataset, RandomSampler, SequentialSampler from . import __version__ from .configuration_utils import PreTrainedConfig @@ -1050,10 +1050,13 @@ def _get_dataloader( if prepared_bs is not None and getattr(prepared_bs, "batch_sampler", None) is sampler: prepared_bs.num_processes = 1 prepared_bs.process_index = 0 - elif self._expert_parallel_dispatch: - # accelerate hands every `tp` rank the same batch; under token dispatch each rank gets its own. - dataloader.batch_sampler.num_processes = self.args.world_size - dataloader.batch_sampler.process_index = self.args.process_index + elif isinstance(sampler, DistributedSampler): + # The sampler already splits the dataset across every rank (expert-parallel token dispatch); the + # `BatchSamplerShard` wrapper accelerate adds when it sees data-parallel ranks must not split again. + prepared_bs = dataloader.batch_sampler + if getattr(prepared_bs, "batch_sampler", None) is not None: + prepared_bs.num_processes = 1 + prepared_bs.process_index = 0 # Store the prepared dataloader for subsequent evaluations if using persistent workers. if dataloader_key is not None and self.args.dataloader_persistent_workers: @@ -1138,6 +1141,11 @@ def _get_train_sampler(self, train_dataset: Dataset | None = None) -> torch.util elif self.args.train_sampling_strategy == "sequential": return SequentialSampler(train_dataset) else: + if self._expert_parallel_dispatch: + # accelerate hands every `tp` rank the same batch; under token dispatch each rank gets its own. + return DistributedSampler( + train_dataset, num_replicas=self.args.world_size, rank=self.args.process_index, seed=self.args.seed + ) return RandomSampler(train_dataset) def _get_eval_sampler(self, eval_dataset: Dataset) -> torch.utils.data.Sampler | None: From 3ab94ff1e90f90471f1574c91999cd8e3b10b957 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 01:09:13 +0000 Subject: [PATCH 18/52] Simplify the dispatch: torch's differentiable all-to-all, the grouped forward for the local experts, ids from the count exchange Also reject expert parallel plans that shard anything but the experts, since the ranks no longer see the same batch. --- .../distributed/tensor_parallel.py | 17 ++- src/transformers/integrations/moe.py | 101 ++++++------------ src/transformers/trainer.py | 17 +-- 3 files changed, 53 insertions(+), 82 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 8326b8fab978..f6ad79f83308 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -619,6 +619,9 @@ def install_forward(self, module, mesh, *, is_expert_parallel=False, expert_para if expert_parallel_dispatch: from ..integrations.moe import dispatch_experts_forward + ep_mesh = mesh if mesh.ndim == 1 else mesh["tp"] + ep_group, ep_size = ep_mesh.get_group(), ep_mesh.size() + def tp_forward(*args, **kwargs): if expert_parallel_dispatch: # Every rank routes its own tokens to the experts' owners; nothing is replicated across @@ -626,9 +629,10 @@ def tp_forward(*args, **kwargs): hidden_states, top_k_index, top_k_weights = args[:3] 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) + return dispatch_experts_forward( + module, hidden_states, top_k_index, top_k_weights, ep_group, ep_size + ) args, kwargs = self.transform_inputs_pre_forward( module, args, kwargs, mesh, is_expert_parallel=is_expert_parallel ) @@ -820,6 +824,15 @@ def _validate_tp_plan_styles(tp_plan: dict[str, str] | None) -> None: def apply_tensor_parallelism(model, tp_mesh): """DTensor backend: shard params as placeholders and install TP forward hooks.""" + if model.config.distributed_config.expert_parallel_dispatch: + # Every rank trains on its own part of the batch, which tensor parallelism cannot do: only the experts can + # be sharded across the group. + other_styles = set(model.tp_plan.values()) - {"ep_router", "grouped_gemm", "moe_tp_experts"} + if other_styles: + raise ValueError( + "`expert_parallel_dispatch=True` needs an expert parallel plan that only shards the experts, but " + f"this model's plan also uses {sorted(other_styles)}." + ) _validate_tp_plan_styles(model.tp_plan) diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 4725e182c9ce..012acc8a3f6d 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -374,35 +374,6 @@ def _grouped_linear( return out -class _AllToAll(torch.autograd.Function): - """`all_to_all_single` with uneven splits whose backward is the reverse exchange.""" - - @staticmethod - def forward(ctx, tensor, output_split_sizes, input_split_sizes, group): - ctx.group, ctx.output_split_sizes, ctx.input_split_sizes = group, output_split_sizes, input_split_sizes - output = tensor.new_empty(sum(output_split_sizes), *tensor.shape[1:]) - torch.distributed.all_to_all_single( - output, - tensor.contiguous(), - output_split_sizes=output_split_sizes, - input_split_sizes=input_split_sizes, - group=group, - ) - return output - - @staticmethod - def backward(ctx, grad_output): - grad_input = grad_output.new_empty(sum(ctx.input_split_sizes), *grad_output.shape[1:]) - torch.distributed.all_to_all_single( - grad_input, - grad_output.contiguous(), - output_split_sizes=ctx.input_split_sizes, - input_split_sizes=ctx.output_split_sizes, - group=ctx.group, - ) - return grad_input, None, None, None - - class _ScaleGrad(torch.autograd.Function): """Identity whose backward scales the gradient.""" @@ -421,63 +392,57 @@ def dispatch_experts_forward( hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor, - ep_mesh, + ep_group, + ep_size: int, ) -> torch.Tensor: """ Expert-parallel forward by token dispatch. Every rank routes its own tokens, sends each selected (token, expert) pair to the rank that owns the expert with an all-to-all, runs its local experts on what it receives, sends the results back and combines them with the routing weights. Each rank trains on its own part of the batch, so the parameters outside the experts are data-parallel across the group and averaged by FSDP2; the all-to-all backward - already sums every rank's contribution to the expert gradients, which are scaled by `1 / ep_size` to match. + sums every rank's contribution to the expert gradients, which are scaled by `1 / ep_size` to match. """ - ep_group, ep_size = ep_mesh.get_group(), ep_mesh.size() + from torch.distributed.nn.functional import all_to_all_single + num_tokens, hidden_dim = hidden_states.shape num_top_k = top_k_index.size(-1) num_local_experts = self.num_experts # the sharding leaves the module with its local expert count - # Group the selected pairs by the rank that owns their expert. + # Sorting the selected pairs by expert groups them by owner rank, since each rank owns a contiguous range of + # experts, and the per-expert counts tell every receiver which expert each token it gets is for. expert_ids = top_k_index.reshape(-1) - owner = expert_ids // num_local_experts - order = torch.argsort(owner, stable=True) - send_tokens = hidden_states.repeat_interleave(num_top_k, dim=0)[order] - send_expert_ids = (expert_ids % num_local_experts)[order] - send_counts = torch.bincount(owner, minlength=ep_size) + order = torch.argsort(expert_ids, stable=True) + send_tokens = hidden_states[order // num_top_k] + send_counts = torch.bincount(expert_ids, minlength=num_local_experts * ep_size).view(ep_size, num_local_experts) recv_counts = torch.empty_like(send_counts) torch.distributed.all_to_all_single(recv_counts, send_counts, group=ep_group) - send_sizes, recv_sizes = send_counts.tolist(), recv_counts.tolist() - - # Exchange the tokens (with gradient) and their local expert ids. - recv_tokens = _AllToAll.apply(send_tokens, recv_sizes, send_sizes, ep_group) - recv_expert_ids = send_expert_ids.new_empty(sum(recv_sizes)) - torch.distributed.all_to_all_single( - recv_expert_ids, send_expert_ids, output_split_sizes=recv_sizes, input_split_sizes=send_sizes, group=ep_group + send_sizes, recv_sizes = send_counts.sum(dim=1).tolist(), recv_counts.sum(dim=1).tolist() + recv_tokens = all_to_all_single( + send_tokens.new_empty(sum(recv_sizes), hidden_dim), + send_tokens, + output_split_sizes=recv_sizes, + input_split_sizes=send_sizes, + group=ep_group, ) + recv_expert_ids = torch.arange(num_local_experts, device=hidden_states.device).repeat(ep_size) + recv_expert_ids = recv_expert_ids.repeat_interleave(recv_counts.reshape(-1)) - # Local experts: sort by expert, grouped GEMMs, unsort. - expert_ids_g, perm = torch.sort(recv_expert_ids, stable=True) - tokens_g = recv_tokens[perm] - offsets = torch.cumsum(torch.bincount(expert_ids_g, minlength=num_local_experts), dim=0, dtype=torch.int32) - grad_scale = 1.0 / ep_size - if self.has_gate: - weight = _ScaleGrad.apply(self.gate_up_proj, grad_scale) - bias = _ScaleGrad.apply(self.gate_up_proj_bias, grad_scale)[expert_ids_g] if self.has_bias else None - else: - weight = _ScaleGrad.apply(self.up_proj, grad_scale) - bias = _ScaleGrad.apply(self.up_proj_bias, grad_scale)[expert_ids_g] if self.has_bias else None - proj_out = _grouped_linear(tokens_g, weight, offsets, bias=bias, is_transposed=self.is_transposed) - proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) - weight = _ScaleGrad.apply(self.down_proj, grad_scale) - bias = _ScaleGrad.apply(self.down_proj_bias, grad_scale)[expert_ids_g] if self.has_bias else None - proj_out = _grouped_linear(proj_out, weight, offsets, bias=bias, is_transposed=self.is_transposed) - inverse_perm = torch.empty_like(perm) - inverse_perm[perm] = torch.arange(perm.numel(), device=perm.device) - expert_out = proj_out[inverse_perm] + # The local experts, as a top-1 routing with unit weights. Scaling the gradient of the output by `1 / ep_size` + # and of the input by `ep_size` leaves the token gradients untouched and scales the expert gradients. + recv_tokens = _ScaleGrad.apply(recv_tokens, ep_size) + unit_weights = torch.ones_like(recv_expert_ids, dtype=recv_tokens.dtype).unsqueeze(-1) + expert_out = grouped_mm_experts_forward(self, recv_tokens, recv_expert_ids.unsqueeze(-1), unit_weights) + expert_out = _ScaleGrad.apply(expert_out, 1.0 / ep_size) # Send the results back to the owners of the tokens and combine them with the routing weights. - recv_out = _AllToAll.apply(expert_out, send_sizes, recv_sizes, ep_group) - combined = recv_out.new_zeros(num_tokens * num_top_k, hidden_dim) - combined[order] = recv_out - combined = combined * top_k_weights.reshape(-1, 1) + recv_out = all_to_all_single( + expert_out.new_empty(send_tokens.size(0), hidden_dim), + expert_out, + output_split_sizes=send_sizes, + input_split_sizes=recv_sizes, + group=ep_group, + ) + combined = recv_out[torch.argsort(order)] * top_k_weights.reshape(-1, 1) return combined.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 63b0ec80aad3..a0a1e5ed756b 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -1042,19 +1042,12 @@ def _get_dataloader( dataloader = self.accelerator.prepare(DataLoader(dataset, **dataloader_params)) - # `BatchRebalanceSampler` is already rank-aware, so the `BatchSamplerShard` wrapper - # added by `accelerator.prepare` would re-shard it and silently drop samples. Neutralise - # it by making the wrapper a passthrough (num_processes=1) - if isinstance(sampler, BatchRebalanceSampler): + # `BatchRebalanceSampler` and the `DistributedSampler` of expert-parallel token dispatch are already + # rank-aware, so the `BatchSamplerShard` wrapper added by `accelerator.prepare` would re-shard them and + # silently drop samples. Neutralise it by making the wrapper a passthrough (num_processes=1) + if isinstance(sampler, (BatchRebalanceSampler, DistributedSampler)): prepared_bs = getattr(dataloader, "batch_sampler", None) - if prepared_bs is not None and getattr(prepared_bs, "batch_sampler", None) is sampler: - prepared_bs.num_processes = 1 - prepared_bs.process_index = 0 - elif isinstance(sampler, DistributedSampler): - # The sampler already splits the dataset across every rank (expert-parallel token dispatch); the - # `BatchSamplerShard` wrapper accelerate adds when it sees data-parallel ranks must not split again. - prepared_bs = dataloader.batch_sampler - if getattr(prepared_bs, "batch_sampler", None) is not None: + if prepared_bs is not None and getattr(prepared_bs, "batch_sampler", None) is not None: prepared_bs.num_processes = 1 prepared_bs.process_index = 0 From d1881444164e5a8b8c113585b81e19d0882a7c57 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 01:11:02 +0000 Subject: [PATCH 19/52] Dispatch as plan-level styles; one notion of the ranks sharing a batch in the Trainer ep_dispatch_router and ep_dispatch_experts are substituted into the expert parallel plan when expert_parallel_dispatch is set, so no style carries a flag it ignores and nothing is stored on the modules. get_tp_size() is 1 under dispatch, which the loss scale, the token count and the total batch size all go through; non-random sampling strategies are rejected explicitly. --- src/transformers/distributed/fsdp.py | 11 ++- src/transformers/distributed/mixin.py | 4 +- .../distributed/tensor_parallel.py | 69 ++++++++++--------- src/transformers/trainer.py | 27 +++++--- 4 files changed, 61 insertions(+), 50 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 1445964f2a71..95477df8905b 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -188,7 +188,6 @@ def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None) def apply_fully_sharded_data_parallelism( model: nn.Module, fsdp_mesh: torch.distributed.device_mesh.DeviceMesh, - expert_parallel_dispatch: bool = False, expert_mesh: torch.distributed.device_mesh.DeviceMesh | None = None, ) -> nn.Module: """ @@ -197,9 +196,9 @@ def apply_fully_sharded_data_parallelism( Torch availability, distributed initialization and the version requirement are asserted upstream by `initialize_fully_sharded_data_parallelism`. - With expert-parallel token dispatch (`expert_parallel_dispatch=True`) `fsdp_mesh` spans the expert-parallel - ranks, which the experts are already sharded across: they are sharded across `expert_mesh` in their own group - instead, or left as they are when there is no such mesh. + With expert-parallel token dispatch `fsdp_mesh` spans the expert-parallel ranks, which the experts are already + sharded across: they are sharded across `expert_mesh` in their own group instead, or left as they are when there + is no such mesh. """ fsdp_plan = dict(getattr(model, "_fsdp_plan", None) or {}) if not fsdp_plan: @@ -214,8 +213,8 @@ def apply_fully_sharded_data_parallelism( adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, model) reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) - if expert_parallel_dispatch: - # The DTensor parameters are the expert-parallel experts. + if distributed_config is not None and distributed_config.expert_parallel_dispatch: + # The DTensor parameters are the expert-parallel experts: the plan is restricted to them under dispatch. expert_modules = [ module for module in model.modules() diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index af0b3ad4c150..32241239ead1 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -216,9 +216,7 @@ def maybe_distribute_model( # their gradient reduction. The experts stay sharded across `tp` and, if any, across `fsdp`. trunk_mesh = device_mesh._flatten("fsdp_tp") if device_mesh.ndim > 1 else device_mesh expert_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else None - model = apply_fully_sharded_data_parallelism( - model, trunk_mesh, expert_parallel_dispatch=True, expert_mesh=expert_mesh - ) + model = apply_fully_sharded_data_parallelism(model, trunk_mesh, expert_mesh=expert_mesh) elif 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) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index f6ad79f83308..ff0619ad2d4f 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -144,7 +144,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, *, expert_parallel_dispatch=False): + def install_forward(self, module, mesh): """Install pre / around / post transforms by replacing module.forward.""" original_forward = module.forward @@ -344,7 +344,7 @@ class ReplicatedWithGradAllReduce(TensorParallelLayer): summed across the mesh. """ - def install_forward(self, module, mesh, *, expert_parallel_dispatch=False): + def install_forward(self, module, mesh): # 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 +415,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, *, expert_parallel_dispatch=False): + def install_forward(self, module, mesh): # 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( @@ -604,7 +604,7 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_ return (hidden_states, *routing_args), kwargs - def install_forward(self, module, mesh, *, is_expert_parallel=False, expert_parallel_dispatch=False): + def install_forward(self, module, mesh, *, is_expert_parallel=False): """Install the transforms but pass `is_expert_parallel` in the forward call.""" original_forward = module.forward output_source = ( @@ -616,23 +616,7 @@ def install_forward(self, module, mesh, *, is_expert_parallel=False, expert_para else Replicate() ) - if expert_parallel_dispatch: - from ..integrations.moe import dispatch_experts_forward - - ep_mesh = mesh if mesh.ndim == 1 else mesh["tp"] - ep_group, ep_size = ep_mesh.get_group(), ep_mesh.size() - def tp_forward(*args, **kwargs): - if expert_parallel_dispatch: - # Every rank routes its own tokens to the experts' owners; nothing is replicated across - # the group, so there is neither an input gradient sum nor an output all-reduce. - hidden_states, top_k_index, top_k_weights = args[:3] - if isinstance(hidden_states, DTensor): - hidden_states = hidden_states.to_local() - with self.context_around_forward(module, mesh): - return dispatch_experts_forward( - module, hidden_states, top_k_index, top_k_weights, ep_group, ep_size - ) args, kwargs = self.transform_inputs_pre_forward( module, args, kwargs, mesh, is_expert_parallel=is_expert_parallel ) @@ -711,14 +695,7 @@ class EpRouterParallel(TensorParallelLayer): Scores and indices stay paired element-wise in `(seq, top_k)` shape. """ - def install_forward(self, module, mesh, *, expert_parallel_dispatch=False): - module.expert_parallel_dispatch = expert_parallel_dispatch - return super().install_forward(module, mesh) - def transform_output_post_forward(self, module, output, mesh): - if module.expert_parallel_dispatch: - # The experts forward sends every token to its experts' owners, so it needs the global ids. - return output ep_rank, ep_size = mesh.get_local_rank(), mesh.size() num_experts = getattr(module, "num_experts", None) if num_experts is None: @@ -760,6 +737,32 @@ def transform_output_post_forward(self, module, output, mesh): return output +class EpDispatchRouterParallel(TensorParallelLayer): + """Router of expert-parallel token dispatch: the experts forward needs the global expert ids and scores.""" + + +class EpDispatchExpertsParallel(MoeExpertsParallel): + """ + Experts of expert-parallel token dispatch: every rank routes its own tokens to the experts' owners, so nothing is + replicated across the group and the forward has neither an input gradient sum nor an output all-reduce. + """ + + def install_forward(self, module, mesh, *, is_expert_parallel=False): + from ..integrations.moe import dispatch_experts_forward + + ep_mesh = mesh if mesh.ndim == 1 else mesh["tp"] + ep_group, ep_size = ep_mesh.get_group(), ep_mesh.size() + + def tp_forward(hidden_states, top_k_index, top_k_weights, *args, **kwargs): + if isinstance(hidden_states, DTensor): + hidden_states = hidden_states.to_local() + with self.context_around_forward(module, mesh): + return dispatch_experts_forward(module, hidden_states, top_k_index, top_k_weights, ep_group, ep_size) + + module.forward = tp_forward + return module + + class MoeTensorParalellMegaMoeExperts(MoeExpertsParallel): """TP layer for DeepGEMM Mega MoE experts. @@ -797,6 +800,8 @@ class ParallelInterface(GeneralInterface): "sequence_parallel": SequenceParallel(use_local_output=True), "grouped_gemm": MoEParamShard(Shard(0), shards_expert_dim=True), "ep_router": EpRouterParallel(), + "ep_dispatch_router": EpDispatchRouterParallel(), + "ep_dispatch_experts": EpDispatchExpertsParallel(), "megamoe_router": RouterParallelMegaMoe(), "moe_tp_experts": MoeExpertsParallel(), "megamoe_experts": MoeTensorParalellMegaMoeExperts(), @@ -826,13 +831,17 @@ def apply_tensor_parallelism(model, tp_mesh): """DTensor backend: shard params as placeholders and install TP forward hooks.""" if model.config.distributed_config.expert_parallel_dispatch: # Every rank trains on its own part of the batch, which tensor parallelism cannot do: only the experts can - # be sharded across the group. + # be sharded across the group, and the router and experts styles become their dispatch versions. other_styles = set(model.tp_plan.values()) - {"ep_router", "grouped_gemm", "moe_tp_experts"} if other_styles: raise ValueError( "`expert_parallel_dispatch=True` needs an expert parallel plan that only shards the experts, but " f"this model's plan also uses {sorted(other_styles)}." ) + model.tp_plan = { + name: {"ep_router": "ep_dispatch_router", "moe_tp_experts": "ep_dispatch_experts"}.get(style, style) + for name, style in model.tp_plan.items() + } _validate_tp_plan_styles(model.tp_plan) @@ -853,9 +862,7 @@ 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, expert_parallel_dispatch=model.config.distributed_config.expert_parallel_dispatch - ) + ALL_PARALLEL_STYLES[style_name].install_forward(module, tp_mesh) module._is_hooked = True return model diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index a0a1e5ed756b..413567459456 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -621,6 +621,11 @@ def __init__( # treats the `tp` ranks as one data-parallel rank: the batches and the token counts are handled here. distributed_config = getattr(getattr(model, "config", None), "distributed_config", None) self._expert_parallel_dispatch = distributed_config is not None and distributed_config.expert_parallel_dispatch + if self._expert_parallel_dispatch and args.train_sampling_strategy != "random": + raise ValueError( + "`expert_parallel_dispatch=True` splits the batches across every rank with a `DistributedSampler`, " + f"which `train_sampling_strategy='{args.train_sampling_strategy}'` does not go through." + ) if ( getattr(model, "_device_mesh", None) is not None and args.save_strategy != SaveStrategy.NO @@ -1137,7 +1142,11 @@ def _get_train_sampler(self, train_dataset: Dataset | None = None) -> torch.util if self._expert_parallel_dispatch: # accelerate hands every `tp` rank the same batch; under token dispatch each rank gets its own. return DistributedSampler( - train_dataset, num_replicas=self.args.world_size, rank=self.args.process_index, seed=self.args.seed + train_dataset, + num_replicas=self.args.world_size, + rank=self.args.process_index, + seed=self.args.seed, + drop_last=self.args.dataloader_drop_last, ) return RandomSampler(train_dataset) @@ -2166,10 +2175,7 @@ def compute_loss( ): # TP and EP-as-TP ranks see replicated batches; `num_processes` over-counts # them by `tp_size`. Mirror the divisor used in `_get_num_items_in_batch`. - loss_scale = self.accelerator.num_processes - pc = getattr(self.accelerator, "parallelism_config", None) - if pc is not None and not self._expert_parallel_dispatch: - loss_scale //= pc.tp_size + loss_scale = self.accelerator.num_processes // self.get_tp_size() loss *= loss_scale if self.args.n_gpu <= 1 else self.args.n_gpu return (loss, outputs) if return_outputs else loss @@ -2318,10 +2324,9 @@ def _get_num_items_in_batch(self, batch_samples: list, device: torch.device) -> # In the DataParallel case, convert the scalar tensor into a 2-dim tensor with the same value repeated num_items_in_batch = num_items_in_batch.unsqueeze(0).expand(self.args.n_gpu, -1) # Divide by number of devices with the same batch - if ( - pc := getattr(self.accelerator, "parallelism_config", None) - ) and not self._expert_parallel_dispatch: - num_items_in_batch = num_items_in_batch // pc.non_data_parallel_size + num_items_in_batch = num_items_in_batch // ( + self.get_tp_size() * self.get_cp_size() * self.get_sp_size() + ) return num_items_in_batch @@ -2558,7 +2563,9 @@ def get_cp_size(self) -> int: def get_tp_size(self) -> int: """Get the tensor parallel size from either the model or DeepSpeed config.""" - # 1. Check model.tp_size first + # 1. Check model.tp_size first; with expert-parallel token dispatch the `tp` ranks train on their own batches + if self._expert_parallel_dispatch: + return 1 if (model_tp := getattr(self.model, "_tp_size", None)) is not None: return model_tp From b9bf58bdce75b45dcecbb3f9231a3ef75c972895 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 01:12:59 +0000 Subject: [PATCH 20/52] Select the dispatch styles in the mixin, guard the torch version and dispatch_batches --- docs/source/en/expert_parallelism.md | 2 +- src/transformers/distributed/fsdp.py | 2 ++ src/transformers/distributed/mixin.py | 11 +++++++++++ src/transformers/distributed/tensor_parallel.py | 13 ------------- src/transformers/trainer.py | 5 +++++ 5 files changed, 19 insertions(+), 14 deletions(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index d62332ae4c08..ec8d4d680dd2 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -62,7 +62,7 @@ distributed_config = DistributedConfig( ) ``` -Every rank trains on its own part of the batch: at each MoE layer it routes its tokens, sends the selected (token, expert) pairs to their owners with an all-to-all, runs its local experts, and gets the results back with a second all-to-all. Only the routed tokens travel. The parameters outside the experts are data-parallel across the expert-parallel group, so they are sharded with [FSDP2](./fsdp) across every rank (the `fsdp` and `tp` dimensions together when both are set) and FSDP2 reduces their gradients; the experts stay sharded across `tp` and, with `fsdp_size > 1`, across `fsdp`. The [`Trainer`] gives each rank its own batches and counts tokens across all of them. +Every rank trains on its own part of the batch: at each MoE layer it routes its tokens, sends the selected (token, expert) pairs to their owners with an all-to-all, runs its local experts, and gets the results back with a second all-to-all. Only the routed tokens travel. The parameters outside the experts are data-parallel across the expert-parallel group, so they are sharded with [FSDP2](./fsdp) across every rank (the `fsdp` and `tp` dimensions together when both are set) and FSDP2 reduces their gradients; the experts stay sharded across `tp` and, with `fsdp_size > 1`, across `fsdp`. The [`Trainer`] gives each rank its own batches and counts tokens across all of them. With `fsdp_size=1` the experts stay outside FSDP2, so `fsdp_mixed_precision` and `fsdp_cpu_offload` do not apply to them. ## Combining with FSDP2 diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 95477df8905b..fcba3076db10 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -214,6 +214,8 @@ def apply_fully_sharded_data_parallelism( reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) if distributed_config is not None and distributed_config.expert_parallel_dispatch: + if not is_torch_greater_or_equal("2.7"): + raise OSError("Expert-parallel token dispatch requires `torch>=2.7`.") # The DTensor parameters are the expert-parallel experts: the plan is restricted to them under dispatch. expert_modules = [ module diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 32241239ead1..1bc73f44b846 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -208,6 +208,17 @@ def maybe_distribute_model( 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 + if distributed_config.expert_parallel_dispatch: + # Every rank trains on its own part of the batch, which tensor parallelism cannot do: only the + # experts can be sharded across the group, through the dispatch versions of their styles. + other_styles = set(model.tp_plan.values()) - {"ep_router", "grouped_gemm", "moe_tp_experts"} + if other_styles: + raise ValueError( + "`expert_parallel_dispatch=True` needs an expert parallel plan that only shards the " + f"experts, but this model's plan also uses {sorted(other_styles)}." + ) + dispatch_styles = {"ep_router": "ep_dispatch_router", "moe_tp_experts": "ep_dispatch_experts"} + model.tp_plan = {name: dispatch_styles.get(style, style) for name, style in model.tp_plan.items()} model = apply_tensor_parallelism(model, tp_mesh) if distributed_config.expert_parallel_dispatch: diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index ff0619ad2d4f..7044b26bc2b8 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -829,19 +829,6 @@ def _validate_tp_plan_styles(tp_plan: dict[str, str] | None) -> None: def apply_tensor_parallelism(model, tp_mesh): """DTensor backend: shard params as placeholders and install TP forward hooks.""" - if model.config.distributed_config.expert_parallel_dispatch: - # Every rank trains on its own part of the batch, which tensor parallelism cannot do: only the experts can - # be sharded across the group, and the router and experts styles become their dispatch versions. - other_styles = set(model.tp_plan.values()) - {"ep_router", "grouped_gemm", "moe_tp_experts"} - if other_styles: - raise ValueError( - "`expert_parallel_dispatch=True` needs an expert parallel plan that only shards the experts, but " - f"this model's plan also uses {sorted(other_styles)}." - ) - model.tp_plan = { - name: {"ep_router": "ep_dispatch_router", "moe_tp_experts": "ep_dispatch_experts"}.get(style, style) - for name, style in model.tp_plan.items() - } _validate_tp_plan_styles(model.tp_plan) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 413567459456..71e628e9edac 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -1052,6 +1052,11 @@ def _get_dataloader( # silently drop samples. Neutralise it by making the wrapper a passthrough (num_processes=1) if isinstance(sampler, (BatchRebalanceSampler, DistributedSampler)): prepared_bs = getattr(dataloader, "batch_sampler", None) + if prepared_bs is None and isinstance(sampler, DistributedSampler): + raise ValueError( + "`expert_parallel_dispatch=True` splits the batches across every rank with the batch sampler, " + "which `dispatch_batches=True` replaces." + ) if prepared_bs is not None and getattr(prepared_bs, "batch_sampler", None) is not None: prepared_bs.num_processes = 1 prepared_bs.process_index = 0 From 39cdce508c99452296623ec9e039b741b94602d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 01:15:35 +0000 Subject: [PATCH 21/52] Rewrite the expert parallel plan the model actually reads for the dispatch styles --- src/transformers/distributed/mixin.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 1bc73f44b846..f4b24ae32c73 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -218,7 +218,8 @@ def maybe_distribute_model( f"experts, but this model's plan also uses {sorted(other_styles)}." ) dispatch_styles = {"ep_router": "ep_dispatch_router", "moe_tp_experts": "ep_dispatch_experts"} - model.tp_plan = {name: dispatch_styles.get(style, style) for name, style in model.tp_plan.items()} + # `tp_plan` reads `_ep_plan` under expert parallelism, so that is the plan to rewrite. + model._ep_plan = {name: dispatch_styles.get(style, style) for name, style in model.tp_plan.items()} model = apply_tensor_parallelism(model, tp_mesh) if distributed_config.expert_parallel_dispatch: From 9f86230812c4835591954046cd965e7fa827b8bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 01:19:17 +0000 Subject: [PATCH 22/52] Seed the dispatch sampler with data_seed; require the grouped_mm experts implementation --- docs/source/en/expert_parallelism.md | 2 +- src/transformers/distributed/mixin.py | 5 +++++ src/transformers/trainer.py | 2 +- 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index ec8d4d680dd2..c397f5a0d48b 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -62,7 +62,7 @@ distributed_config = DistributedConfig( ) ``` -Every rank trains on its own part of the batch: at each MoE layer it routes its tokens, sends the selected (token, expert) pairs to their owners with an all-to-all, runs its local experts, and gets the results back with a second all-to-all. Only the routed tokens travel. The parameters outside the experts are data-parallel across the expert-parallel group, so they are sharded with [FSDP2](./fsdp) across every rank (the `fsdp` and `tp` dimensions together when both are set) and FSDP2 reduces their gradients; the experts stay sharded across `tp` and, with `fsdp_size > 1`, across `fsdp`. The [`Trainer`] gives each rank its own batches and counts tokens across all of them. With `fsdp_size=1` the experts stay outside FSDP2, so `fsdp_mixed_precision` and `fsdp_cpu_offload` do not apply to them. +Every rank trains on its own part of the batch: at each MoE layer it routes its tokens, sends the selected (token, expert) pairs to their owners with an all-to-all, runs its local experts, and gets the results back with a second all-to-all. Only the routed tokens travel. The parameters outside the experts are data-parallel across the expert-parallel group, so they are sharded with [FSDP2](./fsdp) across every rank (the `fsdp` and `tp` dimensions together when both are set) and FSDP2 reduces their gradients; the experts stay sharded across `tp` and, with `fsdp_size > 1`, across `fsdp`. The [`Trainer`] gives each rank its own training batches and counts tokens across all of them; evaluation runs as under plain expert parallelism, with the same batches on every `tp` rank. With `fsdp_size=1` the experts stay outside FSDP2, so `fsdp_mixed_precision` and `fsdp_cpu_offload` do not apply to them. ## Combining with FSDP2 diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index f4b24ae32c73..dd589c3d9765 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -217,6 +217,11 @@ def maybe_distribute_model( "`expert_parallel_dispatch=True` needs an expert parallel plan that only shards the " f"experts, but this model's plan also uses {sorted(other_styles)}." ) + if model.config._experts_implementation != "grouped_mm": + raise ValueError( + "`expert_parallel_dispatch=True` runs the local experts with the `grouped_mm` experts " + f"implementation, not `{model.config._experts_implementation}`." + ) dispatch_styles = {"ep_router": "ep_dispatch_router", "moe_tp_experts": "ep_dispatch_experts"} # `tp_plan` reads `_ep_plan` under expert parallelism, so that is the plan to rewrite. model._ep_plan = {name: dispatch_styles.get(style, style) for name, style in model.tp_plan.items()} diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 71e628e9edac..6f83626a06fb 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -1150,7 +1150,7 @@ def _get_train_sampler(self, train_dataset: Dataset | None = None) -> torch.util train_dataset, num_replicas=self.args.world_size, rank=self.args.process_index, - seed=self.args.seed, + seed=self.args.data_seed if self.args.data_seed is not None else self.args.seed, drop_last=self.args.dataloader_drop_last, ) return RandomSampler(train_dataset) From 930c146a0a2b783b67e521a1704d839caf437942 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 01:26:00 +0000 Subject: [PATCH 23/52] Run the local experts directly: the grouped forward's masking and combine passes cost 15-25% of throughput --- src/transformers/integrations/moe.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 012acc8a3f6d..fb2e559c9a85 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -427,12 +427,21 @@ def dispatch_experts_forward( recv_expert_ids = torch.arange(num_local_experts, device=hidden_states.device).repeat(ep_size) recv_expert_ids = recv_expert_ids.repeat_interleave(recv_counts.reshape(-1)) - # The local experts, as a top-1 routing with unit weights. Scaling the gradient of the output by `1 / ep_size` + # The local experts: sort by expert, grouped GEMMs, unsort. Scaling the gradient of the output by `1 / ep_size` # and of the input by `ep_size` leaves the token gradients untouched and scales the expert gradients. recv_tokens = _ScaleGrad.apply(recv_tokens, ep_size) - unit_weights = torch.ones_like(recv_expert_ids, dtype=recv_tokens.dtype).unsqueeze(-1) - expert_out = grouped_mm_experts_forward(self, recv_tokens, recv_expert_ids.unsqueeze(-1), unit_weights) - expert_out = _ScaleGrad.apply(expert_out, 1.0 / ep_size) + expert_ids_g, perm = torch.sort(recv_expert_ids, stable=True) + tokens_g = recv_tokens[perm] + offsets = torch.cumsum(torch.bincount(expert_ids_g, minlength=num_local_experts), dim=0, dtype=torch.int32) + if self.has_gate: + weight, bias = self.gate_up_proj, self.gate_up_proj_bias[expert_ids_g] if self.has_bias else None + else: + weight, bias = self.up_proj, self.up_proj_bias[expert_ids_g] if self.has_bias else None + proj_out = _grouped_linear(tokens_g, weight, offsets, bias=bias, is_transposed=self.is_transposed) + proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) + bias = self.down_proj_bias[expert_ids_g] if self.has_bias else None + proj_out = _grouped_linear(proj_out, self.down_proj, offsets, bias=bias, is_transposed=self.is_transposed) + expert_out = _ScaleGrad.apply(proj_out[torch.argsort(perm)], 1.0 / ep_size) # Send the results back to the owners of the tokens and combine them with the routing weights. recv_out = all_to_all_single( From 15588e00d9518a55b7a9eb61f96fbb2722e49c5b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 01:32:02 +0000 Subject: [PATCH 24/52] Run the local experts through the module's own forward The dispatch kernel no longer assumes the grouped-GEMM parameter layout: the experts module's forward is called as a top-1 routing with unit weights, so any experts implementation (and any model-specific step inside it) works under dispatch. --- .../distributed/tensor_parallel.py | 6 +++- src/transformers/integrations/moe.py | 30 +++++++------------ 2 files changed, 16 insertions(+), 20 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 7044b26bc2b8..1ddc4b194d52 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -750,6 +750,7 @@ class EpDispatchExpertsParallel(MoeExpertsParallel): def install_forward(self, module, mesh, *, is_expert_parallel=False): from ..integrations.moe import dispatch_experts_forward + original_forward = module.forward ep_mesh = mesh if mesh.ndim == 1 else mesh["tp"] ep_group, ep_size = ep_mesh.get_group(), ep_mesh.size() @@ -757,7 +758,10 @@ def tp_forward(hidden_states, top_k_index, top_k_weights, *args, **kwargs): if isinstance(hidden_states, DTensor): hidden_states = hidden_states.to_local() with self.context_around_forward(module, mesh): - return dispatch_experts_forward(module, hidden_states, top_k_index, top_k_weights, ep_group, ep_size) + # The sharding leaves the module with its local expert count. + return dispatch_experts_forward( + original_forward, module.num_experts, hidden_states, top_k_index, top_k_weights, ep_group, ep_size + ) module.forward = tp_forward return module diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index fb2e559c9a85..6e80de23ad8a 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -388,7 +388,8 @@ def backward(ctx, grad_output): def dispatch_experts_forward( - self: torch.nn.Module, + experts_forward: Callable, + num_local_experts: int, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor, @@ -397,16 +398,16 @@ def dispatch_experts_forward( ) -> torch.Tensor: """ Expert-parallel forward by token dispatch. Every rank routes its own tokens, sends each selected (token, expert) - pair to the rank that owns the expert with an all-to-all, runs its local experts on what it receives, sends the - results back and combines them with the routing weights. Each rank trains on its own part of the batch, so the - parameters outside the experts are data-parallel across the group and averaged by FSDP2; the all-to-all backward - sums every rank's contribution to the expert gradients, which are scaled by `1 / ep_size` to match. + pair to the rank that owns the expert with an all-to-all, runs its local experts on what it receives with + `experts_forward` (the experts module's own forward, as a top-1 routing with unit weights), sends the results + back and combines them with the routing weights. Each rank trains on its own part of the batch, so the parameters + outside the experts are data-parallel across the group and averaged by FSDP2; the all-to-all backward sums every + rank's contribution to the expert gradients, which are scaled by `1 / ep_size` to match. """ from torch.distributed.nn.functional import all_to_all_single num_tokens, hidden_dim = hidden_states.shape num_top_k = top_k_index.size(-1) - num_local_experts = self.num_experts # the sharding leaves the module with its local expert count # Sorting the selected pairs by expert groups them by owner rank, since each rank owns a contiguous range of # experts, and the per-expert counts tell every receiver which expert each token it gets is for. @@ -427,21 +428,12 @@ def dispatch_experts_forward( recv_expert_ids = torch.arange(num_local_experts, device=hidden_states.device).repeat(ep_size) recv_expert_ids = recv_expert_ids.repeat_interleave(recv_counts.reshape(-1)) - # The local experts: sort by expert, grouped GEMMs, unsort. Scaling the gradient of the output by `1 / ep_size` + # The local experts, as a top-1 routing with unit weights. Scaling the gradient of the output by `1 / ep_size` # and of the input by `ep_size` leaves the token gradients untouched and scales the expert gradients. recv_tokens = _ScaleGrad.apply(recv_tokens, ep_size) - expert_ids_g, perm = torch.sort(recv_expert_ids, stable=True) - tokens_g = recv_tokens[perm] - offsets = torch.cumsum(torch.bincount(expert_ids_g, minlength=num_local_experts), dim=0, dtype=torch.int32) - if self.has_gate: - weight, bias = self.gate_up_proj, self.gate_up_proj_bias[expert_ids_g] if self.has_bias else None - else: - weight, bias = self.up_proj, self.up_proj_bias[expert_ids_g] if self.has_bias else None - proj_out = _grouped_linear(tokens_g, weight, offsets, bias=bias, is_transposed=self.is_transposed) - proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) - bias = self.down_proj_bias[expert_ids_g] if self.has_bias else None - proj_out = _grouped_linear(proj_out, self.down_proj, offsets, bias=bias, is_transposed=self.is_transposed) - expert_out = _ScaleGrad.apply(proj_out[torch.argsort(perm)], 1.0 / ep_size) + unit_weights = torch.ones_like(recv_expert_ids, dtype=recv_tokens.dtype).unsqueeze(-1) + expert_out = experts_forward(recv_tokens, recv_expert_ids.unsqueeze(-1), unit_weights) + expert_out = _ScaleGrad.apply(expert_out, 1.0 / ep_size) # Send the results back to the owners of the tokens and combine them with the routing weights. recv_out = all_to_all_single( From 51f281d9f7b34375165aa89c1d18cdc7cbebc9da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 01:53:50 +0000 Subject: [PATCH 25/52] Dispatch: one host sync per layer, comments say what the code does --- .../distributed/configuration_utils.py | 6 +++--- src/transformers/distributed/fsdp.py | 6 +++--- src/transformers/distributed/mixin.py | 4 ++-- src/transformers/distributed/tensor_parallel.py | 5 +---- src/transformers/integrations/moe.py | 16 ++++++++++------ 5 files changed, 19 insertions(+), 18 deletions(-) diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index b81611a171c6..88bda47b4434 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -35,9 +35,9 @@ class DistributedConfig: enable_expert_parallel (`bool`, *optional*, defaults to `False`): Route MoE models through the expert-parallel path (``base_model_ep_plan``). expert_parallel_dispatch (`bool`, *optional*, defaults to `False`): - Send each token to the rank that owns its experts with an all-to-all instead of running the whole batch - on every expert-parallel rank. Each rank then trains on its own part of the batch, and the parameters - that are not expert-parallel are sharded with FSDP2 across every rank. Requires `enable_expert_parallel`. + Send each token to the rank that owns its experts with an all-to-all. Each rank trains on its own part of + the batch, and the parameters that are not expert-parallel are sharded with FSDP2 across every rank. + Requires `enable_expert_parallel`. fsdp_size (`int`, *optional*): Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1. fsdp_cpu_offload (`bool`, *optional*, defaults to `False`): diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index fcba3076db10..38162eb74aec 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -197,8 +197,8 @@ def apply_fully_sharded_data_parallelism( are asserted upstream by `initialize_fully_sharded_data_parallelism`. With expert-parallel token dispatch `fsdp_mesh` spans the expert-parallel ranks, which the experts are already - sharded across: they are sharded across `expert_mesh` in their own group instead, or left as they are when there - is no such mesh. + sharded across: the experts are fully sharded across `expert_mesh` in their own group, and passed to FSDP2 as + `ignored_params` when `expert_mesh` is `None`. """ fsdp_plan = dict(getattr(model, "_fsdp_plan", None) or {}) if not fsdp_plan: @@ -216,7 +216,7 @@ def apply_fully_sharded_data_parallelism( if distributed_config is not None and distributed_config.expert_parallel_dispatch: if not is_torch_greater_or_equal("2.7"): raise OSError("Expert-parallel token dispatch requires `torch>=2.7`.") - # The DTensor parameters are the expert-parallel experts: the plan is restricted to them under dispatch. + # The DTensor parameters are the expert-parallel experts (the expert parallel plan shards only them). expert_modules = [ module for module in model.modules() diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index dd589c3d9765..d14aa57aff9a 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -209,8 +209,8 @@ def maybe_distribute_model( if isinstance(distributed_config.tp_plan, dict): model.tp_plan = distributed_config.tp_plan if distributed_config.expert_parallel_dispatch: - # Every rank trains on its own part of the batch, which tensor parallelism cannot do: only the - # experts can be sharded across the group, through the dispatch versions of their styles. + # Every rank trains on its own part of the batch, so only the experts can be sharded across the + # group, through the dispatch versions of their styles. other_styles = set(model.tp_plan.values()) - {"ep_router", "grouped_gemm", "moe_tp_experts"} if other_styles: raise ValueError( diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 1ddc4b194d52..a6bfd2e2226d 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -742,10 +742,7 @@ class EpDispatchRouterParallel(TensorParallelLayer): class EpDispatchExpertsParallel(MoeExpertsParallel): - """ - Experts of expert-parallel token dispatch: every rank routes its own tokens to the experts' owners, so nothing is - replicated across the group and the forward has neither an input gradient sum nor an output all-reduce. - """ + """Experts of expert-parallel token dispatch: every rank sends its own tokens to the experts' owners.""" def install_forward(self, module, mesh, *, is_expert_parallel=False): from ..integrations.moe import dispatch_experts_forward diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 6e80de23ad8a..e1236b0a5238 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -410,14 +410,16 @@ def dispatch_experts_forward( num_top_k = top_k_index.size(-1) # Sorting the selected pairs by expert groups them by owner rank, since each rank owns a contiguous range of - # experts, and the per-expert counts tell every receiver which expert each token it gets is for. + # experts, and the per-expert counts tell every receiver which expert each token it gets is for. The split + # sizes are the one host sync of the layer. expert_ids = top_k_index.reshape(-1) - order = torch.argsort(expert_ids, stable=True) + order = torch.argsort(expert_ids) send_tokens = hidden_states[order // num_top_k] - send_counts = torch.bincount(expert_ids, minlength=num_local_experts * ep_size).view(ep_size, num_local_experts) + send_counts = torch.zeros(num_local_experts * ep_size, dtype=torch.long, device=hidden_states.device) + send_counts = send_counts.scatter_add_(0, expert_ids, torch.ones_like(expert_ids)).view(ep_size, num_local_experts) recv_counts = torch.empty_like(send_counts) torch.distributed.all_to_all_single(recv_counts, send_counts, group=ep_group) - send_sizes, recv_sizes = send_counts.sum(dim=1).tolist(), recv_counts.sum(dim=1).tolist() + send_sizes, recv_sizes = torch.stack([send_counts.sum(dim=1), recv_counts.sum(dim=1)]).tolist() recv_tokens = all_to_all_single( send_tokens.new_empty(sum(recv_sizes), hidden_dim), send_tokens, @@ -426,7 +428,7 @@ def dispatch_experts_forward( group=ep_group, ) recv_expert_ids = torch.arange(num_local_experts, device=hidden_states.device).repeat(ep_size) - recv_expert_ids = recv_expert_ids.repeat_interleave(recv_counts.reshape(-1)) + recv_expert_ids = recv_expert_ids.repeat_interleave(recv_counts.reshape(-1), output_size=sum(recv_sizes)) # The local experts, as a top-1 routing with unit weights. Scaling the gradient of the output by `1 / ep_size` # and of the input by `ep_size` leaves the token gradients untouched and scales the expert gradients. @@ -443,7 +445,9 @@ def dispatch_experts_forward( input_split_sizes=recv_sizes, group=ep_group, ) - combined = recv_out[torch.argsort(order)] * top_k_weights.reshape(-1, 1) + inverse_order = torch.empty_like(order) + inverse_order[order] = torch.arange(order.numel(), device=order.device) + combined = recv_out[inverse_order] * top_k_weights.reshape(-1, 1) return combined.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype) From 101bfea9fd3b41e6e79662406b200e56cb8c571f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 02:05:52 +0000 Subject: [PATCH 26/52] Dispatch: drop the router style, reuse the named mesh and is_dtensor, check the batches at init --- src/transformers/distributed/fsdp.py | 7 ++----- src/transformers/distributed/mixin.py | 15 ++++++++++----- src/transformers/distributed/tensor_parallel.py | 12 +++--------- src/transformers/integrations/moe.py | 4 ++-- src/transformers/trainer.py | 15 ++++++++------- 5 files changed, 25 insertions(+), 28 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 38162eb74aec..c216da458d10 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -20,7 +20,7 @@ from ..utils import is_torch_available, is_torch_distributed_available, is_torch_greater_or_equal, logging, strtobool from ..utils.quantization_config import QuantizationMethod from .tensor_parallel import replace_layer_number_by_wildcard -from .utils import _is_torch_distributed_initialized +from .utils import _is_torch_distributed_initialized, is_dtensor if TYPE_CHECKING: @@ -34,7 +34,6 @@ if is_torch_distributed_available() and is_torch_greater_or_equal("2.6"): from torch.distributed._composable.fsdp import fully_shard from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy - from torch.distributed.tensor import DTensor logger = logging.get_logger(__name__) @@ -218,9 +217,7 @@ def apply_fully_sharded_data_parallelism( raise OSError("Expert-parallel token dispatch requires `torch>=2.7`.") # The DTensor parameters are the expert-parallel experts (the expert parallel plan shards only them). expert_modules = [ - module - for module in model.modules() - if any(isinstance(p, DTensor) for p in module.parameters(recurse=False)) + module for module in model.modules() if any(is_dtensor(p) for p in module.parameters(recurse=False)) ] if expert_mesh is not None: for module in expert_modules: diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index d14aa57aff9a..2eea82bb6538 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -57,6 +57,7 @@ class DistributedMixin: _ep_plan: dict[str, str] | None = None _tp_size = None _fsdp_size = None + _expert_parallel_dispatch = False _pp_plan: dict[str, tuple[str, str]] | None = None _fsdp_plan: dict[str, str] | None = None @@ -201,6 +202,7 @@ def maybe_distribute_model( # 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 + model._expert_parallel_dispatch = distributed_config.expert_parallel_dispatch # Both may apply: the tensor/expert parallel plan shards across `tp` first, then FSDP2 # shards every parameter (the `tp`-sharded ones included) across `fsdp`. @@ -210,7 +212,7 @@ def maybe_distribute_model( model.tp_plan = distributed_config.tp_plan if distributed_config.expert_parallel_dispatch: # Every rank trains on its own part of the batch, so only the experts can be sharded across the - # group, through the dispatch versions of their styles. + # group: the experts get the dispatch style, and the router keeps its global ids and scores. other_styles = set(model.tp_plan.values()) - {"ep_router", "grouped_gemm", "moe_tp_experts"} if other_styles: raise ValueError( @@ -222,16 +224,19 @@ def maybe_distribute_model( "`expert_parallel_dispatch=True` runs the local experts with the `grouped_mm` experts " f"implementation, not `{model.config._experts_implementation}`." ) - dispatch_styles = {"ep_router": "ep_dispatch_router", "moe_tp_experts": "ep_dispatch_experts"} # `tp_plan` reads `_ep_plan` under expert parallelism, so that is the plan to rewrite. - model._ep_plan = {name: dispatch_styles.get(style, style) for name, style in model.tp_plan.items()} + model._ep_plan = { + name: "ep_dispatch_experts" if style == "moe_tp_experts" else style + for name, style in model.tp_plan.items() + if style != "ep_router" + } model = apply_tensor_parallelism(model, tp_mesh) if distributed_config.expert_parallel_dispatch: # Every expert-parallel rank trains on its own part of the batch, so the parameters outside the # experts are data-parallel across the whole mesh: FSDP2 shards them across all of it and owns # their gradient reduction. The experts stay sharded across `tp` and, if any, across `fsdp`. - trunk_mesh = device_mesh._flatten("fsdp_tp") if device_mesh.ndim > 1 else device_mesh + trunk_mesh = device_mesh["fsdp_tp"] if device_mesh.ndim > 1 else device_mesh expert_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else None model = apply_fully_sharded_data_parallelism(model, trunk_mesh, expert_mesh=expert_mesh) elif distributed_config.fsdp_size > 1: @@ -301,7 +306,7 @@ def gather_sharded_state_dict_for_save( if distributed_config is None: return state_dict - if is_fsdp_managed_module(model_to_save): + if distributed_config.fsdp_size > 1 or distributed_config.expert_parallel_dispatch: # Also covers the 2-D (fsdp, tp) mesh and token dispatch: every parameter is FSDP-managed, and # the full state dict is only materialized on rank 0. if not _is_torch_distributed_initialized(): diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index a6bfd2e2226d..c15cc772f9e9 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -737,21 +737,16 @@ def transform_output_post_forward(self, module, output, mesh): return output -class EpDispatchRouterParallel(TensorParallelLayer): - """Router of expert-parallel token dispatch: the experts forward needs the global expert ids and scores.""" - - class EpDispatchExpertsParallel(MoeExpertsParallel): """Experts of expert-parallel token dispatch: every rank sends its own tokens to the experts' owners.""" - def install_forward(self, module, mesh, *, is_expert_parallel=False): + def install_forward(self, module, mesh): from ..integrations.moe import dispatch_experts_forward original_forward = module.forward - ep_mesh = mesh if mesh.ndim == 1 else mesh["tp"] - ep_group, ep_size = ep_mesh.get_group(), ep_mesh.size() + ep_group, ep_size = mesh.get_group(), mesh.size() - def tp_forward(hidden_states, top_k_index, top_k_weights, *args, **kwargs): + def tp_forward(hidden_states, top_k_index, top_k_weights): if isinstance(hidden_states, DTensor): hidden_states = hidden_states.to_local() with self.context_around_forward(module, mesh): @@ -801,7 +796,6 @@ class ParallelInterface(GeneralInterface): "sequence_parallel": SequenceParallel(use_local_output=True), "grouped_gemm": MoEParamShard(Shard(0), shards_expert_dim=True), "ep_router": EpRouterParallel(), - "ep_dispatch_router": EpDispatchRouterParallel(), "ep_dispatch_experts": EpDispatchExpertsParallel(), "megamoe_router": RouterParallelMegaMoe(), "moe_tp_experts": MoeExpertsParallel(), diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index e1236b0a5238..9946bec6ba5a 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -401,8 +401,8 @@ def dispatch_experts_forward( pair to the rank that owns the expert with an all-to-all, runs its local experts on what it receives with `experts_forward` (the experts module's own forward, as a top-1 routing with unit weights), sends the results back and combines them with the routing weights. Each rank trains on its own part of the batch, so the parameters - outside the experts are data-parallel across the group and averaged by FSDP2; the all-to-all backward sums every - rank's contribution to the expert gradients, which are scaled by `1 / ep_size` to match. + outside the experts are data-parallel across the group and averaged by FSDP2; the local experts run on every + rank's tokens, so their gradients already sum every rank's contribution and are scaled by `1 / ep_size` to match. """ from torch.distributed.nn.functional import all_to_all_single diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 6f83626a06fb..72a5c2d172ed 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -619,13 +619,19 @@ def __init__( self._mixed_mesh_grads: bool | None = None # With expert-parallel token dispatch every rank trains on its own part of the batch, while accelerate # treats the `tp` ranks as one data-parallel rank: the batches and the token counts are handled here. - distributed_config = getattr(getattr(model, "config", None), "distributed_config", None) - self._expert_parallel_dispatch = distributed_config is not None and distributed_config.expert_parallel_dispatch + self._expert_parallel_dispatch = getattr(model, "_expert_parallel_dispatch", False) if self._expert_parallel_dispatch and args.train_sampling_strategy != "random": raise ValueError( "`expert_parallel_dispatch=True` splits the batches across every rank with a `DistributedSampler`, " f"which `train_sampling_strategy='{args.train_sampling_strategy}'` does not go through." ) + if self._expert_parallel_dispatch and ( + self.accelerator.dispatch_batches or (train_dataset is not None and not has_length(train_dataset)) + ): + raise ValueError( + "`expert_parallel_dispatch=True` splits the batches across every rank with a `DistributedSampler`, " + "which needs a sized training dataset and `dispatch_batches=False`." + ) if ( getattr(model, "_device_mesh", None) is not None and args.save_strategy != SaveStrategy.NO @@ -1052,11 +1058,6 @@ def _get_dataloader( # silently drop samples. Neutralise it by making the wrapper a passthrough (num_processes=1) if isinstance(sampler, (BatchRebalanceSampler, DistributedSampler)): prepared_bs = getattr(dataloader, "batch_sampler", None) - if prepared_bs is None and isinstance(sampler, DistributedSampler): - raise ValueError( - "`expert_parallel_dispatch=True` splits the batches across every rank with the batch sampler, " - "which `dispatch_batches=True` replaces." - ) if prepared_bs is not None and getattr(prepared_bs, "batch_sampler", None) is not None: prepared_bs.num_processes = 1 prepared_bs.process_index = 0 From 2d9046bd7a48b150ace6200731f30b685929ae3c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 02:08:16 +0000 Subject: [PATCH 27/52] Dispatch: skip the mixin tests for models without an FSDP plan, drop the experts implementation check --- src/transformers/distributed/mixin.py | 5 ----- tests/test_tensor_parallel_mixin.py | 9 +++++++++ 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 2eea82bb6538..0186f3096385 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -219,11 +219,6 @@ def maybe_distribute_model( "`expert_parallel_dispatch=True` needs an expert parallel plan that only shards the " f"experts, but this model's plan also uses {sorted(other_styles)}." ) - if model.config._experts_implementation != "grouped_mm": - raise ValueError( - "`expert_parallel_dispatch=True` runs the local experts with the `grouped_mm` experts " - f"implementation, not `{model.config._experts_implementation}`." - ) # `tp_plan` reads `_ep_plan` under expert parallelism, so that is the plan to rewrite. model._ep_plan = { name: "ep_dispatch_experts" if style == "moe_tp_experts" else style diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index ea4b8f816027..571097dd17da 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -489,6 +489,11 @@ def _has_ep_plan(self) -> bool: config = self.model_tester.get_config() return hasattr(config, "base_model_ep_plan") and config.base_model_ep_plan is not None + def _has_fsdp_plan(self) -> bool: + """Check if model has an expert parallel plan defined.""" + config = self.model_tester.get_config() + return hasattr(config, "base_model_fsdp_plan") and config.base_model_fsdp_plan is not None + def _has_tp_plan(self) -> bool: """Check if model has a tensor parallel plan defined.""" config = self.model_tester.get_config() @@ -645,6 +650,8 @@ def test_tp_generation_quantized(self): @is_tensor_parallel_test def test_ep_forward(self, tie_word_embeddings, dispatch): self._skip_if_not_supported(expert_parallel=True) + if dispatch and not self._has_fsdp_plan(): + self.skipTest("Model does not have an FSDP plan (base_model_fsdp_plan)") config = self._get_tp_config(tie_word_embeddings=tie_word_embeddings) model_class = self._get_tp_model_class() @@ -664,6 +671,8 @@ def test_ep_forward(self, tie_word_embeddings, dispatch): @is_tensor_parallel_test def test_ep_backward(self, dispatch): self._skip_if_not_supported(expert_parallel=True) + if dispatch and not self._has_fsdp_plan(): + self.skipTest("Model does not have an FSDP plan (base_model_fsdp_plan)") config = self._get_tp_config() model_class = self._get_tp_model_class() From 42606afbf864c69f4dfe015d5c8b64e062b69356 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 02:08:30 +0000 Subject: [PATCH 28/52] Fix the _has_fsdp_plan docstring --- tests/test_tensor_parallel_mixin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 571097dd17da..ffc485e75002 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -490,7 +490,7 @@ def _has_ep_plan(self) -> bool: return hasattr(config, "base_model_ep_plan") and config.base_model_ep_plan is not None def _has_fsdp_plan(self) -> bool: - """Check if model has an expert parallel plan defined.""" + """Check if model has an FSDP plan defined.""" config = self.model_tester.get_config() return hasattr(config, "base_model_fsdp_plan") and config.base_model_fsdp_plan is not None From d8eb1fc1a3abff3b8496fc389306f3c8d7b5e596 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 02:23:38 +0000 Subject: [PATCH 29/52] Drop the FSDP plan skip from the mixin tests: every config carries a default plan --- tests/test_tensor_parallel_mixin.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index ffc485e75002..bdb978eca42f 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -489,10 +489,6 @@ def _has_ep_plan(self) -> bool: config = self.model_tester.get_config() return hasattr(config, "base_model_ep_plan") and config.base_model_ep_plan is not None - def _has_fsdp_plan(self) -> bool: - """Check if model has an FSDP plan defined.""" - config = self.model_tester.get_config() - return hasattr(config, "base_model_fsdp_plan") and config.base_model_fsdp_plan is not None def _has_tp_plan(self) -> bool: """Check if model has a tensor parallel plan defined.""" @@ -650,8 +646,6 @@ def test_tp_generation_quantized(self): @is_tensor_parallel_test def test_ep_forward(self, tie_word_embeddings, dispatch): self._skip_if_not_supported(expert_parallel=True) - if dispatch and not self._has_fsdp_plan(): - self.skipTest("Model does not have an FSDP plan (base_model_fsdp_plan)") config = self._get_tp_config(tie_word_embeddings=tie_word_embeddings) model_class = self._get_tp_model_class() @@ -671,8 +665,6 @@ def test_ep_forward(self, tie_word_embeddings, dispatch): @is_tensor_parallel_test def test_ep_backward(self, dispatch): self._skip_if_not_supported(expert_parallel=True) - if dispatch and not self._has_fsdp_plan(): - self.skipTest("Model does not have an FSDP plan (base_model_fsdp_plan)") config = self._get_tp_config() model_class = self._get_tp_model_class() From 23920e3ad2dc11aa48a81fcd0a1a242f8436c99a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 02:23:52 +0000 Subject: [PATCH 30/52] Format --- tests/test_tensor_parallel_mixin.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index bdb978eca42f..ea4b8f816027 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -489,7 +489,6 @@ def _has_ep_plan(self) -> bool: config = self.model_tester.get_config() return hasattr(config, "base_model_ep_plan") and config.base_model_ep_plan is not None - def _has_tp_plan(self) -> bool: """Check if model has a tensor parallel plan defined.""" config = self.model_tester.get_config() From e14986967558d09b1594df70403f7d401ac90256 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 02:32:50 +0000 Subject: [PATCH 31/52] Keep the base install_forward signature on the dispatch experts style --- src/transformers/distributed/tensor_parallel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index c15cc772f9e9..53f148b5cd63 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -740,7 +740,7 @@ def transform_output_post_forward(self, module, output, mesh): class EpDispatchExpertsParallel(MoeExpertsParallel): """Experts of expert-parallel token dispatch: every rank sends its own tokens to the experts' owners.""" - def install_forward(self, module, mesh): + def install_forward(self, module, mesh, *, is_expert_parallel=False): from ..integrations.moe import dispatch_experts_forward original_forward = module.forward From e139e8a5fd6a3d0199b327f4be3a6fb9c2f55436 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 4 Sep 2026 03:03:46 +0000 Subject: [PATCH 32/52] Docs: shorter token dispatch overview --- docs/source/en/expert_parallelism.md | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index c397f5a0d48b..ddaeb60c3ec6 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -62,7 +62,13 @@ distributed_config = DistributedConfig( ) ``` -Every rank trains on its own part of the batch: at each MoE layer it routes its tokens, sends the selected (token, expert) pairs to their owners with an all-to-all, runs its local experts, and gets the results back with a second all-to-all. Only the routed tokens travel. The parameters outside the experts are data-parallel across the expert-parallel group, so they are sharded with [FSDP2](./fsdp) across every rank (the `fsdp` and `tp` dimensions together when both are set) and FSDP2 reduces their gradients; the experts stay sharded across `tp` and, with `fsdp_size > 1`, across `fsdp`. The [`Trainer`] gives each rank its own training batches and counts tokens across all of them; evaluation runs as under plain expert parallelism, with the same batches on every `tp` rank. With `fsdp_size=1` the experts stay outside FSDP2, so `fsdp_mixed_precision` and `fsdp_cpu_offload` do not apply to them. +Each rank trains on its own part of the batch. At every MoE layer it routes its tokens, sends each (token, expert) pair to the rank that owns the expert with an all-to-all, runs its local experts, and gets the results back with a second all-to-all. Only the routed tokens travel. + +The rest of the model is handled as follows: + +- The parameters outside the experts are data-parallel across the whole group, so they are sharded with [FSDP2](./fsdp) across every rank (`fsdp` and `tp` together when both are set), and FSDP2 reduces their gradients. +- The experts stay sharded across `tp`, and across `fsdp` too when `fsdp_size > 1`. With `fsdp_size=1` they are outside FSDP2, so `fsdp_mixed_precision` and `fsdp_cpu_offload` do not apply to them. +- The [`Trainer`] gives each rank its own training batches and counts tokens across all of them. Evaluation is unchanged from plain expert parallelism: every `tp` rank sees the same batches. ## Combining with FSDP2 From cf74929d3a2fbe7e951f7ce563c7c0d475f69cf1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Mon, 7 Sep 2026 21:37:36 +0000 Subject: [PATCH 33/52] Replicate the rest of the plan under expert parallel token dispatch `expert_parallel_dispatch=True` raised when the expert parallel plan sharded anything but the experts. DeepseekV4's plan does: it shards the lightning indexer colwise and all-reduces the scorer, so `test_ep_forward_2` and `test_ep_backward_1` died in `from_pretrained`. Drop those entries with a warning instead. Under dispatch every rank trains on its own part of the batch, so the trunk is replicated and FSDP2-sharded across the whole mesh anyway, which is what dropping them gives. --- src/transformers/distributed/fsdp.py | 3 ++- src/transformers/distributed/mixin.py | 19 ++++++++++++------- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index c216da458d10..905a35e52959 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -215,7 +215,8 @@ def apply_fully_sharded_data_parallelism( if distributed_config is not None and distributed_config.expert_parallel_dispatch: if not is_torch_greater_or_equal("2.7"): raise OSError("Expert-parallel token dispatch requires `torch>=2.7`.") - # The DTensor parameters are the expert-parallel experts (the expert parallel plan shards only them). + # The DTensor parameters are the expert-parallel experts: `maybe_distribute_model` rewrote the expert + # parallel plan to shard only them. expert_modules = [ module for module in model.modules() if any(is_dtensor(p) for p in module.parameters(recurse=False)) ] diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 0186f3096385..20f43ee3e832 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -212,18 +212,23 @@ def maybe_distribute_model( model.tp_plan = distributed_config.tp_plan if distributed_config.expert_parallel_dispatch: # Every rank trains on its own part of the batch, so only the experts can be sharded across the - # group: the experts get the dispatch style, and the router keeps its global ids and scores. - other_styles = set(model.tp_plan.values()) - {"ep_router", "grouped_gemm", "moe_tp_experts"} - if other_styles: - raise ValueError( - "`expert_parallel_dispatch=True` needs an expert parallel plan that only shards the " - f"experts, but this model's plan also uses {sorted(other_styles)}." + # group: the experts get the dispatch style, the router keeps its global ids and scores, and + # whatever else the plan shards stays replicated, data-parallel like the rest of the trunk. + replicated = sorted( + name + for name, style in model.tp_plan.items() + if style not in ("ep_router", "grouped_gemm", "moe_tp_experts") + ) + if replicated: + logger.warning( + "`expert_parallel_dispatch=True` shards only the experts, so these expert parallel plan " + f"entries are ignored and their modules stay replicated: {replicated}." ) # `tp_plan` reads `_ep_plan` under expert parallelism, so that is the plan to rewrite. model._ep_plan = { name: "ep_dispatch_experts" if style == "moe_tp_experts" else style for name, style in model.tp_plan.items() - if style != "ep_router" + if style in ("grouped_gemm", "moe_tp_experts") } model = apply_tensor_parallelism(model, tp_mesh) From c9a4216715357b9e910f8871e12349086c1efa11 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:02:50 -0400 Subject: [PATCH 34/52] Update docs/source/en/expert_parallelism.md Co-authored-by: VI-Arthur --- docs/source/en/expert_parallelism.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index ddaeb60c3ec6..881caa3d85c4 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -52,7 +52,7 @@ torchrun --nproc-per-node 8 your_script.py ## Token dispatch -By default every expert-parallel rank runs the whole batch, keeps only the experts it owns, and the group all-reduces the expert outputs at the end of every MoE layer. The parameters outside the experts do `tp_size` times the same work, and the all-reduce moves the full activations. `expert_parallel_dispatch=True` sends each token to the rank that owns its experts instead. +By default every expert-parallel rank runs the whole batch, keeps only the experts it owns, and the group all-reduces the expert outputs at the end of every MoE layer. The parameters outside the experts do `tp_size` times the same work, and the all-reduce moves the full activations. `expert_parallel_dispatch=True` sends each token to the rank that owns its experts instead: this means a lot less data needs to travel between GPUs / nodes in large scale trainings. ```py distributed_config = DistributedConfig( From 8f2f9cb1e1edb71b0ab9670d09efbe3d19f06c34 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:48:58 -0400 Subject: [PATCH 35/52] Update docs/source/en/expert_parallelism.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/expert_parallelism.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index 881caa3d85c4..d2434818a069 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -69,7 +69,7 @@ The rest of the model is handled as follows: - The parameters outside the experts are data-parallel across the whole group, so they are sharded with [FSDP2](./fsdp) across every rank (`fsdp` and `tp` together when both are set), and FSDP2 reduces their gradients. - The experts stay sharded across `tp`, and across `fsdp` too when `fsdp_size > 1`. With `fsdp_size=1` they are outside FSDP2, so `fsdp_mixed_precision` and `fsdp_cpu_offload` do not apply to them. - The [`Trainer`] gives each rank its own training batches and counts tokens across all of them. Evaluation is unchanged from plain expert parallelism: every `tp` rank sees the same batches. - +- Training needs a sized (map-style_ dataset, `dispatch_batches=False`, and `train_sampling_strategy="random"`. Iterable datasets and other sampling strategies are rejected. ## 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. From 3c8f600e71ac103ba29969824b9b0a863ead88d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:49:16 -0400 Subject: [PATCH 36/52] Update docs/source/en/expert_parallelism.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/expert_parallelism.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index d2434818a069..a6d10e607530 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -55,6 +55,9 @@ torchrun --nproc-per-node 8 your_script.py By default every expert-parallel rank runs the whole batch, keeps only the experts it owns, and the group all-reduces the expert outputs at the end of every MoE layer. The parameters outside the experts do `tp_size` times the same work, and the all-reduce moves the full activations. `expert_parallel_dispatch=True` sends each token to the rank that owns its experts instead: this means a lot less data needs to travel between GPUs / nodes in large scale trainings. ```py +from transformers import AutoModelForCausalLM +from transformers.distributed import DistributedConfig + distributed_config = DistributedConfig( tp_size=8, enable_expert_parallel=True, From e7be30a9c2ea9dae5790df556f4413f08bff5f64 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:49:26 -0400 Subject: [PATCH 37/52] Update docs/source/en/expert_parallelism.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/expert_parallelism.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index a6d10e607530..b590b7c255b2 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -52,7 +52,7 @@ torchrun --nproc-per-node 8 your_script.py ## Token dispatch -By default every expert-parallel rank runs the whole batch, keeps only the experts it owns, and the group all-reduces the expert outputs at the end of every MoE layer. The parameters outside the experts do `tp_size` times the same work, and the all-reduce moves the full activations. `expert_parallel_dispatch=True` sends each token to the rank that owns its experts instead: this means a lot less data needs to travel between GPUs / nodes in large scale trainings. +By default, every expert parallel rank runs the whole batch, keeps only the experts it own, and all-reduces expert outputs after every MoE layer. Dense layers then do `tp_size` times the same work, and the all-reduce moves full activations. Set `expert_parallel_dispatch=True` to send each token to the rank that owns its experts. Each rank trains on its own batch shard, and a lot less data is required to travel between GPUs/nodes during larg-scale training. ```py from transformers import AutoModelForCausalLM From 8dbccad41e8e1501973f67ee4e6d9df902d4224b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:49:35 -0400 Subject: [PATCH 38/52] Update docs/source/en/expert_parallelism.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/expert_parallelism.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index b590b7c255b2..78cd5ae033df 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -67,7 +67,7 @@ distributed_config = DistributedConfig( Each rank trains on its own part of the batch. At every MoE layer it routes its tokens, sends each (token, expert) pair to the rank that owns the expert with an all-to-all, runs its local experts, and gets the results back with a second all-to-all. Only the routed tokens travel. -The rest of the model is handled as follows: +For the rest of the model: - The parameters outside the experts are data-parallel across the whole group, so they are sharded with [FSDP2](./fsdp) across every rank (`fsdp` and `tp` together when both are set), and FSDP2 reduces their gradients. - The experts stay sharded across `tp`, and across `fsdp` too when `fsdp_size > 1`. With `fsdp_size=1` they are outside FSDP2, so `fsdp_mixed_precision` and `fsdp_cpu_offload` do not apply to them. From 0817429f032d76a7b85b3a905692060674fa934b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Tue, 8 Sep 2026 21:40:36 +0000 Subject: [PATCH 39/52] Keep replicated_with_grad_allreduce under expert-parallel dispatch Dropping it left the parameter with no gradient reduction at all: FSDP2 treats the experts module as expert-owned and skips it, and every rank saw different tokens. fp32 certification on a tiny Muse-shaped model: post_expert_norm.weight at rel 0.80 and 0.50 against a 1-GPU reference under dispatch, 0.75 and 0.59 under dispatch2d; masked and 2-D at 2e-6. A ~0.75 error is the local quarter of the sum that never happened. --- src/transformers/distributed/mixin.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 20f43ee3e832..9eed97deada5 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -214,10 +214,11 @@ def maybe_distribute_model( # Every rank trains on its own part of the batch, so only the experts can be sharded across the # group: the experts get the dispatch style, the router keeps its global ids and scores, and # whatever else the plan shards stays replicated, data-parallel like the rest of the trunk. + # Replicated parameters inside the experts module keep their gradient all-reduce: FSDP2 treats + # that module as expert-owned and does not reduce them, and each rank saw different tokens. + kept = ("grouped_gemm", "moe_tp_experts", "replicated_with_grad_allreduce") replicated = sorted( - name - for name, style in model.tp_plan.items() - if style not in ("ep_router", "grouped_gemm", "moe_tp_experts") + name for name, style in model.tp_plan.items() if style not in ("ep_router", *kept) ) if replicated: logger.warning( @@ -228,7 +229,7 @@ def maybe_distribute_model( model._ep_plan = { name: "ep_dispatch_experts" if style == "moe_tp_experts" else style for name, style in model.tp_plan.items() - if style in ("grouped_gemm", "moe_tp_experts") + if style in kept } model = apply_tensor_parallelism(model, tp_mesh) From 34e63c9d8a3dc0a376125938d5e2630c1d6495ba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Wed, 9 Sep 2026 01:00:35 +0000 Subject: [PATCH 40/52] Compare gradients, not only the loss, in the EP backward test A missing or doubled gradient reduction leaves the forward and the loss untouched, so the loss-only check passed on both. Gather each sharded gradient back to the full parameter and compare every parameter's gradient against the single-process reference, and require the same set of parameters to have received one. --- tests/test_tensor_parallel_mixin.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index ea4b8f816027..4a3fda34db8c 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -454,6 +454,27 @@ def _test_ep_backward_impl(_rank, model_path, model_class, atol, rtol, dispatch= f"Diff: {(loss_ref - loss_ep).abs().item()}" ) + # A missing or doubled gradient reduction leaves the forward, and so the loss, untouched: only the parameter + # gradients show it. Sharded gradients are gathered back to the full parameter before comparing. + from torch.distributed.tensor import DTensor + + grads_ref = {name: param.grad for name, param in model_ref.named_parameters() if param.grad is not None} + grads_ep = {name: param.grad for name, param in model_ep.named_parameters() if param.grad is not None} + assert grads_ep.keys() == grads_ref.keys(), ( + f"Parameters with a gradient differ. Only in EP: {sorted(grads_ep.keys() - grads_ref.keys())}, " + f"only in reference: {sorted(grads_ref.keys() - grads_ep.keys())}" + ) + mismatched = [] + for name, grad in grads_ep.items(): + grad = grad.full_tensor() if isinstance(grad, DTensor) else grad + ref = grads_ref[name] + if not torch.allclose(ref, grad.to(ref.device), atol=atol, rtol=rtol): + mismatched.append( + f"{name}: max abs diff {(ref - grad).abs().max().item():.3e}, " + f"ref norm {ref.norm().item():.3e}, EP norm {grad.norm().item():.3e}" + ) + assert not mismatched, "EP and non-EP model gradients differ:\n" + "\n".join(mismatched) + dist.barrier() From 75fff7166656c08b66b264247783c68fc6e5edf1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Wed, 9 Sep 2026 23:31:09 -0400 Subject: [PATCH 41/52] Update docs/source/en/expert_parallelism.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/expert_parallelism.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index c62eef296119..ba2a704c8676 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -65,7 +65,7 @@ model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_c 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: +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 | |---|---|---| From acba6bc97a0e809178ac43aecdae5de9c0062ea8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Wed, 9 Sep 2026 23:31:17 -0400 Subject: [PATCH 42/52] Update docs/source/en/expert_parallelism.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/expert_parallelism.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index ba2a704c8676..170ea6ba9813 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -55,6 +55,9 @@ torchrun --nproc-per-node 8 your_script.py 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 +from transformers import AutoModelForCausalLM +from transformers.distributed import DistributedConfig + distributed_config = DistributedConfig( tp_size=4, # expert parallel size fsdp_size=2, # data parallel shards From 69668c3c29413ce5b01e8d2a2470d8f46c812795 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Wed, 9 Sep 2026 23:31:25 -0400 Subject: [PATCH 43/52] Update docs/source/en/expert_parallelism.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/expert_parallelism.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index 170ea6ba9813..ec673f61f56a 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -66,7 +66,9 @@ distributed_config = DistributedConfig( 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. +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`. From 9cfb2d3b09035e380713b157380c1105e5bcca21 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Wed, 9 Sep 2026 23:31:42 -0400 Subject: [PATCH 44/52] Update docs/source/en/expert_parallelism.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/expert_parallelism.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index ec673f61f56a..307d8f8d9ee1 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -52,7 +52,7 @@ 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. +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 28cb4827ff97f6ebba86a07673d50f7e62bc1646 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 10 Sep 2026 14:11:18 +0000 Subject: [PATCH 45/52] Declare dispatch_experts_forward next to its only caller `dispatch_experts_forward` and the `_ScaleGrad` helper it uses move from `integrations/moe.py` to `distributed/tensor_parallel.py`, above `EpDispatchExpertsParallel`, so the local import inside `install_forward` goes away. Also say in the FSDP2 section that it describes expert parallelism without token dispatch, and what changes when dispatch is on, which read as a contradiction with the section above it. --- docs/source/en/expert_parallelism.md | 9 ++- .../distributed/tensor_parallel.py | 80 ++++++++++++++++++- src/transformers/integrations/moe.py | 77 ------------------ 3 files changed, 83 insertions(+), 83 deletions(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index 0cad407eac47..07d45d546163 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -52,7 +52,7 @@ torchrun --nproc-per-node 8 your_script.py ## Token dispatch -By default, every expert parallel rank runs the whole batch, keeps only the experts it own, and all-reduces expert outputs after every MoE layer. Dense layers then do `tp_size` times the same work, and the all-reduce moves full activations. Set `expert_parallel_dispatch=True` to send each token to the rank that owns its experts. Each rank trains on its own batch shard, and a lot less data is required to travel between GPUs/nodes during larg-scale training. +By default, every expert parallel rank runs the whole batch, keeps only the experts it owns, and all-reduces expert outputs after every MoE layer. Dense layers then do `tp_size` times the same work, and the all-reduce moves full activations. Set `expert_parallel_dispatch=True` to send each token to the rank that owns its experts. Each rank trains on its own batch shard, and a lot less data is required to travel between GPUs/nodes during large-scale training. ```py from transformers import AutoModelForCausalLM @@ -72,10 +72,11 @@ For the rest of the model: - The parameters outside the experts are data-parallel across the whole group, so they are sharded with [FSDP2](./fsdp) across every rank (`fsdp` and `tp` together when both are set), and FSDP2 reduces their gradients. - The experts stay sharded across `tp`, and across `fsdp` too when `fsdp_size > 1`. With `fsdp_size=1` they are outside FSDP2, so `fsdp_mixed_precision` and `fsdp_cpu_offload` do not apply to them. - The [`Trainer`] gives each rank its own training batches and counts tokens across all of them. Evaluation is unchanged from plain expert parallelism: every `tp` rank sees the same batches. -- Training needs a sized (map-style_ dataset, `dispatch_batches=False`, and `train_sampling_strategy="random"`. Iterable datasets and other sampling strategies are rejected. +- Training needs a sized (map-style) dataset, `dispatch_batches=False`, and `train_sampling_strategy="random"`. Iterable datasets and other sampling strategies are rejected. + ## 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). +Without token dispatch, 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 @@ -89,7 +90,7 @@ distributed_config = DistributedConfig( 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. +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. With `expert_parallel_dispatch=True` the non-expert parameters are data-parallel across the whole mesh rather than replicated across `tp`, so FSDP2 shards them across `fsdp` and `tp` together. 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`]. diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 53f148b5cd63..fcc1c563edbf 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -15,6 +15,7 @@ import contextlib import re +from collections.abc import Callable from ..utils import logging from ..utils.generic import GeneralInterface @@ -737,12 +738,87 @@ def transform_output_post_forward(self, module, output, mesh): return output +class _ScaleGrad(torch.autograd.Function): + """Identity whose backward scales the gradient.""" + + @staticmethod + def forward(ctx, tensor, scale): + ctx.scale = scale + return tensor + + @staticmethod + def backward(ctx, grad_output): + return grad_output * ctx.scale, None + + +def dispatch_experts_forward( + experts_forward: Callable, + num_local_experts: int, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, + ep_group, + ep_size: int, +) -> torch.Tensor: + """ + Expert-parallel forward by token dispatch. Every rank routes its own tokens, sends each selected (token, expert) + pair to the rank that owns the expert with an all-to-all, runs its local experts on what it receives with + `experts_forward` (the experts module's own forward, as a top-1 routing with unit weights), sends the results + back and combines them with the routing weights. Each rank trains on its own part of the batch, so the parameters + outside the experts are data-parallel across the group and averaged by FSDP2; the local experts run on every + rank's tokens, so their gradients already sum every rank's contribution and are scaled by `1 / ep_size` to match. + """ + from torch.distributed.nn.functional import all_to_all_single + + num_tokens, hidden_dim = hidden_states.shape + num_top_k = top_k_index.size(-1) + + # Sorting the selected pairs by expert groups them by owner rank, since each rank owns a contiguous range of + # experts, and the per-expert counts tell every receiver which expert each token it gets is for. The split + # sizes are the one host sync of the layer. + expert_ids = top_k_index.reshape(-1) + order = torch.argsort(expert_ids) + send_tokens = hidden_states[order // num_top_k] + send_counts = torch.zeros(num_local_experts * ep_size, dtype=torch.long, device=hidden_states.device) + send_counts = send_counts.scatter_add_(0, expert_ids, torch.ones_like(expert_ids)).view(ep_size, num_local_experts) + recv_counts = torch.empty_like(send_counts) + torch.distributed.all_to_all_single(recv_counts, send_counts, group=ep_group) + send_sizes, recv_sizes = torch.stack([send_counts.sum(dim=1), recv_counts.sum(dim=1)]).tolist() + recv_tokens = all_to_all_single( + send_tokens.new_empty(sum(recv_sizes), hidden_dim), + send_tokens, + output_split_sizes=recv_sizes, + input_split_sizes=send_sizes, + group=ep_group, + ) + recv_expert_ids = torch.arange(num_local_experts, device=hidden_states.device).repeat(ep_size) + recv_expert_ids = recv_expert_ids.repeat_interleave(recv_counts.reshape(-1), output_size=sum(recv_sizes)) + + # The local experts, as a top-1 routing with unit weights. Scaling the gradient of the output by `1 / ep_size` + # and of the input by `ep_size` leaves the token gradients untouched and scales the expert gradients. + recv_tokens = _ScaleGrad.apply(recv_tokens, ep_size) + unit_weights = torch.ones_like(recv_expert_ids, dtype=recv_tokens.dtype).unsqueeze(-1) + expert_out = experts_forward(recv_tokens, recv_expert_ids.unsqueeze(-1), unit_weights) + expert_out = _ScaleGrad.apply(expert_out, 1.0 / ep_size) + + # Send the results back to the owners of the tokens and combine them with the routing weights. + recv_out = all_to_all_single( + expert_out.new_empty(send_tokens.size(0), hidden_dim), + expert_out, + output_split_sizes=send_sizes, + input_split_sizes=recv_sizes, + group=ep_group, + ) + inverse_order = torch.empty_like(order) + inverse_order[order] = torch.arange(order.numel(), device=order.device) + combined = recv_out[inverse_order] * top_k_weights.reshape(-1, 1) + return combined.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype) + + class EpDispatchExpertsParallel(MoeExpertsParallel): """Experts of expert-parallel token dispatch: every rank sends its own tokens to the experts' owners.""" def install_forward(self, module, mesh, *, is_expert_parallel=False): - from ..integrations.moe import dispatch_experts_forward - original_forward = module.forward ep_group, ep_size = mesh.get_group(), mesh.size() diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 9946bec6ba5a..fb3b736c0301 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -374,83 +374,6 @@ def _grouped_linear( return out -class _ScaleGrad(torch.autograd.Function): - """Identity whose backward scales the gradient.""" - - @staticmethod - def forward(ctx, tensor, scale): - ctx.scale = scale - return tensor - - @staticmethod - def backward(ctx, grad_output): - return grad_output * ctx.scale, None - - -def dispatch_experts_forward( - experts_forward: Callable, - num_local_experts: int, - hidden_states: torch.Tensor, - top_k_index: torch.Tensor, - top_k_weights: torch.Tensor, - ep_group, - ep_size: int, -) -> torch.Tensor: - """ - Expert-parallel forward by token dispatch. Every rank routes its own tokens, sends each selected (token, expert) - pair to the rank that owns the expert with an all-to-all, runs its local experts on what it receives with - `experts_forward` (the experts module's own forward, as a top-1 routing with unit weights), sends the results - back and combines them with the routing weights. Each rank trains on its own part of the batch, so the parameters - outside the experts are data-parallel across the group and averaged by FSDP2; the local experts run on every - rank's tokens, so their gradients already sum every rank's contribution and are scaled by `1 / ep_size` to match. - """ - from torch.distributed.nn.functional import all_to_all_single - - num_tokens, hidden_dim = hidden_states.shape - num_top_k = top_k_index.size(-1) - - # Sorting the selected pairs by expert groups them by owner rank, since each rank owns a contiguous range of - # experts, and the per-expert counts tell every receiver which expert each token it gets is for. The split - # sizes are the one host sync of the layer. - expert_ids = top_k_index.reshape(-1) - order = torch.argsort(expert_ids) - send_tokens = hidden_states[order // num_top_k] - send_counts = torch.zeros(num_local_experts * ep_size, dtype=torch.long, device=hidden_states.device) - send_counts = send_counts.scatter_add_(0, expert_ids, torch.ones_like(expert_ids)).view(ep_size, num_local_experts) - recv_counts = torch.empty_like(send_counts) - torch.distributed.all_to_all_single(recv_counts, send_counts, group=ep_group) - send_sizes, recv_sizes = torch.stack([send_counts.sum(dim=1), recv_counts.sum(dim=1)]).tolist() - recv_tokens = all_to_all_single( - send_tokens.new_empty(sum(recv_sizes), hidden_dim), - send_tokens, - output_split_sizes=recv_sizes, - input_split_sizes=send_sizes, - group=ep_group, - ) - recv_expert_ids = torch.arange(num_local_experts, device=hidden_states.device).repeat(ep_size) - recv_expert_ids = recv_expert_ids.repeat_interleave(recv_counts.reshape(-1), output_size=sum(recv_sizes)) - - # The local experts, as a top-1 routing with unit weights. Scaling the gradient of the output by `1 / ep_size` - # and of the input by `ep_size` leaves the token gradients untouched and scales the expert gradients. - recv_tokens = _ScaleGrad.apply(recv_tokens, ep_size) - unit_weights = torch.ones_like(recv_expert_ids, dtype=recv_tokens.dtype).unsqueeze(-1) - expert_out = experts_forward(recv_tokens, recv_expert_ids.unsqueeze(-1), unit_weights) - expert_out = _ScaleGrad.apply(expert_out, 1.0 / ep_size) - - # Send the results back to the owners of the tokens and combine them with the routing weights. - recv_out = all_to_all_single( - expert_out.new_empty(send_tokens.size(0), hidden_dim), - expert_out, - output_split_sizes=send_sizes, - input_split_sizes=recv_sizes, - group=ep_group, - ) - inverse_order = torch.empty_like(order) - inverse_order[order] = torch.arange(order.numel(), device=order.device) - combined = recv_out[inverse_order] * top_k_weights.reshape(-1, 1) - return combined.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype) - - def grouped_mm_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, From a6e36c09864a6642e86e0c3cd58b56de651643e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Thu, 10 Sep 2026 15:15:10 +0000 Subject: [PATCH 46/52] Name the dispatch strategy instead of switching it with a bool `expert_parallel_dispatch: bool` becomes `experts_dispatch: str`, defaulting to `"all-reduce"`, the path where every rank runs the whole batch and the expert outputs are all-reduced. `"all-to-all"` is the token dispatch. The strategy maps to the parallel style the experts take, so another backend is an entry in that map plus a registered style rather than a second flag. Everything that only means "each rank routes and trains on its own part of the batch" reads the `dispatches_tokens` property, so a new strategy does not have to touch those call sites. --- docs/source/en/expert_parallelism.md | 6 ++-- .../distributed/configuration_utils.py | 33 +++++++++++++++---- src/transformers/distributed/fsdp.py | 2 +- src/transformers/distributed/mixin.py | 16 +++++---- src/transformers/trainer.py | 9 ++--- tests/test_fsdp_mixin.py | 5 ++- tests/test_tensor_parallel_mixin.py | 4 ++- 7 files changed, 51 insertions(+), 24 deletions(-) diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index 07d45d546163..bc487dcb1a41 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -52,7 +52,7 @@ torchrun --nproc-per-node 8 your_script.py ## Token dispatch -By default, every expert parallel rank runs the whole batch, keeps only the experts it owns, and all-reduces expert outputs after every MoE layer. Dense layers then do `tp_size` times the same work, and the all-reduce moves full activations. Set `expert_parallel_dispatch=True` to send each token to the rank that owns its experts. Each rank trains on its own batch shard, and a lot less data is required to travel between GPUs/nodes during large-scale training. +By default, every expert parallel rank runs the whole batch, keeps only the experts it owns, and all-reduces expert outputs after every MoE layer. Dense layers then do `tp_size` times the same work, and the all-reduce moves full activations. Set `experts_dispatch="all-to-all"` to send each token to the rank that owns its experts. Each rank trains on its own batch shard, and a lot less data is required to travel between GPUs/nodes during large-scale training. ```py from transformers import AutoModelForCausalLM @@ -61,7 +61,7 @@ from transformers.distributed import DistributedConfig distributed_config = DistributedConfig( tp_size=8, enable_expert_parallel=True, - expert_parallel_dispatch=True, + experts_dispatch="all-to-all", ) ``` @@ -90,7 +90,7 @@ distributed_config = DistributedConfig( 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. With `expert_parallel_dispatch=True` the non-expert parameters are data-parallel across the whole mesh rather than replicated across `tp`, so FSDP2 shards them across `fsdp` and `tp` together. +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. With `experts_dispatch="all-to-all"` the non-expert parameters are data-parallel across the whole mesh rather than replicated across `tp`, so FSDP2 shards them across `fsdp` and `tp` together. 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`]. diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 88bda47b4434..db2f37da2888 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -18,6 +18,12 @@ from typing import Literal +# How the expert outputs get back to the tokens that need them, mapped to the parallel style the experts take. +# `all-reduce` runs the whole batch on every rank and needs no style of its own. A new backend such as DeepEP is a +# new entry here plus the matching style in `ParallelInterface`, not a new flag. +EXPERTS_DISPATCH_STRATEGIES = {"all-reduce": None, "all-to-all": "ep_dispatch_experts"} + + @dataclass class DistributedConfig: """ @@ -34,10 +40,12 @@ class DistributedConfig: Reserved for sequence parallelism. Not wired up yet. enable_expert_parallel (`bool`, *optional*, defaults to `False`): Route MoE models through the expert-parallel path (``base_model_ep_plan``). - expert_parallel_dispatch (`bool`, *optional*, defaults to `False`): - Send each token to the rank that owns its experts with an all-to-all. Each rank trains on its own part of - the batch, and the parameters that are not expert-parallel are sharded with FSDP2 across every rank. - Requires `enable_expert_parallel`. + experts_dispatch (`str`, *optional*, defaults to `"all-reduce"`): + How the expert outputs get back to the tokens that need them. `"all-reduce"` runs the whole batch on + every rank and all-reduces the expert outputs. `"all-to-all"` sends each token to the rank that owns its + experts instead, so each rank trains on its own part of the batch and the parameters that are not + expert-parallel are sharded with FSDP2 across every rank. Anything but `"all-reduce"` requires + `enable_expert_parallel`. fsdp_size (`int`, *optional*): Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1. fsdp_cpu_offload (`bool`, *optional*, defaults to `False`): @@ -52,12 +60,18 @@ class DistributedConfig: tp_plan: dict[str, str] | Literal["auto"] | None = None enable_sequence_parallel: bool = False enable_expert_parallel: bool = False - expert_parallel_dispatch: bool = False + experts_dispatch: str = "all-reduce" fsdp_size: int | None = None fsdp_cpu_offload: bool = False fsdp_mixed_precision: bool = False pp_size: int | None = None + @property + def dispatches_tokens(self) -> bool: + """Whether each rank routes and trains on its own part of the batch, which is every strategy but the + `all-reduce` default.""" + return self.experts_dispatch != "all-reduce" + def __post_init__(self): if self.tp_plan is None and self.tp_size is None and self.fsdp_size is None and self.pp_size is None: return @@ -78,8 +92,13 @@ def __post_init__(self): elif self.tp_size is None: self.tp_size = 1 - if self.expert_parallel_dispatch and not self.enable_expert_parallel: - raise ValueError("`expert_parallel_dispatch=True` requires `enable_expert_parallel=True`.") + if self.experts_dispatch not in EXPERTS_DISPATCH_STRATEGIES: + raise ValueError( + f"Unknown `experts_dispatch={self.experts_dispatch!r}`, expected one of " + f"{sorted(EXPERTS_DISPATCH_STRATEGIES)}." + ) + if self.dispatches_tokens and not self.enable_expert_parallel: + raise ValueError(f"`experts_dispatch={self.experts_dispatch!r}` requires `enable_expert_parallel=True`.") if self.pp_size > 1 and (self.tp_size > 1 or self.fsdp_size > 1): raise ValueError( diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 905a35e52959..07782b026941 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -212,7 +212,7 @@ def apply_fully_sharded_data_parallelism( adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, model) reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) - if distributed_config is not None and distributed_config.expert_parallel_dispatch: + if distributed_config is not None and distributed_config.dispatches_tokens: if not is_torch_greater_or_equal("2.7"): raise OSError("Expert-parallel token dispatch requires `torch>=2.7`.") # The DTensor parameters are the expert-parallel experts: `maybe_distribute_model` rewrote the expert diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 9eed97deada5..f30021bd4fc4 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -20,7 +20,7 @@ from ..utils import is_torch_greater_or_equal, logging from ..utils.hub import create_and_tag_model_card -from .configuration_utils import DistributedConfig +from .configuration_utils import EXPERTS_DISPATCH_STRATEGIES, DistributedConfig from .fsdp import apply_fully_sharded_data_parallelism, is_fsdp_managed_module from .pipeline_parallel import apply_pipeline_parallelism from .tensor_parallel import ( @@ -202,7 +202,7 @@ def maybe_distribute_model( # 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 - model._expert_parallel_dispatch = distributed_config.expert_parallel_dispatch + model._expert_parallel_dispatch = distributed_config.dispatches_tokens # Both may apply: the tensor/expert parallel plan shards across `tp` first, then FSDP2 # shards every parameter (the `tp`-sharded ones included) across `fsdp`. @@ -210,7 +210,7 @@ def maybe_distribute_model( 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 - if distributed_config.expert_parallel_dispatch: + if distributed_config.dispatches_tokens: # Every rank trains on its own part of the batch, so only the experts can be sharded across the # group: the experts get the dispatch style, the router keeps its global ids and scores, and # whatever else the plan shards stays replicated, data-parallel like the rest of the trunk. @@ -222,18 +222,20 @@ def maybe_distribute_model( ) if replicated: logger.warning( - "`expert_parallel_dispatch=True` shards only the experts, so these expert parallel plan " + f"`experts_dispatch={distributed_config.experts_dispatch!r}` shards only the experts, " + "so these expert parallel plan " f"entries are ignored and their modules stay replicated: {replicated}." ) # `tp_plan` reads `_ep_plan` under expert parallelism, so that is the plan to rewrite. + dispatch_style = EXPERTS_DISPATCH_STRATEGIES[distributed_config.experts_dispatch] model._ep_plan = { - name: "ep_dispatch_experts" if style == "moe_tp_experts" else style + name: dispatch_style if style == "moe_tp_experts" else style for name, style in model.tp_plan.items() if style in kept } model = apply_tensor_parallelism(model, tp_mesh) - if distributed_config.expert_parallel_dispatch: + if distributed_config.dispatches_tokens: # Every expert-parallel rank trains on its own part of the batch, so the parameters outside the # experts are data-parallel across the whole mesh: FSDP2 shards them across all of it and owns # their gradient reduction. The experts stay sharded across `tp` and, if any, across `fsdp`. @@ -307,7 +309,7 @@ def gather_sharded_state_dict_for_save( if distributed_config is None: return state_dict - if distributed_config.fsdp_size > 1 or distributed_config.expert_parallel_dispatch: + if distributed_config.fsdp_size > 1 or distributed_config.dispatches_tokens: # Also covers the 2-D (fsdp, tp) mesh and token dispatch: every parameter is FSDP-managed, and # the full state dict is only materialized on rank 0. if not _is_torch_distributed_initialized(): diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 72a5c2d172ed..f0ef147c469c 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -622,15 +622,16 @@ def __init__( self._expert_parallel_dispatch = getattr(model, "_expert_parallel_dispatch", False) if self._expert_parallel_dispatch and args.train_sampling_strategy != "random": raise ValueError( - "`expert_parallel_dispatch=True` splits the batches across every rank with a `DistributedSampler`, " - f"which `train_sampling_strategy='{args.train_sampling_strategy}'` does not go through." + "`experts_dispatch` other than 'all-reduce' splits the batches across every rank with a " + f"`DistributedSampler`, which `train_sampling_strategy='{args.train_sampling_strategy}'` does not " + "go through." ) if self._expert_parallel_dispatch and ( self.accelerator.dispatch_batches or (train_dataset is not None and not has_length(train_dataset)) ): raise ValueError( - "`expert_parallel_dispatch=True` splits the batches across every rank with a `DistributedSampler`, " - "which needs a sized training dataset and `dispatch_batches=False`." + "`experts_dispatch` other than 'all-reduce' splits the batches across every rank with a " + "`DistributedSampler`, which needs a sized training dataset and `dispatch_batches=False`." ) if ( getattr(model, "_device_mesh", None) is not None diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index 413ce8bff6c0..842391d59729 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -540,7 +540,10 @@ def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, init_model_dir, torch_dtype=dtype, distributed_config=DistributedConfig( - tp_size=2, fsdp_size=dp, enable_expert_parallel=True, expert_parallel_dispatch=dispatch + tp_size=2, + fsdp_size=dp, + enable_expert_parallel=True, + experts_dispatch="all-to-all" if dispatch else "all-reduce", ), ) assert model.tp_size == 2 and model.fsdp_size == dp diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 4a3fda34db8c..7ce074aec793 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -394,7 +394,9 @@ def _load_ep_and_reference_models(model_path, model_class, dispatch=False): model_ep = model_class.from_pretrained( model_path, distributed_config=DistributedConfig( - tp_size=dist.get_world_size(), enable_expert_parallel=True, expert_parallel_dispatch=dispatch + tp_size=dist.get_world_size(), + enable_expert_parallel=True, + experts_dispatch="all-to-all" if dispatch else "all-reduce", ), ) dist.barrier() From 6a86e1812f839126cf3fa07aa18ee454af26b266 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Fri, 11 Sep 2026 19:24:12 +0000 Subject: [PATCH 47/52] Pass the ignored params through the policy helper `ignored_params` was set on the dict the helper had just returned. It is an argument now, so the helper is the only place that builds those kwargs. --- src/transformers/distributed/fsdp.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 07782b026941..b9ea5b1e9005 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -61,12 +61,16 @@ def is_fsdp_managed_module(module: nn.Module) -> bool: return isinstance(module, FullyShardedDataParallel) -def _get_fsdp_policy_kwargs(distributed_config: DistributedConfig | None) -> dict[str, Any]: +def _get_fsdp_policy_kwargs( + distributed_config: DistributedConfig | None, ignored_params: set[torch.nn.Parameter] | None = None +) -> dict[str, Any]: """Build ``fully_shard`` policy kwargs from ``DistributedConfig`` runtime flags.""" + fsdp_policy_kwargs = {} + if ignored_params: + fsdp_policy_kwargs["ignored_params"] = ignored_params if distributed_config is None: - return {} + return fsdp_policy_kwargs - fsdp_policy_kwargs = {} if distributed_config.fsdp_cpu_offload: fsdp_policy_kwargs["offload_policy"] = CPUOffloadPolicy() if distributed_config.fsdp_mixed_precision: @@ -207,11 +211,11 @@ def apply_fully_sharded_data_parallelism( ) distributed_config = getattr(model.config, "distributed_config", None) - fsdp_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config) adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, model) reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) + ignored_params = None if distributed_config is not None and distributed_config.dispatches_tokens: if not is_torch_greater_or_equal("2.7"): raise OSError("Expert-parallel token dispatch requires `torch>=2.7`.") @@ -221,10 +225,13 @@ def apply_fully_sharded_data_parallelism( module for module in model.modules() if any(is_dtensor(p) for p in module.parameters(recurse=False)) ] if expert_mesh is not None: + expert_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config) for module in expert_modules: - fully_shard(module, mesh=expert_mesh, reshard_after_forward=True, **fsdp_policy_kwargs) + fully_shard(module, mesh=expert_mesh, reshard_after_forward=True, **expert_policy_kwargs) else: - fsdp_policy_kwargs["ignored_params"] = {p for module in expert_modules for p in module.parameters()} + ignored_params = {p for module in expert_modules for p in module.parameters()} + + fsdp_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config, ignored_params=ignored_params) for module_name, module in reshard_targets: fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=True, **fsdp_policy_kwargs) From 5120748873b981d38c5f48087ed1ac0b0f8e5375 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Sat, 12 Sep 2026 14:40:52 +0000 Subject: [PATCH 48/52] Drop the comment on the parallel sizes --- src/transformers/distributed/mixin.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 957289aeffa9..646b8ac102cc 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -172,8 +172,6 @@ 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 From 0147840fa50ef57164bcab05794232511a4538fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Sat, 12 Sep 2026 14:45:10 +0000 Subject: [PATCH 49/52] Trim the parallelism config comment --- src/transformers/trainer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index b08ba795bfda..a4c927f34c12 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -766,8 +766,7 @@ def _build_accelerator_args(self, **kwargs) -> dict[str, Any]: 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, or it sees unaccounted ranks and wraps - # the DTensor model in DDP, which raises. + # 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 3439765fbca910df559cc1fbff53ca2aa4b37017 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Sun, 13 Sep 2026 02:17:47 +0000 Subject: [PATCH 50/52] Address review: guard dispatch + PP, derive the mesh name, move the guards - `experts_dispatch` with `pp_size > 1` raised `KeyError: fsdp_tp` when looking up the flattened sub-mesh, since the mesh is then `(pp, tp)`. Rejected in `__post_init__`, next to the FSDP + PP guard. - The flattened sub-mesh is named after the mesh's own dimensions rather than hardcoded, so it follows `initialize_distributed_mesh`. - `_ScaleGrad` moves inside the `is_torch_distributed_available()` block with the other two autograd functions: at module level it made the import fail without torch. - The `torch>=2.7` requirement is asserted in `_prepare_distributed_setup` with the FSDP2 one, before the mesh is built, instead of after tensor parallelism has been applied. --- .../distributed/configuration_utils.py | 6 +++++ src/transformers/distributed/fsdp.py | 2 -- src/transformers/distributed/mixin.py | 5 +++- .../distributed/tensor_parallel.py | 25 +++++++++---------- 4 files changed, 22 insertions(+), 16 deletions(-) diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 942bdeafbfc3..88d1753e836d 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -100,6 +100,12 @@ def __post_init__(self): if self.dispatches_tokens and not self.enable_expert_parallel: raise ValueError(f"`experts_dispatch={self.experts_dispatch!r}` requires `enable_expert_parallel=True`.") + if self.dispatches_tokens and self.pp_size > 1: + raise ValueError( + f"Combining `experts_dispatch={self.experts_dispatch!r}` with pipeline parallelism is not " + "supported yet." + ) + if self.fsdp_size > 1 and self.pp_size > 1: raise ValueError( "Combining FSDP with pipeline parallelism is not supported yet. " diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 3586038bcfc8..f477820e820c 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -217,8 +217,6 @@ def apply_fully_sharded_data_parallelism( ignored_params = None if distributed_config is not None and distributed_config.dispatches_tokens: - if not is_torch_greater_or_equal("2.7"): - raise OSError("Expert-parallel token dispatch requires `torch>=2.7`.") # The DTensor parameters are the expert-parallel experts: `maybe_distribute_model` rewrote the expert # parallel plan to shard only them. expert_modules = [ diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index bce189d83959..550ac7c8973d 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -157,6 +157,8 @@ def prepare_distribute_model( raise ValueError("Tensor parallelism and `device_map` are mutually exclusive.") if distributed_config.fsdp_size > 1 and not is_torch_greater_or_equal("2.7"): raise OSError("FSDP2 requires `torch>=2.7` (distributed checkpoint save/load).") + if distributed_config.dispatches_tokens and not is_torch_greater_or_equal("2.7"): + raise OSError("Expert-parallel token dispatch requires `torch>=2.7`.") device_map, device_mesh = initialize_distributed_mesh(distributed_config) @@ -216,7 +218,8 @@ def maybe_distribute_model( # Every expert-parallel rank trains on its own part of the batch, so the parameters outside the # experts are data-parallel across the whole mesh: FSDP2 shards them across all of it and owns # their gradient reduction. The experts stay sharded across `tp` and, if any, across `fsdp`. - trunk_mesh = device_mesh["fsdp_tp"] if device_mesh.ndim > 1 else device_mesh + flattened = "_".join(device_mesh.mesh_dim_names) + trunk_mesh = device_mesh[flattened] if device_mesh.ndim > 1 else device_mesh expert_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else None model = apply_fully_sharded_data_parallelism(model, trunk_mesh, expert_mesh=expert_mesh) elif distributed_config.fsdp_size > 1: diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 4fcf0d7c56fb..d5eb842a7296 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -585,6 +585,18 @@ def backward(ctx, grad): dist.all_reduce(grad, group=ctx.process_group) return grad, None + class _ScaleGrad(torch.autograd.Function): + """Identity whose backward scales the gradient.""" + + @staticmethod + def forward(ctx, tensor, scale): + ctx.scale = scale + return tensor + + @staticmethod + def backward(ctx, grad_output): + return grad_output * ctx.scale, None + class MoeExpertsParallel(TensorParallelLayer): def should_use_local_tensors(self, module): @@ -740,19 +752,6 @@ def transform_output_post_forward(self, module, output, mesh): return output -class _ScaleGrad(torch.autograd.Function): - """Identity whose backward scales the gradient.""" - - @staticmethod - def forward(ctx, tensor, scale): - ctx.scale = scale - return tensor - - @staticmethod - def backward(ctx, grad_output): - return grad_output * ctx.scale, None - - def dispatch_experts_forward( experts_forward: Callable, num_local_experts: int, From 4f2a6b9110be8db157b11852bfdce30c5d995ef6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Sun, 13 Sep 2026 02:17:52 +0000 Subject: [PATCH 51/52] Keep reading the tensor parallel size from accelerate in `get_tp_size` `_get_num_items_in_batch` and the loss scale now go through `get_tp_size`, which only reads `model._tp_size`. Tensor parallelism configured through `TrainingArguments(parallelism_config=...)` rather than `DistributedConfig` leaves that unset, so both lost their `tp_size` divisor. Fall back to `parallelism_config.tp_size`, which is what they read before. --- src/transformers/trainer.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 529a0ed47210..0a35c81a6bdb 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2579,7 +2579,11 @@ def get_tp_size(self) -> int: if self.is_deepspeed_enabled and (deepspeed_config := getattr(self.args, "hf_deepspeed_config", None)): return deepspeed_config.config.get("tensor_parallel", {}).get("autotp_size", 1) - # 3. Default fallback + # 3. Fall back to accelerate, for tensor parallelism configured outside `DistributedConfig` + if (pc := getattr(self.accelerator, "parallelism_config", None)) is not None: + return pc.tp_size + + # 4. Default fallback return 1 def _wrap_model(self, model: nn.Module, training: bool = True, dataloader: DataLoader | None = None) -> nn.Module: From 86561b69c95cf0906fba686896240b2ad8235808 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Sun, 13 Sep 2026 02:43:01 +0000 Subject: [PATCH 52/52] Give the token dispatch weight comparison its own tolerance `test_fsdp2_expert_parallel_2d_vs_ddp` failed on qwen3_moe under dispatch, on `layers.1.self_attn.k_proj.weight`, 2.6e-5 against an atol of 1e-5. The gradients are not the problem: the same single step with SGD, whose update is proportional to the gradient, matches to 6e-11 and passes. Adam normalises each step to about `lr * sign(grad)`, so an element whose gradient is near zero turns the last bits of a reduction over every rank rather than over `fsdp` into an `lr`-sized weight difference, flat across 1, 5 and 20 steps. Losses and gradient norms keep the tight tolerance, matching to 1e-6 at all 20 steps, and they are what says the reduction is right. --- tests/test_fsdp_mixin.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index 842391d59729..54e1f3635d5d 100644 --- a/tests/test_fsdp_mixin.py +++ b/tests/test_fsdp_mixin.py @@ -583,13 +583,19 @@ def _test_fsdp2_expert_parallel_2d_vs_ddp_impl(rank, config_class, config_dict, msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, FSDP2+EP={grad_norms[step]}", ) + # Adam normalises each step to about `lr * sign(grad)`, so an element whose gradient is near zero turns a + # rounding-level difference into an `lr`-sized weight difference. Token dispatch reduces the gradients over + # every rank rather than over `fsdp`, which changes those last bits: the same single step with SGD, whose + # update is proportional to the gradient, matches to 6e-11. Losses and gradient norms keep the tight + # tolerance, and they are what says the reduction is right. + weight_atol = 1e-4 if dispatch else DDP_FSDP_ATOL for key in ddp_state_dict: assert key in state_dict, f"Key {key} missing from FSDP2+EP state dict" torch.testing.assert_close( ddp_state_dict[key], state_dict[key], rtol=DDP_FSDP_RTOL, - atol=DDP_FSDP_ATOL, + atol=weight_atol, msg=f"Weight mismatch for {key}: DDP vs FSDP2+EP", )