[spmd_types] fix spmd->DTensor translation on partial mesh

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.

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

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:

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)