diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md index 307d8f8d9ee1..bc487dcb1a41 100644 --- a/docs/source/en/expert_parallelism.md +++ b/docs/source/en/expert_parallelism.md @@ -50,9 +50,33 @@ 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 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 +from transformers.distributed import DistributedConfig + +distributed_config = DistributedConfig( + tp_size=8, + enable_expert_parallel=True, + experts_dispatch="all-to-all", +) +``` + +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. + +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. + ## 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 @@ -66,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 `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 55e7c85013dc..88d1753e836d 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,6 +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``). + 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`): @@ -48,11 +60,18 @@ class DistributedConfig: tp_plan: dict[str, str] | Literal["auto"] | None = None enable_sequence_parallel: bool = False enable_expert_parallel: 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 @@ -73,6 +92,20 @@ def __post_init__(self): elif self.tp_size is None: self.tp_size = 1 + 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.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 aef3a5979a8c..f477820e820c 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: @@ -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: @@ -185,13 +189,19 @@ 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_distributed_mesh`. + + With expert-parallel token dispatch `fsdp_mesh` spans the expert-parallel ranks, which the experts are already + 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: @@ -201,11 +211,26 @@ 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: + # 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)) + ] + 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, **expert_policy_kwargs) + else: + 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) 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 646b8ac102cc..550ac7c8973d 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 ( @@ -53,6 +53,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 @@ -156,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) @@ -174,6 +177,7 @@ def maybe_distribute_model( model._device_mesh = device_mesh model._tp_size = distributed_config.tp_size model._fsdp_size = distributed_config.fsdp_size + model._expert_parallel_dispatch = distributed_config.dispatches_tokens if distributed_config.pp_size > 1: pp_mesh = device_mesh["pp"] if device_mesh.ndim > 1 else device_mesh @@ -185,9 +189,40 @@ 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.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. + # 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", *kept) + ) + if replicated: + logger.warning( + 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: 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.fsdp_size > 1: + 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`. + 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: fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh model = apply_fully_sharded_data_parallelism(model, fsdp_mesh) return model @@ -250,9 +285,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 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(): 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 3151f6264a12..d5eb842a7296 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 @@ -584,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): @@ -739,6 +752,90 @@ def transform_output_post_forward(self, module, output, mesh): return output +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): + original_forward = module.forward + ep_group, ep_size = mesh.get_group(), mesh.size() + + 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): + # 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 + + class MoeTensorParalellMegaMoeExperts(MoeExpertsParallel): """TP layer for DeepGEMM Mega MoE experts. @@ -776,6 +873,7 @@ class ParallelInterface(GeneralInterface): "sequence_parallel": SequenceParallel(use_local_output=True), "grouped_gemm": MoEParamShard(Shard(0), shards_expert_dim=True), "ep_router": EpRouterParallel(), + "ep_dispatch_experts": EpDispatchExpertsParallel(), "megamoe_router": RouterParallelMegaMoe(), "moe_tp_experts": MoeExpertsParallel(), "megamoe_experts": MoeTensorParalellMegaMoeExperts(), diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index a4c927f34c12..0a35c81a6bdb 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 @@ -617,6 +617,22 @@ 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. + self._expert_parallel_dispatch = getattr(model, "_expert_parallel_dispatch", False) + if self._expert_parallel_dispatch and args.train_sampling_strategy != "random": + raise ValueError( + "`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( + "`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 and args.save_strategy != SaveStrategy.NO @@ -1037,12 +1053,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: + 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 @@ -1129,6 +1145,15 @@ 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.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) def _get_eval_sampler(self, eval_dataset: Dataset) -> torch.utils.data.Sampler | None: @@ -2156,9 +2181,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 - if (pc := getattr(self.accelerator, "parallelism_config", None)) is not None: - 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 @@ -2307,8 +2330,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): - 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 @@ -2545,7 +2569,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 @@ -2553,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: diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py index e780b74f56c9..54e1f3635d5d 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,24 @@ 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, + experts_dispatch="all-to-all" if dispatch else "all-reduce", + ), ) 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 +564,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)): @@ -573,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", ) @@ -731,15 +747,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 22d92bdd0948..357ec4e86bda 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -390,11 +390,15 @@ 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, + experts_dispatch="all-to-all" if dispatch else "all-reduce", + ), ) dist.barrier() @@ -405,11 +409,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, experts_implementation): +def _test_ep_forward_impl(_rank, model_path, model_class, atol, rtol, experts_implementation, 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() @@ -432,11 +436,11 @@ def _test_ep_forward_impl(_rank, model_path, model_class, atol, rtol, experts_im dist.barrier() -def _test_ep_backward_impl(_rank, model_path, model_class, atol, rtol, experts_implementation): +def _test_ep_backward_impl(_rank, model_path, model_class, atol, rtol, experts_implementation, 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() @@ -459,6 +463,27 @@ def _test_ep_backward_impl(_rank, model_path, model_class, atol, rtol, experts_i 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() @@ -647,15 +672,12 @@ def test_tp_generation_quantized(self): ) @parameterized.expand( - list( - product( - [False, True], # tie_word_embeddings - ["eager", "grouped_mm", "batched_mm"], # experts_implementation - ) - ) + [(tie, impl, False) for tie, impl in product([False, True], ["eager", "grouped_mm", "batched_mm"])] + # Token dispatch is orthogonal to the implementation, so it adds the one combination on its own. + + [(False, "eager", True)] ) @is_tensor_parallel_test - def test_ep_forward(self, tie_word_embeddings, experts_implementation): + def test_ep_forward(self, tie_word_embeddings, experts_implementation, dispatch): self._skip_if_not_supported(expert_parallel=True) config = self._get_tp_config(tie_word_embeddings=tie_word_embeddings) @@ -669,12 +691,12 @@ def test_ep_forward(self, tie_word_embeddings, experts_implementation): 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, experts_implementation + tmp_dir, model_class, atol, rtol, experts_implementation, dispatch=dispatch ) - @parameterized.expand([("eager",), ("grouped_mm",), ("batched_mm",)]) + @parameterized.expand([("eager", False), ("grouped_mm", False), ("batched_mm", False), ("eager", True)]) @is_tensor_parallel_test - def test_ep_backward(self, experts_implementation): + def test_ep_backward(self, experts_implementation, dispatch): self._skip_if_not_supported(expert_parallel=True) config = self._get_tp_config() @@ -688,5 +710,5 @@ def test_ep_backward(self, experts_implementation): 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, experts_implementation + tmp_dir, model_class, atol, rtol, experts_implementation, dispatch=dispatch )