diff --git a/autoparallel/_testing/models/dsv3.py b/autoparallel/_testing/models/dsv3.py index 5a897b71..0294a060 100644 --- a/autoparallel/_testing/models/dsv3.py +++ b/autoparallel/_testing/models/dsv3.py @@ -14,10 +14,22 @@ import triton.language as tl from torch import nn from torch.distributed.tensor import DeviceMesh, DTensor -from torch.distributed.tensor.placement_types import Partial, Replicate, Shard +from torch.distributed.tensor.placement_types import ( + Partial, + Placement, + Replicate, + Shard, +) from torch.nn.attention import SDPBackend, sdpa_kernel -from autoparallel.collectives import all_to_all, axis_size, local_map +from autoparallel.collectives import ( + all_gather, + all_reduce, + all_to_all, + axis_size, + get_mesh_from_global, + local_map, +) _MODULE_FQN = "module_fqn" @@ -271,12 +283,184 @@ def _run_experts_for_loop( return torch.cat(out_experts_splits, dim=0) +PlacementTuple = tuple[Placement, ...] + + +def _mesh_axis_names(mesh: DeviceMesh) -> tuple[str, ...]: + mesh_axis_names = mesh.mesh_dim_names + if mesh_axis_names is None: + raise ValueError("MoE local_map requires named DeviceMesh axes") + return tuple(mesh_axis_names) + + +def _validate_axis_name(mesh_axis_names: tuple[str, ...], axis_name: str) -> None: + if axis_name not in mesh_axis_names: + raise ValueError( + f"MoE axis {axis_name!r} is not in mesh axes {mesh_axis_names}" + ) + + +def _axis_size_from_mesh(mesh: DeviceMesh, axis_name: str) -> int: + mesh_axis_names = _mesh_axis_names(mesh) + _validate_axis_name(mesh_axis_names, axis_name) + return mesh.size(mesh_axis_names.index(axis_name)) + + +def _default_token_axis_names( + mesh_axis_names: tuple[str, ...], + expert_tensor_parallel_axis_name: str | None, +) -> tuple[str, ...]: + return tuple( + axis_name + for axis_name in mesh_axis_names + if axis_name != expert_tensor_parallel_axis_name + ) + + +def _placement_from_axis_map( + mesh_axis_names: tuple[str, ...], + axis_placements: dict[str, Placement], +) -> PlacementTuple: + return tuple( + axis_placements.get(axis_name, Replicate()) for axis_name in mesh_axis_names + ) + + +def _resolve_expert_fsdp_shard_dim( + mesh: DeviceMesh, + ep_axis_name: str, + expert_tensor_parallel_axis_name: str | None, + expert_fsdp_axis_name: str | None, + expert_fsdp_shard_dim: int | None, + num_experts: int | None, +) -> int | None: + mesh_axis_names = _mesh_axis_names(mesh) + if expert_fsdp_axis_name is None: + if expert_fsdp_shard_dim is not None: + raise ValueError("MoE EFSDP shard dim requires an EFSDP axis") + return None + + _validate_axis_name(mesh_axis_names, expert_fsdp_axis_name) + if expert_fsdp_axis_name == ep_axis_name: + raise ValueError("MoE EFSDP and EP axes must differ") + if expert_fsdp_axis_name == expert_tensor_parallel_axis_name: + raise ValueError("MoE EFSDP and expert tensor parallel axes must differ") + + if expert_fsdp_shard_dim is None: + if num_experts is None: + raise ValueError("MoE EFSDP placement requires num_experts") + ep_size = _axis_size_from_mesh(mesh, ep_axis_name) + efsdp_size = _axis_size_from_mesh(mesh, expert_fsdp_axis_name) + resolved = 1 if efsdp_size * ep_size > num_experts else 0 + elif expert_fsdp_shard_dim in (0, 1): + resolved = expert_fsdp_shard_dim + else: + raise ValueError("MoE EFSDP shard dim must be 0 or 1") + + if expert_tensor_parallel_axis_name is not None and resolved != 0: + raise ValueError("MoE EFSDP Shard(1) is not supported with ETP") + return resolved + + +def _moe_local_map_placements( + mesh: DeviceMesh, + ep_axis_name: str, + token_axis_names: tuple[str, ...] | None, + expert_tensor_parallel_axis_name: str | None, + expert_fsdp_axis_name: str | None = None, + expert_fsdp_shard_dim: int | None = None, + num_experts: int | None = None, +) -> tuple[ + PlacementTuple, + PlacementTuple, + PlacementTuple, + PlacementTuple, + PlacementTuple, +]: + mesh_axis_names = _mesh_axis_names(mesh) + _validate_axis_name(mesh_axis_names, ep_axis_name) + if expert_tensor_parallel_axis_name is not None: + _validate_axis_name(mesh_axis_names, expert_tensor_parallel_axis_name) + if expert_tensor_parallel_axis_name == ep_axis_name: + raise ValueError("MoE EP and expert tensor parallel axes must differ") + + resolved_expert_fsdp_shard_dim = _resolve_expert_fsdp_shard_dim( + mesh, + ep_axis_name, + expert_tensor_parallel_axis_name, + expert_fsdp_axis_name, + expert_fsdp_shard_dim, + num_experts, + ) + + if token_axis_names is None: + token_axis_names = _default_token_axis_names( + mesh_axis_names, expert_tensor_parallel_axis_name + ) + for axis_name in token_axis_names: + _validate_axis_name(mesh_axis_names, axis_name) + + token_axes = set(token_axis_names) + if expert_tensor_parallel_axis_name in token_axes: + raise ValueError( + "MoE token axes must not include the expert tensor parallel axis" + ) + + token_placement = _placement_from_axis_map( + mesh_axis_names, {axis_name: Shard(0) for axis_name in token_axes} + ) + expert_weight_axis_placements = {ep_axis_name: Shard(0)} + if expert_fsdp_axis_name is not None: + assert resolved_expert_fsdp_shard_dim is not None + expert_weight_axis_placements[expert_fsdp_axis_name] = Shard( + resolved_expert_fsdp_shard_dim + ) + + expert_w1_placement = _placement_from_axis_map( + mesh_axis_names, + { + **expert_weight_axis_placements, + **( + {expert_tensor_parallel_axis_name: Shard(1)} + if expert_tensor_parallel_axis_name is not None + else {} + ), + }, + ) + expert_w3_placement = expert_w1_placement + expert_w2_placement = _placement_from_axis_map( + mesh_axis_names, + { + **expert_weight_axis_placements, + **( + {expert_tensor_parallel_axis_name: Shard(2)} + if expert_tensor_parallel_axis_name is not None + else {} + ), + }, + ) + token_count_placement = _placement_from_axis_map( + mesh_axis_names, + {axis_name: Partial(reduce_op="sum") for axis_name in token_axes}, + ) + return ( + token_placement, + expert_w1_placement, + expert_w3_placement, + expert_w2_placement, + token_count_placement, + ) + + def _run_experts_grouped_mm( w1: torch.Tensor, w2: torch.Tensor, w3: torch.Tensor, x: torch.Tensor, num_tokens_per_expert: torch.Tensor, + expert_tensor_parallel_axis_name: str | None, + expert_fsdp_axis_name: str | None, + expert_fsdp_shard_dim: int | None, ) -> torch.Tensor: if isinstance(w1, DTensor): assert isinstance(w2, DTensor) @@ -285,6 +469,12 @@ def _run_experts_grouped_mm( w2 = w2.to_local() w3 = w3.to_local() + if expert_fsdp_axis_name is not None: + assert expert_fsdp_shard_dim is not None + w1 = all_gather(w1, expert_fsdp_shard_dim, expert_fsdp_axis_name) + w2 = all_gather(w2, expert_fsdp_shard_dim, expert_fsdp_axis_name) + w3 = all_gather(w3, expert_fsdp_shard_dim, expert_fsdp_axis_name) + offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) # grouped mm between a 2D tensor and a 3D tensor assert x.dim() == 2 @@ -296,6 +486,8 @@ def _run_experts_grouped_mm( x.bfloat16(), w3.bfloat16().transpose(-2, -1), offs=offsets ) out = torch._grouped_mm(h, w2.bfloat16().transpose(-2, -1), offs=offsets).type_as(x) + if expert_tensor_parallel_axis_name is not None: + out = all_reduce(out, expert_tensor_parallel_axis_name) return out @@ -322,7 +514,14 @@ def forward( ) -> torch.Tensor: if self.use_grouped_mm: return _run_experts_grouped_mm( - self.w1, self.w2, self.w3, x, num_tokens_per_expert + self.w1, + self.w2, + self.w3, + x, + num_tokens_per_expert, + None, + None, + None, ) else: return _run_experts_for_loop( @@ -629,6 +828,9 @@ def local_mapped_region( num_experts: int, score_before_experts: bool, axis_name: str, + expert_tensor_parallel_axis_name: str | None, + expert_fsdp_axis_name: str | None, + expert_fsdp_shard_dim: int | None, ) -> tuple[torch.Tensor, torch.Tensor]: # assert False, f"{x.shape}, {selected_experts_indices.shape}, {top_scores.shape}, {out.shape}" @@ -677,6 +879,9 @@ def local_mapped_region( experts_w3, routed_input, num_tokens_per_expert_group, + expert_tensor_parallel_axis_name, + expert_fsdp_axis_name, + expert_fsdp_shard_dim, ) routed_output = _token_combine( @@ -809,6 +1014,10 @@ def _moe_forward( reorderer: TokenReorderer, mesh: Optional[DeviceMesh], axis_name: str, + token_axis_names: tuple[str, ...] | None, + expert_tensor_parallel_axis_name: str | None, + expert_fsdp_axis_name: str | None, + expert_fsdp_shard_dim: int | None, score_before_experts: bool, compute_dtype: torch.dtype | None = None, ): @@ -855,23 +1064,52 @@ def _moe_forward( # This is in the local_map region ###################################################### - # expert_placements = ((Replicate(), Shard(0)),) * 3 - # in_placements = ( - # (Shard(0), Shard(0)), - # (Shard(0), Shard(0)), - # (Shard(0), Shard(0)), - # (Shard(0), Shard(0)), - # ) + if mesh is None: + mesh = get_mesh_from_global() + if mesh is None: + raise ValueError("MoE local_map requires an explicit DeviceMesh") + ep_size = _axis_size_from_mesh(mesh, axis_name) + if router.num_experts % ep_size != 0: + raise ValueError( + f"num_experts ({router.num_experts}) must be divisible by " + f"EP axis {axis_name!r} size ({ep_size})" + ) + resolved_expert_fsdp_shard_dim = _resolve_expert_fsdp_shard_dim( + mesh, + axis_name, + expert_tensor_parallel_axis_name, + expert_fsdp_axis_name, + expert_fsdp_shard_dim, + router.num_experts, + ) + ( + token_placement, + expert_w1_placement, + expert_w3_placement, + expert_w2_placement, + token_count_placement, + ) = _moe_local_map_placements( + mesh, + axis_name, + token_axis_names, + expert_tensor_parallel_axis_name, + expert_fsdp_axis_name, + resolved_expert_fsdp_shard_dim, + router.num_experts, + ) # Dynamo reorders captured variables (lifted freevars) before explicit # arguments, so x must come first in the input order and placements. reordered_placements = ( - (Shard(0), Shard(0)), - (Shard(0), Shard(0)), - (Shard(0), Shard(0)), - (Replicate(), Shard(0)), - (Replicate(), Shard(0)), - (Replicate(), Shard(0)), - (Shard(0), Shard(0)), + token_placement, + token_placement, + token_placement, + expert_w1_placement, + expert_w3_placement, + expert_w2_placement, + token_placement, + None, + None, + None, None, None, None, @@ -881,8 +1119,8 @@ def _moe_forward( out, num_tokens_per_expert = local_map( local_mapped_region, out_placements=( - (Shard(0), Shard(0)), - (Partial(reduce_op="sum"), Partial(reduce_op="sum")), + token_placement, + token_count_placement, ), in_placements=reordered_placements, redistribute_inputs=True, @@ -900,6 +1138,9 @@ def _moe_forward( router.num_experts, score_before_experts, axis_name, + expert_tensor_parallel_axis_name, + expert_fsdp_axis_name, + resolved_expert_fsdp_shard_dim, ) # assert False, f"there: {out.shape}, {num_tokens_per_expert.shape}" @@ -933,12 +1174,21 @@ def __init__( use_grouped_mm: bool = True, load_balance_coeff: float | None = 1e-3, mesh: DeviceMesh | None = None, + ep_axis_name: str = "ep", + token_axis_names: tuple[str, ...] | None = None, + expert_tensor_parallel_axis_name: str | None = None, + expert_fsdp_axis_name: str | None = None, + expert_fsdp_shard_dim: int | None = None, compute_dtype: torch.dtype | None = None, ): super().__init__() self.mesh = mesh - self.axis_name = "ep" + self.axis_name = ep_axis_name + self.token_axis_names = token_axis_names + self.expert_tensor_parallel_axis_name = expert_tensor_parallel_axis_name + self.expert_fsdp_axis_name = expert_fsdp_axis_name + self.expert_fsdp_shard_dim = expert_fsdp_shard_dim self.compute_dtype = compute_dtype self.experts = GroupedExperts( dim=dim, @@ -999,6 +1249,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self.reorderer, self.mesh, self.axis_name, + self.token_axis_names, + self.expert_tensor_parallel_axis_name, + self.expert_fsdp_axis_name, + self.expert_fsdp_shard_dim, self.score_before_experts, self.compute_dtype, ) @@ -1555,6 +1809,11 @@ def __init__( layer_config, model_config, mesh: DeviceMesh | None = None, + ep_axis_name: str = "ep", + token_axis_names: tuple[str, ...] | None = None, + expert_tensor_parallel_axis_name: str | None = None, + expert_fsdp_axis_name: str | None = None, + expert_fsdp_shard_dim: int | None = None, compute_dtype: torch.dtype | None = None, ): super().__init__() @@ -1584,6 +1843,11 @@ def __init__( use_grouped_mm=moe_cfg.experts.use_grouped_mm, load_balance_coeff=moe_cfg.load_balance_coeff, mesh=mesh, + ep_axis_name=ep_axis_name, + token_axis_names=token_axis_names, + expert_tensor_parallel_axis_name=expert_tensor_parallel_axis_name, + expert_fsdp_axis_name=expert_fsdp_axis_name, + expert_fsdp_shard_dim=expert_fsdp_shard_dim, compute_dtype=compute_dtype, ) else: @@ -1639,6 +1903,11 @@ def __init__( self, config, mesh: DeviceMesh | None = None, + ep_axis_name: str = "ep", + token_axis_names: tuple[str, ...] | None = None, + expert_tensor_parallel_axis_name: str | None = None, + expert_fsdp_axis_name: str | None = None, + expert_fsdp_shard_dim: int | None = None, compute_dtype: torch.dtype | None = None, ): # Explicitly call nn.Module.__init__ to avoid MRO issues when this class @@ -1658,6 +1927,11 @@ def __init__( layer_config, config, mesh, + ep_axis_name=ep_axis_name, + token_axis_names=token_axis_names, + expert_tensor_parallel_axis_name=expert_tensor_parallel_axis_name, + expert_fsdp_axis_name=expert_fsdp_axis_name, + expert_fsdp_shard_dim=expert_fsdp_shard_dim, compute_dtype=compute_dtype, ) diff --git a/autoparallel/shardings/placement_options.py b/autoparallel/shardings/placement_options.py index e2d3496a..4e220e2f 100644 --- a/autoparallel/shardings/placement_options.py +++ b/autoparallel/shardings/placement_options.py @@ -397,8 +397,11 @@ def get_local_map_placement_option( mesh, None, ), "Not yet implemented" - assert "call_local_map" in str(node.target) or "call_local_map_backward" in str( - node.target + target_name = str(node.target) + assert ( + "call_local_map" in target_name + or "call_local_map_backward" in target_name + or "local_map_hop" in target_name ) in_specs = [] num_activation_inputs = len(user_args) - len(in_placements) diff --git a/autoparallel/shardings/propagation_rules.py b/autoparallel/shardings/propagation_rules.py index 0c5a9849..28814125 100644 --- a/autoparallel/shardings/propagation_rules.py +++ b/autoparallel/shardings/propagation_rules.py @@ -286,6 +286,51 @@ def view_rule(mesh, op_schema): return OpStrategy(strats) +@register_rule(torch.ops.aten.expand.default) +def expand_rule(mesh, op_schema): + op_spec = op_schema.args_schema[0] + shape = op_schema.args_schema[1] + strats = [] + dim_map = dim_maps[torch.Tensor.expand] + rules = dim_map(op_spec, shape) + global_shape = op_spec.shape + in_tensor = _build_meta_tensor(op_spec.strategies[0].output_specs.tensor_meta) + out_tensor = torch.ops.aten.expand.default(in_tensor, shape) + out_tensor_meta = _gen_tensor_meta(out_tensor) + for strat in op_spec.strategies: + input_specs = strat.output_specs + try: + input_tgt_placements, output_placements = propagate_shape_and_sharding( + input_specs.placements, + global_shape, + rules, + mesh.shape, + strict_view=False, + ) + except (AssertionError, KeyError): + continue + + input_tgt_spec = DTensorSpec( + placements=tuple(input_tgt_placements), + mesh=mesh, + tensor_meta=input_specs.tensor_meta, + ) + output_spec = DTensorSpec( + mesh=mesh, + placements=tuple(output_placements), + tensor_meta=out_tensor_meta, + ) + + redistribute_costs = [generate_redistribute_costs(op_spec, input_tgt_spec)] + s = OpSpec( + output_spec, + input_specs=(input_tgt_spec,), + redistribute_cost=redistribute_costs, + ) + strats.append(s) + return OpStrategy(strats) + + @register_rule(torch.ops.aten.view.dtype) def view_dtype_rule(mesh, op_schema): """view(dtype=...) reinterprets the last dim's bytes as a different dtype. diff --git a/docs/local_map_and_moe.md b/docs/local_map_and_moe.md index 80ad0e5f..a54f05bd 100644 --- a/docs/local_map_and_moe.md +++ b/docs/local_map_and_moe.md @@ -285,6 +285,130 @@ The MoE dispatch is a fixed cost in the optimization: AutoParallel doesn't try to optimize the communication inside `local_map`, but it does account for the redistribution cost to get inputs into the required placements. +## User-facing MoE axis roles + +The MoE `local_map` helper is not an arbitrary placement DSL. It represents one +constrained family of EP boundaries: + +- token-like tensors shard on a configured set of token axes, +- expert weights shard on one EP axis, +- optional expert FSDP shards expert weights on one EFSDP storage axis at the + HOP boundary and all-gathers them inside the local body, +- optional expert tensor parallelism shards expert MLP dimensions on one ETP + axis, +- every other axis is replicated for expert weights. + +Axis names are lookup handles into the `DeviceMesh`. The implementation does +not infer semantics from names like `"dp"`, `"tp"`, `"cp"`, `"efsdp"`, or +`"dp_replicate"`. Roles are assigned by the MoE arguments: + +```python +MoE( + ..., + mesh=mesh, + ep_axis_name="ep", + token_axis_names=None, + expert_tensor_parallel_axis_name=None, + expert_fsdp_axis_name=None, + expert_fsdp_shard_dim=None, +) +``` + +The defaults mean: + +- `ep_axis_name="ep"`: the mesh axis named `"ep"` owns experts and runs token + dispatch/combine. +- `token_axis_names=None`: all mesh axes except the ETP axis are token axes. +- `expert_tensor_parallel_axis_name=None`: ETP is disabled. +- `expert_fsdp_axis_name=None`: EFSDP storage sharding is disabled. Axis names + such as `"efsdp"` do not enable it automatically. +- `expert_fsdp_shard_dim=None`: when EFSDP is enabled, pick TorchTitan's + storage shard dim rule: shard expert dim `0` unless `efsdp * ep` is larger + than the number of experts, then shard dim `1`. + +For the DSv3 example mesh: + +```python +mesh = init_device_mesh("cuda", (dp, ep), mesh_dim_names=("dp", "ep")) +DeepSeekV3Model(config, mesh=mesh) +``` + +the defaults create the 2D boundary `token=(Shard(0), Shard(0))` and expert +weights `(Replicate(), Shard(0))`. + +For a TorchTitan sparse mesh view: + +```python +mesh = parallel_dims.get_mesh(["efsdp", "ep"]) +DeepSeekV3Model(config, mesh=mesh, expert_fsdp_axis_name="efsdp") +``` + +the defaults still make both `"efsdp"` and `"ep"` token axes, while +`expert_fsdp_axis_name="efsdp"` makes expert weights shard on `"efsdp"` and +`"ep"` at the boundary. The local body all-gathers expert weights on `"efsdp"` +before grouped-mm; autograd lowers the matching backward path to +reduce-scatter. + +For explicit ETP, the user must keep ETP out of the token axes: + +```python +MoE( + ..., + mesh=mesh, # e.g. ("dp", "tp", "ep", "etp") + ep_axis_name="ep", + token_axis_names=("dp", "tp", "ep"), + expert_tensor_parallel_axis_name="etp", +) +``` + +This gives `w1`/`w3` `Shard(1)` on ETP, `w2` `Shard(2)` on ETP, and an +ETP all-reduce after local expert compute. + +Current `DeepSeekV3Model` wiring forwards these role bindings into `MoE` as +constructor arguments. They are not config fields. + +## EP `local_map` variants on 3+D meshes + +Higher-rank meshes do not require AutoParallel to understand MoE. They only +change the `local_map` boundary placements and the local implementation inside +the HOP. AutoParallel should continue to see a normal `local_map` node with +one placement tuple per tensor input/output. + +Axis names are handles for process groups and placement tuple order; they are +not semantic by themselves. The MoE code assigns roles to axes: + +| Role | Boundary meaning | +|---|---| +| Token axes | Activations, router outputs, and combined MoE output use `Shard(0)` after flattening `(batch, seq)` to tokens. Token-count outputs use `Partial(sum)`. | +| EP axis | Expert ownership and token dispatch/combine axis. Expert weights use `Shard(0)` on this axis. | +| EFSDP axis | Optional expert storage axis. Expert weights use `Shard(0)` or `Shard(1)` at the HOP boundary and are all-gathered inside the local body. | +| ETP axis | Optional expert-compute tensor-parallel axis. `w1`/`w3` shard hidden dim, `w2` shards the matching input dim, and the local body reduces the result on this axis. | +| Replica/storage axes | Expert weights are usually `Replicate()` at the HOP boundary unless that axis is also assigned a token or ETP role for a specific tensor. | +| PP axis | Outside current AutoParallel MoE `local_map` scope; pass a stage-local mesh without PP. | + +Useful variants include: + +| Mesh view | Intended MoE boundary | +|---|---| +| `(dp, ep)` | 2D compatibility: token tensors shard on both axes; expert weights shard only on EP. | +| `(dp_replicate, efsdp, ep)` | Sparse EP view from TorchTitan-style parallel dims. Token tensors can shard on the data/storage-derived axes and EP; pass `expert_fsdp_axis_name="efsdp"` to shard expert weights on EFSDP and EP. | +| `(dp, tp, ep)` or `(dp, cp, ep)` | Logical dense + EP view. Dense token axes and EP are all exposed in the HOP placement tuple. | +| `(dp, cp, tp, ep)` | Full logical dense + EP view. After flattening, DP/CP/SP-style TP token sharding is represented as token-dim `Shard(0)`. | +| `(dp_replicate, dp_shard, cp, tp, ep)` | Storage-visible dense + EP view for backends that keep DP storage axes explicit. | +| `(..., ep, etp)` | EP + explicit expert tensor parallelism. ETP is enabled by role assignment, not by assuming any particular axis name. | + +EFSDP is different from TorchTitan's `expert_param_placement_sparse()` boundary. +TorchTitan can write EFSDP as replicated at the local_map boundary because +`fully_shard()` reshards expert parameters outside local_map. AutoParallel only +gets HOP-level redistribute for this DSv3 path, so EFSDP storage sharding must +be visible in the HOP input placements and the local body must contain the +all-gather/reduce-scatter pair. + +The implementation rule is: no MoE-specific branch should be needed in +AutoParallel's optimizer, sharding propagation, or lowering. If a new variant +is needed, update the MoE boundary placement construction and the code inside +the `local_map` region. + ## Helpers in `autoparallel.collectives` AutoParallel provides collective wrappers that work inside `local_map` diff --git a/tests/test_moe_local_map_variants.py b/tests/test_moe_local_map_variants.py new file mode 100644 index 00000000..718c6c0d --- /dev/null +++ b/tests/test_moe_local_map_variants.py @@ -0,0 +1,204 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. +# +# This source code is licensed under the BSD license found in the +# LICENSE file in the root directory of this source tree. + +import pytest +import torch +from torch.distributed._tensor.placement_types import DTensorSpec, TensorMeta +from torch.distributed.tensor._op_schema import OpSpec, OpStrategy +from torch.distributed.tensor.placement_types import Partial, Replicate, Shard + +from autoparallel._testing.models.dsv3 import _moe_local_map_placements +from autoparallel.shardings.placement_options import get_placement_options +from conftest import apply_cuda_patches + + +DS3_EXAMPLE_WORLD_SIZE = 256 +DS3_EXAMPLE_EP_DEGREE = 64 +DS3_EXAMPLE_DP_DEGREE = DS3_EXAMPLE_WORLD_SIZE // DS3_EXAMPLE_EP_DEGREE +DS3_EXAMPLE_NUM_EXPERTS = 64 + + +def _mesh(shape, names): + return torch.distributed.device_mesh.init_device_mesh( + "cuda", + shape, + mesh_dim_names=names, + ) + + +@apply_cuda_patches +def test_ds3_attention_expand_skips_invalid_expanded_dim_shard_on_3d_mesh(): + mesh = _mesh( + (DS3_EXAMPLE_DP_DEGREE, 1, DS3_EXAMPLE_EP_DEGREE), + ("dp", "efsdp", "ep"), + ) + x = torch.empty((512, 128, 1, 64), dtype=torch.bfloat16, device="meta") + tensor_meta = TensorMeta(x.shape, x.stride(), x.dtype) + input_strategy = OpStrategy( + [ + OpSpec( + DTensorSpec(mesh, placements, tensor_meta=tensor_meta), + input_specs=(), + redistribute_cost=[], + ) + for placements in ( + (Shard(0), Shard(0), Shard(2)), + (Shard(0), Shard(0), Replicate()), + ) + ] + ) + + result = get_placement_options( + mesh, + torch.ops.aten.expand.default, + (input_strategy, [-1, -1, 16, -1]), + (x, [-1, -1, 16, -1]), + {}, + ) + + input_placements = {s.input_specs[0].placements for s in result.strategies} + assert (Shard(0), Shard(0), Shard(2)) not in input_placements + assert (Shard(0), Shard(0), Replicate()) in input_placements + + +@apply_cuda_patches +def test_ds3_example_2d_ep_boundary_matches_local_map_contract(): + mesh = _mesh((DS3_EXAMPLE_DP_DEGREE, DS3_EXAMPLE_EP_DEGREE), ("dp", "ep")) + + token, w1, w3, w2, counts = _moe_local_map_placements(mesh, "ep", None, None) + + assert token == (Shard(0), Shard(0)) + assert w1 == w3 == w2 == (Replicate(), Shard(0)) + assert counts == (Partial(reduce_op="sum"), Partial(reduce_op="sum")) + + +@apply_cuda_patches +def test_torchtitan_sparse_mesh_boundary_matches_parallel_dims(): + mesh = _mesh( + (DS3_EXAMPLE_DP_DEGREE, 1, DS3_EXAMPLE_EP_DEGREE), + ("dp_replicate", "efsdp", "ep"), + ) + + token, w1, w3, w2, counts = _moe_local_map_placements( + mesh, + "ep", + None, + None, + expert_fsdp_axis_name="efsdp", + num_experts=DS3_EXAMPLE_NUM_EXPERTS, + ) + + assert token == (Shard(0), Shard(0), Shard(0)) + assert w1 == w3 == w2 == (Replicate(), Shard(0), Shard(0)) + assert counts == ( + Partial(reduce_op="sum"), + Partial(reduce_op="sum"), + Partial(reduce_op="sum"), + ) + + +@apply_cuda_patches +def test_torchtitan_sparse_mesh_efsdp_fallback_matches_fsdp_shard_dim(): + mesh = _mesh( + (1, 4, DS3_EXAMPLE_EP_DEGREE), + ("dp_replicate", "efsdp", "ep"), + ) + + token, w1, w3, w2, counts = _moe_local_map_placements( + mesh, + "ep", + None, + None, + expert_fsdp_axis_name="efsdp", + num_experts=DS3_EXAMPLE_NUM_EXPERTS, + ) + + assert token == (Shard(0), Shard(0), Shard(0)) + assert w1 == w3 == w2 == (Replicate(), Shard(1), Shard(0)) + assert counts == ( + Partial(reduce_op="sum"), + Partial(reduce_op="sum"), + Partial(reduce_op="sum"), + ) + + +@apply_cuda_patches +def test_torchtitan_sparse_mesh_etp_boundary_is_future_gated_by_axis_role(): + mesh = _mesh( + (DS3_EXAMPLE_DP_DEGREE, 1, DS3_EXAMPLE_EP_DEGREE, 1), + ("dp_replicate", "efsdp", "ep", "etp"), + ) + + token, w1, w3, w2, counts = _moe_local_map_placements( + mesh, + "ep", + None, + "etp", + expert_fsdp_axis_name="efsdp", + num_experts=DS3_EXAMPLE_NUM_EXPERTS, + ) + + assert token == (Shard(0), Shard(0), Shard(0), Replicate()) + assert w1 == w3 == (Replicate(), Shard(0), Shard(0), Shard(1)) + assert w2 == (Replicate(), Shard(0), Shard(0), Shard(2)) + assert counts == ( + Partial(reduce_op="sum"), + Partial(reduce_op="sum"), + Partial(reduce_op="sum"), + Replicate(), + ) + + +@apply_cuda_patches +def test_torchtitan_full_dtensor_dense_ep_boundary_matches_parallel_dims(): + mesh = _mesh( + (DS3_EXAMPLE_DP_DEGREE, 1, 1, 1, DS3_EXAMPLE_EP_DEGREE), + ("dp_replicate", "dp_shard", "cp", "tp", "ep"), + ) + + token, w1, w3, w2, counts = _moe_local_map_placements(mesh, "ep", None, None) + + assert token == (Shard(0), Shard(0), Shard(0), Shard(0), Shard(0)) + assert w1 == w3 == w2 == ( + Replicate(), + Replicate(), + Replicate(), + Replicate(), + Shard(0), + ) + assert counts == ( + Partial(reduce_op="sum"), + Partial(reduce_op="sum"), + Partial(reduce_op="sum"), + Partial(reduce_op="sum"), + Partial(reduce_op="sum"), + ) + + +@apply_cuda_patches +def test_moe_local_map_rejects_overlapping_ep_and_etp_axes(): + mesh = _mesh((DS3_EXAMPLE_DP_DEGREE, DS3_EXAMPLE_EP_DEGREE), ("dp", "ep")) + + with pytest.raises(ValueError, match="must differ"): + _moe_local_map_placements(mesh, "ep", None, "ep") + + +@apply_cuda_patches +def test_moe_local_map_rejects_efsdp_hidden_dim_shard_with_etp(): + mesh = _mesh( + (DS3_EXAMPLE_DP_DEGREE, 1, DS3_EXAMPLE_EP_DEGREE, 1), + ("dp_replicate", "efsdp", "ep", "etp"), + ) + + with pytest.raises(ValueError, match="Shard\\(1\\).*ETP"): + _moe_local_map_placements( + mesh, + "ep", + None, + "etp", + expert_fsdp_axis_name="efsdp", + expert_fsdp_shard_dim=1, + num_experts=DS3_EXAMPLE_NUM_EXPERTS, + )