diff --git a/tests/unit_tests/test_parallel_dims.py b/tests/unit_tests/test_parallel_dims.py index 68281e37be..575599c574 100644 --- a/tests/unit_tests/test_parallel_dims.py +++ b/tests/unit_tests/test_parallel_dims.py @@ -12,7 +12,7 @@ import torch import torch.distributed as dist from torch.distributed.device_mesh import init_device_mesh -from torch.distributed.tensor import Shard +from torch.distributed.tensor import Replicate from torch.distributed.tensor.debug import CommDebugMode from torch.testing._internal.distributed._tensor.common_dtensor import ( DTensorTestBase, @@ -28,7 +28,6 @@ ) from torchtitan.distributed.spmd_types import ( spmd_distribute_tensor, - spmd_layout_to_dtensor_placements, spmd_redistribute_per_axis, spmd_validate_redistributions, ) @@ -37,7 +36,7 @@ dense_sequence_parallel_placement, ) from torchtitan.models.llama3 import model_registry -from torchtitan.protocols.sharding import ShardingConfig +from torchtitan.protocols.sharding import resolve_placements, ShardingConfig class TestParallelDimsValidation(unittest.TestCase): @@ -236,19 +235,6 @@ class TestSpmdLayout(DTensorTestBase): def world_size(self): return 4 - def test_converts_partition_spec_to_dtensor_shard(self): - """PartitionSpec refines V into concrete DTensor Shard placement.""" - layout = SpmdLayout( - {MeshAxisName.TP: spmd.V}, - partition_spec=spmd.PartitionSpec(MeshAxisName.TP), - ) - - self.assertEqual(layout.per_axis_spmd_types(), {MeshAxisName.TP: spmd.S(0)}) - self.assertEqual( - spmd_layout_to_dtensor_placements(layout), - {MeshAxisName.TP: Shard(0)}, - ) - def test_seq_parallel_activation_per_axis_spmd_types(self): """PartitionSpec can map multiple mesh axes to one tensor dim.""" layout = SpmdLayout( @@ -280,6 +266,21 @@ def test_unfold_dp_axes(self): ["dp_replicate", "dp_shard", "cp", "tp"], ) + @with_comms + def test_resolve_placements_ignores_extra_untranslatable_axes(self): + """Extra layout axes are ignored before converting to DTensor placements.""" + mesh = init_device_mesh( + self.device_type, (self.world_size,), mesh_dim_names=("tp",) + ) + layout = SpmdLayout( + { + MeshAxisName.DP: spmd.V, + MeshAxisName.TP: spmd.I, + } + ) + + self.assertEqual(resolve_placements(layout, mesh), (Replicate(),)) + def test_rejects_partition_spec_reorder_redistribute(self): """((DP, CP), None) -> ((CP, DP), None) not supported by a single redistribute call.""" with self.assertRaises(ValueError) as cm: @@ -332,6 +333,31 @@ def test_rejects_multi_axis_redistribute(self): ) ) + def test_rejects_redistribute_from_varying(self): + for src_dp, dst_dp in ((spmd.V, spmd.R), (spmd.R, spmd.V)): + with self.subTest(src_dp=src_dp, dst_dp=dst_dp): + with self.assertRaisesRegex( + ValueError, + "output: SpmdLayout-based redistribution changes mesh axis " + "'dp' with spmd.V as the source or destination type", + ): + spmd_validate_redistributions( + ShardingConfig( + out_src_shardings=SpmdLayout( + { + MeshAxisName.DP: src_dp, + MeshAxisName.TP: spmd.I, + } + ), + out_dst_shardings=SpmdLayout( + { + MeshAxisName.DP: dst_dp, + MeshAxisName.TP: spmd.I, + } + ), + ) + ) + @with_comms def test_partition_spec_order_controls_state_shard(self): """Test spmd_distribute_tensor follows PartitionSpec order. diff --git a/torchtitan/distributed/parallel_dims.py b/torchtitan/distributed/parallel_dims.py index a57720c388..3f134c2c62 100644 --- a/torchtitan/distributed/parallel_dims.py +++ b/torchtitan/distributed/parallel_dims.py @@ -19,7 +19,13 @@ from torchtitan.tools.utils import device_type -__all__ = ["MeshAxisName", "ParallelDims", "SpmdLayout", "unfold_dp_axes"] +__all__ = [ + "MeshAxisName", + "ParallelDims", + "SpmdLayout", + "unfold_dp_axis", + "unfold_dp_axes", +] class StrEnum(str, Enum): @@ -115,16 +121,19 @@ def per_axis_spmd_types(self) -> dict[MeshAxisName, spmd.PerMeshAxisSpmdType]: return result +def unfold_dp_axis(axis: MeshAxisName | str) -> tuple[MeshAxisName, ...]: + """Expand logical ``dp`` into concrete dense storage mesh axes.""" + axis_name = MeshAxisName(axis) + if axis_name == MeshAxisName.DP: + return (MeshAxisName.DP_REPLICATE, MeshAxisName.DP_SHARD) + return (axis_name,) + + def unfold_dp_axes(axes: Iterable[MeshAxisName | str]) -> list[str]: """Expand logical ``dp`` into concrete dense storage mesh axes.""" - result: list[str] = [] - for axis in axes: - axis_value = axis.value if isinstance(axis, MeshAxisName) else axis - if axis_value == "dp": - result.extend(("dp_replicate", "dp_shard")) - else: - result.append(axis_value) - return result + return [ + concrete_axis.value for axis in axes for concrete_axis in unfold_dp_axis(axis) + ] @dataclass diff --git a/torchtitan/distributed/spmd_types.py b/torchtitan/distributed/spmd_types.py index e0e2b50e62..d843c0e61c 100644 --- a/torchtitan/distributed/spmd_types.py +++ b/torchtitan/distributed/spmd_types.py @@ -16,7 +16,6 @@ import spmd_types as spmd import torch from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import Partial, Placement, Replicate, Shard from torchtitan.distributed.utils import get_spmd_backend @@ -40,7 +39,6 @@ "set_current_spmd_mesh", "set_spmd_meshes", "maybe_set_sparse_mesh", - "spmd_layout_to_dtensor_placements", ] @@ -134,30 +132,6 @@ def maybe_set_sparse_mesh() -> Iterator[None]: yield -def spmd_layout_to_dtensor_placements( - layout: "SpmdLayout", -) -> dict["MeshAxisName", Placement]: - """Convert an SPMD layout to DTensor placements keyed by mesh axis name.""" - from torchtitan.distributed.parallel_dims import MeshAxisName - - result: dict[MeshAxisName, Placement] = {} - for axis_name, axis_type in layout.per_axis_spmd_types().items(): - if axis_type == spmd.R or axis_type == spmd.I: - dtensor_placement: Placement = Replicate() - elif axis_type == spmd.P: - dtensor_placement = Partial() - else: - assert isinstance(axis_type, spmd.Shard) - dtensor_placement = Shard(axis_type.dim) - - if axis_name == MeshAxisName.DP: - result[MeshAxisName.DP_REPLICATE] = dtensor_placement - result[MeshAxisName.DP_SHARD] = dtensor_placement - else: - result[axis_name] = dtensor_placement - return result - - def annotate_input_spmd_types( parallel_dims: "ParallelDims", inputs: torch.Tensor, @@ -259,6 +233,16 @@ def _validate_redistribute_spmd_pair( "spmd_redistribute_per_axis only supports one single-axis " "redistribution." ) + if changed_axes and ( + src_types[changed_axes[0]] is spmd.V or dst_types[changed_axes[0]] is spmd.V + ): + axis = changed_axes[0] + raise ValueError( + f"{name}: SpmdLayout-based redistribution changes mesh axis " + f"{axis.value!r} with spmd.V as the source or destination type. " + "Config-based redistribution requires non-V types; write an " + "explicit collective when the value semantics are unclear." + ) # 2) If neither has PartitionSpec, comparing per_axis_spmd_types() is sufficient. if src.partition_spec is None and dst.partition_spec is None: diff --git a/torchtitan/protocols/sharding.py b/torchtitan/protocols/sharding.py index 118d43d964..4f3af43cfd 100644 --- a/torchtitan/protocols/sharding.py +++ b/torchtitan/protocols/sharding.py @@ -14,12 +14,15 @@ from dataclasses import dataclass, field +import spmd_types as spmd from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import Partial, Placement, Replicate, Shard -from torchtitan.distributed.parallel_dims import MeshAxisName, SpmdLayout - -from torchtitan.distributed.spmd_types import spmd_layout_to_dtensor_placements +from torchtitan.distributed.parallel_dims import ( + MeshAxisName, + SpmdLayout, + unfold_dp_axis, +) __all__ = [ @@ -143,18 +146,22 @@ def resolve_placements( # TODO(fegin): remove the size-1 ``Shard(d)``/``Partial`` to ``Replicate()`` # conversion once FlexShard replaces ``fully_shard``. assert mesh.mesh_dim_names is not None, "DeviceMesh must have named axes" - placements = spmd_layout_to_dtensor_placements(layout) + axis_types = {} + for axis_name, axis_type in layout.per_axis_spmd_types().items(): + for concrete_axis_name in unfold_dp_axis(axis_name): + axis_types[concrete_axis_name] = axis_type + result = [] for i, axis_name in enumerate(mesh.mesh_dim_names): key = MeshAxisName(axis_name) - if key not in placements: + if key not in axis_types: raise ValueError( f"ShardingConfig does not declare a placement for mesh axis " f"{axis_name!r}. Declared: " f"{sorted(k.value for k in layout.axes())}; " f"required: {list(mesh.mesh_dim_names)}." ) - p = placements[key] + p = spmd.spmd_type_to_dtensor_placement(axis_types[key]) if isinstance(p, (Shard, Partial)) and mesh.size(i) == 1: p = Replicate() result.append(p)