From 57428ff170fe4a78660dbb827266d1a9782e19fa Mon Sep 17 00:00:00 2001
From: Pian Pawakapan <pianpwk@meta.com>
Date: Fri, 21 Aug 2026 11:21:58 -0700
Subject: [PATCH] Update

[ghstack-poisoned]
---
 .ci/docker/requirements.txt                  |  2 +-
 pyproject.toml                               |  2 +-
 tests/unit_tests/test_parallel_dims.py       |  4 +-
 torchtitan/components/loss.py                | 63 +++++++----------
 torchtitan/components/quantization/nvfp4.py  |  4 +-
 torchtitan/distributed/linear.py             | 71 +++++++++----------
 torchtitan/distributed/parallel_dims.py      | 73 +++-----------------
 torchtitan/distributed/spmd_types.py         | 51 +++++++++-----
 torchtitan/models/common/attention.py        | 65 +++++++++--------
 torchtitan/models/common/decoder_sharding.py |  2 +-
 torchtitan/models/common/token_dispatcher.py | 26 +++----
 torchtitan/models/deepseek_v3/mtp.py         | 15 +---
 torchtitan/models/gpt_oss/moe.py             | 12 ++--
 torchtitan/overrides/fused_mla.py            | 27 +++-----
 torchtitan/protocols/module.py               | 66 ++++++------------
 torchtitan/protocols/sharding.py             | 27 ++++----
 16 files changed, 204 insertions(+), 306 deletions(-)

diff --git a/.ci/docker/requirements.txt b/.ci/docker/requirements.txt
index 4a6b759f0a..1adba187dd 100644
--- a/.ci/docker/requirements.txt
+++ b/.ci/docker/requirements.txt
@@ -7,4 +7,4 @@ tokenizers >= 0.15.0
 safetensors
 einops
 pillow
-spmd_types==0.2.3
+spmd_types==0.2.4
diff --git a/pyproject.toml b/pyproject.toml
index 2b63009865..2935d9913d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,7 +27,7 @@ dependencies = [
     "wandb",
     "einops",
     "pillow",
-    "spmd_types==0.2.3",
+    "spmd_types==0.2.4",
 ]
 dynamic = ["version"]
 
diff --git a/tests/unit_tests/test_parallel_dims.py b/tests/unit_tests/test_parallel_dims.py
index f0190b5ef3..0ddc0dbdad 100644
--- a/tests/unit_tests/test_parallel_dims.py
+++ b/tests/unit_tests/test_parallel_dims.py
@@ -443,8 +443,8 @@ def test_spmd_redistribute_per_axis_allgather(self):
             result = spmd_redistribute_per_axis(
                 x,
                 mesh,
-                src.per_axis_spmd_types(),
-                dst.per_axis_spmd_types(),
+                src,
+                dst,
             )
 
         self.assertEqual(comm_mode.get_total_counts(), 1)
diff --git a/torchtitan/components/loss.py b/torchtitan/components/loss.py
index bcc0fbe907..10f178c576 100644
--- a/torchtitan/components/loss.py
+++ b/torchtitan/components/loss.py
@@ -62,7 +62,6 @@ def cross_entropy_loss(
     )
 
 
-@spmd.register_autograd_function
 class _LossParallelCrossEntropy(torch.autograd.Function):
     """
     Vocab-parallel cross-entropy on local ``[T, V_local]`` logits.
@@ -80,30 +79,22 @@ class _LossParallelCrossEntropy(torch.autograd.Function):
     """
 
     @staticmethod
-    def typecheck_forward(
+    def spmd_typecheck(
+        result: torch.Tensor,
+        *,
         logits: torch.Tensor,
         labels: torch.Tensor,
         tp_group: dist.ProcessGroup,
-        global_vocab_size: int,
-        reduction: str = "sum",
-    ) -> torch.Tensor:
+    ) -> None:
         """
         SPMD type: logits S(-1)@TP, labels I@TP -> loss I@TP.
         Non-TP axes are passed through from logits to the output.
         """
         spmd.assert_type(logits, {tp_group: spmd.S(logits.dim() - 1)})
         spmd.assert_type(labels, {tp_group: spmd.I})
-        result = _LossParallelCrossEntropy.apply(
-            logits,
-            labels,
-            tp_group,
-            global_vocab_size,
-            reduction,
-        )
         output_type = dict(spmd.get_local_type(logits))
-        output_type[tp_group] = spmd.I
+        output_type[spmd.MeshAxis.of(tp_group)] = spmd.I
         spmd.assert_type(result, output_type)
-        return result
 
     @staticmethod
     # pyrefly: ignore [bad-override]
@@ -261,13 +252,14 @@ def __call__(
         """Return the scaled loss and any metrics computed by the loss."""
         del kwargs
         loss = self.fn(pred, labels)
+        if get_spmd_backend() == "spmd_types":  # loss: V->P, annotate global_valid_tokens
+            spmd.assert_type(loss, {"dp": spmd.P, "cp": spmd.P})
+            spmd.assert_type(
+                global_valid_tokens,
+                {"dp": spmd.R, "cp": spmd.R, "tp": spmd.I},
+            )
         if global_valid_tokens is not None:
-            # TODO(pianpwk): Teach spmd_types that P / scalar preserves P.
-            is_type_checking = spmd.is_type_checking()
-            with spmd.no_typecheck():
-                loss = loss / global_valid_tokens
-                if is_type_checking:
-                    spmd.assert_type(loss, {"dp": spmd.P, "cp": spmd.P, "tp": spmd.I})
+            loss = loss / global_valid_tokens
         return loss, {}
 
 
@@ -293,13 +285,14 @@ def __call__(
     ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
         del kwargs
         loss = self.fn(pred, labels, global_vocab_size=self.global_vocab_size)
+        if get_spmd_backend() == "spmd_types":  # loss: V->P, annotate global_valid_tokens
+            spmd.assert_type(loss, {"dp": spmd.P, "cp": spmd.P})
+            spmd.assert_type(
+                global_valid_tokens,
+                {"dp": spmd.R, "cp": spmd.R, "tp": spmd.I},
+            )
         if global_valid_tokens is not None:
-            # TODO(pianpwk): Teach spmd_types that P / scalar preserves P.
-            is_type_checking = spmd.is_type_checking()
-            with spmd.no_typecheck():
-                loss = loss / global_valid_tokens
-                if is_type_checking:
-                    spmd.assert_type(loss, {"dp": spmd.P, "cp": spmd.P, "tp": spmd.I})
+            loss = loss / global_valid_tokens
         return loss, {}
 
 
@@ -350,15 +343,13 @@ def compute_logprobs(
     elif get_spmd_backend() == "spmd_types" and spmd_mesh_size("tp") > 1:
         # spmd_types returns a plain local vocab shard. Labels are global token
         # ids, so cross_entropy needs full-vocab logits.
-        mesh = current_spmd_mesh()
-        assert mesh is not None
         # dst=I, not R: the vocab all-gather's grad is the replicated upstream
         # grad sliced back to this rank's vocab shard (I's backward), not an
         # all-reduce (R's backward). The latter over-counts by tp_degree and
         # diverges from the DTensor path above, whose redistribute grad slices.
         logits = spmd.redistribute(
             logits,
-            mesh.get_group("tp"),
+            "tp",
             src=spmd.S(-1),
             dst=spmd.I,
         )
@@ -644,15 +635,11 @@ def _chunk(t):
 
             total_loss = hidden_states.new_zeros((), dtype=torch.float32)
             if get_spmd_backend() == "spmd_types" and spmd.is_type_checking():
-                # TODO(pianpwk): would be nice if mutate_type accepted multiple axes.
-                for axis_name, dst in {
-                    "dp": spmd.P,
-                    "cp": spmd.P,
-                    "tp": spmd.I,
-                }.items():
-                    total_loss = spmd.mutate_type(
-                        total_loss, axis_name, src=spmd.R, dst=dst
-                    )
+                total_loss = spmd.mutate_type(
+                    total_loss,
+                    src=spmd.R,
+                    dst={"dp": spmd.P, "cp": spmd.P, "tp": spmd.I},
+                )
             metrics: dict[str, torch.Tensor] = {}
 
             # Disable FSDP reshard on lm_head to keep weight unsharded across
diff --git a/torchtitan/components/quantization/nvfp4.py b/torchtitan/components/quantization/nvfp4.py
index b830a5a366..2b6278f212 100644
--- a/torchtitan/components/quantization/nvfp4.py
+++ b/torchtitan/components/quantization/nvfp4.py
@@ -119,9 +119,7 @@ def build(self, **kwargs):
                 instance = Linear.Config.build(self, **kwargs)
                 if instance._sharding_config is not None:
                     sc = instance._sharding_config
-                    weight_tp = (
-                        sc.state_shardings["weight"].per_axis_spmd_types().get(TP)
-                    )
+                    weight_tp = sc.state_shardings["weight"].local_type.get(TP)
                     rowwise = isinstance(weight_tp, spmd.Shard) and weight_tp.dim == 1
                     if rowwise:
                         in_layout = dense_activation_placement(
diff --git a/torchtitan/distributed/linear.py b/torchtitan/distributed/linear.py
index e56595c547..23d42f7621 100644
--- a/torchtitan/distributed/linear.py
+++ b/torchtitan/distributed/linear.py
@@ -44,7 +44,6 @@ def ensure_symm_mem_ops():
     return symm_mem
 
 
-@spmd.register_autograd_function
 class AllGatherLinear(torch.autograd.Function):
     """All-gather the sequence shard, then apply a column-parallel linear.
 
@@ -80,32 +79,32 @@ class AllGatherLinear(torch.autograd.Function):
     """
 
     @staticmethod
-    def typecheck_forward(
+    def spmd_typecheck(
+        result: torch.Tensor,
+        *,
         x_shard_m: torch.Tensor,
         w_shard_n: torch.Tensor,
         bias_shard_n: torch.Tensor | None,
-        group: dist.ProcessGroup,
         group_name: str,
-    ) -> torch.Tensor:
+    ) -> None:
         """SPMD type: x S(0)@TP, w S(0)@TP -> y S(1)@TP.
 
         The gather consumes the row shard, so the result is full on rows; the
         weight's output-feature shard survives the GEMM. Non-TP axes pass through
         from x.
         """
-        spmd.assert_type(x_shard_m, {group: spmd.S(0)})
+        spmd.assert_type(x_shard_m, {group_name: spmd.S(0)})
         # S(0), not S(1), even though this is the column-parallel direction: torch
         # stores the weight as [N, K] while the mental model of the GEMM is [K, N],
         # so sharding the output features N is dim 0 of what is actually stored.
-        spmd.assert_type(w_shard_n, {group: spmd.S(0)})
+        spmd.assert_type(w_shard_n, {group_name: spmd.S(0)})
         if bias_shard_n is not None:
-            spmd.assert_type(bias_shard_n, {group: spmd.S(0)})
-        result = AllGatherLinear.apply(
-            x_shard_m, w_shard_n, bias_shard_n, group, group_name
+            spmd.assert_type(bias_shard_n, {group_name: spmd.S(0)})
+        spmd.assert_local_type_like(
+            result,
+            x_shard_m,
+            {group_name: spmd.S(1)},  # pyrefly: ignore [bad-argument-type]
         )
-        output_type = {**spmd.get_local_type(x_shard_m), group: spmd.S(1)}
-        spmd.assert_type(result, output_type)
-        return result
 
     @staticmethod
     def forward(  # pyrefly: ignore[bad-override]
@@ -183,7 +182,6 @@ def backward(ctx, grad_y_shard_n: torch.Tensor):  # pyrefly: ignore[bad-override
         return grad_x_shard_m, grad_w_shard_n, grad_bias, None, None
 
 
-@spmd.register_autograd_function
 class AllGatherLinearMulti(torch.autograd.Function):
     """One all-gather feeding a pair of column-parallel linears on the same input.
 
@@ -282,29 +280,28 @@ def backward(  # pyrefly: ignore[bad-override]
         )
 
     @staticmethod
-    def typecheck_forward(
+    def spmd_typecheck(
+        results: tuple[torch.Tensor, torch.Tensor],
+        *,
         x_shard_m: torch.Tensor,
         wa_shard_n: torch.Tensor,
         wb_shard_n: torch.Tensor,
-        group: dist.ProcessGroup,
         group_name: str,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> None:
         """SPMD type: x S(0)@TP, both w S(0)@TP -> both y S(1)@TP."""
-        spmd.assert_type(x_shard_m, {group: spmd.S(0)})
+        spmd.assert_type(x_shard_m, {group_name: spmd.S(0)})
         # S(0) for the column-parallel direction; see AllGatherLinear for why the
         # stored [N, K] layout inverts the dim you would expect.
-        spmd.assert_type(wa_shard_n, {group: spmd.S(0)})
-        spmd.assert_type(wb_shard_n, {group: spmd.S(0)})
-        results = AllGatherLinearMulti.apply(
-            x_shard_m, wa_shard_n, wb_shard_n, group, group_name
-        )
-        output_type = {**spmd.get_local_type(x_shard_m), group: spmd.S(1)}
+        spmd.assert_type(wa_shard_n, {group_name: spmd.S(0)})
+        spmd.assert_type(wb_shard_n, {group_name: spmd.S(0)})
         for result in results:
-            spmd.assert_type(result, output_type)
-        return results
+            spmd.assert_local_type_like(
+                result,
+                x_shard_m,
+                {group_name: spmd.S(1)},  # pyrefly: ignore [bad-argument-type]
+            )
 
 
-@spmd.register_autograd_function
 class LinearReduceScatter(torch.autograd.Function):
     """Apply a row-parallel linear, then reduce-scatter over the sequence.
 
@@ -329,32 +326,32 @@ class LinearReduceScatter(torch.autograd.Function):
     """
 
     @staticmethod
-    def typecheck_forward(
+    def spmd_typecheck(
+        result: torch.Tensor,
+        *,
         x_shard_k: torch.Tensor,
         w_shard_k: torch.Tensor,
         bias: torch.Tensor | None,
-        group: dist.ProcessGroup,
         group_name: str,
-    ) -> torch.Tensor:
+    ) -> None:
         """SPMD type: x S(1)@TP, w S(1)@TP, bias R@TP -> y S(0)@TP.
 
         The local matmul is a partial sum over the sharded K; the reduce-scatter
         completes it and shards rows instead. Non-TP axes pass through from x.
         """
-        spmd.assert_type(x_shard_k, {group: spmd.S(1)})
+        spmd.assert_type(x_shard_k, {group_name: spmd.S(1)})
         # S(1), the mirror of AllGatherLinear's S(0): torch stores the weight as
         # [N, K] while the mental model of the GEMM is [K, N], so sharding the
         # input features K -- the row-parallel direction -- is dim 1 of what is
         # actually stored.
-        spmd.assert_type(w_shard_k, {group: spmd.S(1)})
+        spmd.assert_type(w_shard_k, {group_name: spmd.S(1)})
         if bias is not None:
-            spmd.assert_type(bias, {group: spmd.R})
-        result = LinearReduceScatter.apply(
-            x_shard_k, w_shard_k, bias, group, group_name
+            spmd.assert_type(bias, {group_name: spmd.R})
+        spmd.assert_local_type_like(
+            result,
+            x_shard_k,
+            {group_name: spmd.S(0)},  # pyrefly: ignore [bad-argument-type]
         )
-        output_type = {**spmd.get_local_type(x_shard_k), group: spmd.S(0)}
-        spmd.assert_type(result, output_type)
-        return result
 
     @staticmethod
     def forward(  # pyrefly: ignore[bad-override]
diff --git a/torchtitan/distributed/parallel_dims.py b/torchtitan/distributed/parallel_dims.py
index fceb26f00f..ddc962df87 100644
--- a/torchtitan/distributed/parallel_dims.py
+++ b/torchtitan/distributed/parallel_dims.py
@@ -9,7 +9,7 @@
 from collections.abc import Iterable
 from dataclasses import dataclass, field
 from enum import StrEnum
-from typing import Any, Literal
+from typing import Literal
 
 import spmd_types as spmd
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
@@ -23,6 +23,7 @@
     "MeshAxisName",
     "ParallelDims",
     "SpmdLayout",
+    "layout_axes",
     "unfold_dp_axis",
     "unfold_dp_axes",
 ]
@@ -52,67 +53,15 @@ class MeshAxisName(StrEnum):
     EFSDP = "efsdp"
 
 
-@dataclass(frozen=True, slots=True)
-class SpmdLayout:
-    """Temporary SPMD layout annotations keyed by logical mesh axis name.
+SpmdLayout = spmd.SpmdType
 
-    TODO(pianpwk): Replace this with ``spmd_types.SpmdLayout`` once that API is
-    available in TorchTitan's minimum ``spmd_types`` version.
-    """
-
-    axis_types: dict[MeshAxisName, spmd.PerMeshAxisSpmdType]
-    partition_spec: spmd.PartitionSpec | tuple[Any, ...] | None = None
-
-    def __post_init__(self) -> None:
-        sharded_dims: dict[int, MeshAxisName] = {}
-        for axis_name, axis_type in self.axis_types.items():
-            if not isinstance(axis_type, spmd.Shard):
-                continue
-            if self.partition_spec is not None:
-                raise ValueError(
-                    "SpmdLayout with PartitionSpec should use spmd.V instead "
-                    "of spmd.S(dim) in per-axis-types, and express tensor dim "
-                    "sharding in the provided PartitionSpec."
-                )
-            if axis_type.dim in sharded_dims:
-                raise ValueError(
-                    "SpmdLayout has multiple mesh axes sharding tensor dim "
-                    f"{axis_type.dim}; provide partition_spec to make shard "
-                    "ordering explicit."
-                )
-            sharded_dims[axis_type.dim] = axis_name
-
-    def axes(self) -> tuple[MeshAxisName, ...]:
-        return tuple(self.axis_types)
-
-    def per_axis_spmd_types(self) -> dict[MeshAxisName, spmd.PerMeshAxisSpmdType]:
-        """
-        Return per-axis types with PartitionSpec sharding represented as S(i).
-        e.g. {DP: R, CP: V} + PartitionSpec(None, CP) -> {DP: R, CP: S(1)}
-
-        This is not meant as a minimal description of the SPMD layout; shard order
-        cannot be expressed. Specifically, shard order information will be lost in
-        this representation. This is purely a helper for calling spmd.redistribute,
-        which takes per-axis types (e.g. redistribute(S(1) -> R)).
 
-        This manually handles ``MeshAxisName``, because spmd_types normalization
-        functions often attempt to resolve to concrete runtime mesh axes, even
-        without a set current mesh.
-        """
-        result = dict(self.axis_types)
-        if self.partition_spec is not None:
-            for dim, entry in enumerate(self.partition_spec):
-                if entry is None:
-                    continue
-                axes = entry if isinstance(entry, tuple) else (entry,)
-                for axis_name in axes:
-                    if not isinstance(axis_name, MeshAxisName):
-                        raise TypeError(
-                            f"Expected MeshAxisName in partition_spec, "
-                            f"got {axis_name!r}."
-                        )
-                    result[axis_name] = spmd.S(dim)
-        return result
+def layout_axes(layout: SpmdLayout) -> tuple[MeshAxisName, ...]:
+    """Return and validate the named mesh axes used by a sharding config."""
+    return tuple(
+        MeshAxisName(axis)  # pyrefly: ignore [bad-argument-type]
+        for axis in layout.local_type
+    )
 
 
 def unfold_dp_axis(axis: MeshAxisName | str) -> tuple[MeshAxisName, ...]:
@@ -557,9 +506,9 @@ def resolve_shared_mesh(
         non_none = [p for p in placements if p is not None]
         if not non_none:
             return None
-        axes = non_none[0].axes()
+        axes = layout_axes(non_none[0])
         for p in non_none[1:]:
-            p_axes = p.axes()
+            p_axes = layout_axes(p)
             assert p_axes == axes, (
                 f"Inconsistent mesh axes within a boundary: "
                 f"{sorted(k.value for k in axes)} vs "
diff --git a/torchtitan/distributed/spmd_types.py b/torchtitan/distributed/spmd_types.py
index 919c32a9a3..f11bb4dc69 100644
--- a/torchtitan/distributed/spmd_types.py
+++ b/torchtitan/distributed/spmd_types.py
@@ -51,11 +51,11 @@
 def plain_tensor_to_dtensor_state_dict(
     state_dict: dict[str, Any],
     *,
-    state_dict_layouts: Mapping[str, "SpmdLayout"],
+    state_dict_layouts: Mapping[str, spmd.SpmdType],
     parallel_dims: "ParallelDims",
 ) -> dict[str, Any]:
     """Represent plain local state tensors as DTensors for state transfer."""
-    from torchtitan.distributed.parallel_dims import unfold_dp_axes
+    from torchtitan.distributed.parallel_dims import layout_axes, unfold_dp_axes
     from torchtitan.protocols.sharding import resolve_placements
 
     dtensor_state_dict = dict(state_dict)
@@ -68,7 +68,9 @@ def plain_tensor_to_dtensor_state_dict(
             if layout is None:
                 raise KeyError(f"{name} is missing SPMD layout metadata")
 
-            mesh = parallel_dims.get_activated_mesh(unfold_dp_axes(layout.axes()))
+            mesh = parallel_dims.get_activated_mesh(
+                unfold_dp_axes(layout_axes(layout))
+            )
             if mesh is None:
                 continue
 
@@ -220,6 +222,17 @@ def annotate_input_spmd_types(
     return inputs, labels, extra_kwargs
 
 
+def _per_axis_types(layout: "SpmdLayout") -> spmd.PerMeshAxisSpmdTypes:
+    result = dict(layout.local_type)
+    if layout.partition_spec is not None:
+        for dim, entry in enumerate(layout.partition_spec):
+            for axis in (
+                () if entry is None else entry if isinstance(entry, tuple) else (entry,)
+            ):
+                result[axis] = spmd.S(dim)
+    return result
+
+
 def spmd_validate_redistributions(sharding_config: Any) -> None:
     """Validate that SPMD redistributions fit the current runtime helper.
 
@@ -235,9 +248,10 @@ def spmd_validate_redistributions(sharding_config: Any) -> None:
     or we should write collective-based (not placement-based) redistributions
     once the partial_dtensor backend is removed.
     """
+    from torchtitan.distributed.parallel_dims import MeshAxisName
 
     def _normalize_partition_spec(
-        axis_types: dict["MeshAxisName", spmd.PerMeshAxisSpmdType],
+        axis_types: Mapping["MeshAxisName", spmd.PerMeshAxisSpmdType],
         *,
         ndim: int,
     ) -> tuple[tuple["MeshAxisName", ...], ...]:
@@ -246,13 +260,15 @@ def _normalize_partition_spec(
         for axis_name, axis_type in axis_types.items():
             if not isinstance(axis_type, spmd.Shard):
                 continue
+            if not isinstance(axis_name, str):
+                raise TypeError("ShardingConfig SpmdType axes must be names")
             dim = axis_type.dim if axis_type.dim >= 0 else ndim + axis_type.dim
             if dim < 0 or dim >= ndim:
                 raise ValueError(
                     f"Cannot compare SPMD layout with shard dim {axis_type.dim} "
                     f"against PartitionSpec of rank {ndim}."
                 )
-            entries[dim] = (axis_name,)
+            entries[dim] = (MeshAxisName(axis_name),)
         return tuple(entries)
 
     def _validate_redistribute_spmd_pair(
@@ -262,10 +278,10 @@ def _validate_redistribute_spmd_pair(
         name: str,
     ) -> None:
         """Validate a SPMD redistribution is expressible with one-axis collective."""
-        # 1) Checks based on per_axis_spmd_types(), that only one axis mismatches.
+        # 1) Check that only one axis mismatches.
         # Store the changed_axes so we know what to look for in PartitionSpec.
-        src_types = src.per_axis_spmd_types()
-        dst_types = dst.per_axis_spmd_types()
+        src_types = _per_axis_types(src)
+        dst_types = _per_axis_types(dst)
         if set(src_types) != set(dst_types):
             raise ValueError(
                 "SpmdLayout-based redistribute axis keys do not match for "
@@ -280,7 +296,7 @@ def _validate_redistribute_spmd_pair(
         if len(changed_axes) > 1:
             raise ValueError(
                 f"{name}: SpmdLayout-based redistribution changes multiple mesh "
-                f"axes ({sorted(axis.value for axis in changed_axes)}). "
+                f"axes ({sorted(str(axis) for axis in changed_axes)}). "
                 "spmd_redistribute_per_axis only supports one single-axis "
                 "redistribution."
             )
@@ -290,7 +306,7 @@ def _validate_redistribute_spmd_pair(
             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. "
+                f"{str(axis)!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."
             )
@@ -308,9 +324,9 @@ def _validate_redistribute_spmd_pair(
         )
         src_spec, dst_spec = src.partition_spec, dst.partition_spec
         if src_spec is None:
-            src_spec = _normalize_partition_spec(src.axis_types, ndim=ndim)
+            src_spec = _normalize_partition_spec(src.local_type, ndim=ndim)
         if dst_spec is None:
-            dst_spec = _normalize_partition_spec(dst.axis_types, ndim=ndim)
+            dst_spec = _normalize_partition_spec(dst.local_type, ndim=ndim)
 
         # A one-axis redistribute may only leave each tensor dim's shard axes
         # unchanged, add the changed axis as the innermost shard, or remove it
@@ -363,8 +379,8 @@ def _validate_redistribute_spmd_pair(
 def spmd_redistribute_per_axis(
     x: torch.Tensor,
     mesh: DeviceMesh | None,
-    src_types: spmd.PerMeshAxisSpmdTypes,
-    dst_types: spmd.PerMeshAxisSpmdTypes,
+    src: "SpmdLayout",
+    dst: "SpmdLayout",
 ) -> torch.Tensor:
     """Redistribute a local tensor along axes whose SPMD type changes.
 
@@ -379,6 +395,8 @@ def spmd_redistribute_per_axis(
     if mesh is None:
         return x
 
+    src_types = _per_axis_types(src)
+    dst_types = _per_axis_types(dst)
     assert mesh.mesh_dim_names is not None, "DeviceMesh must have named axes"
     for axis_name, dst_t in dst_types.items():
         src_t = src_types.get(axis_name)
@@ -404,7 +422,7 @@ def spmd_redistribute_per_axis(
 def spmd_distribute_tensor(
     tensor: torch.Tensor,
     mesh: DeviceMesh,
-    layout: SpmdLayout,
+    layout: spmd.SpmdType,
 ) -> torch.Tensor:
     """Materialize local state shards according to the declared SPMD layout.
 
@@ -413,11 +431,10 @@ def spmd_distribute_tensor(
     same tensor dim, e.g. ``(DP, CP)`` means shard by DP, then shard each DP
     slice by CP.
     """
-    shard_types = layout.per_axis_spmd_types()
     if layout.partition_spec is None:
         axis_shard_dims = [
             (axis_name, axis_type.dim)
-            for axis_name, axis_type in shard_types.items()
+            for axis_name, axis_type in layout.local_type.items()
             if isinstance(axis_type, spmd.Shard)
         ]
     else:
diff --git a/torchtitan/models/common/attention.py b/torchtitan/models/common/attention.py
index 7154f2fb42..32a07f71f6 100644
--- a/torchtitan/models/common/attention.py
+++ b/torchtitan/models/common/attention.py
@@ -18,7 +18,6 @@
 import spmd_types as spmd
 import torch
 import torch.nn.functional as F
-from spmd_types.runtime import get_partition_spec
 from torch.distributed.tensor import DTensor, Replicate
 from torch.distributed.tensor.experimental import local_map
 from torch.nn.attention import (
@@ -87,6 +86,21 @@ class VarlenMetadata(NamedTuple):
 )
 
 
+@spmd.no_typecheck(out_types=spmd.PartitionSpec(("dp", "cp"), "tp", None))
+def varlen_attn_fn(*args, **kwargs):
+    return varlen_attn(*args, **kwargs)
+
+
+@spmd.no_typecheck(
+    out_types=(
+        spmd.PartitionSpec(("dp", "cp"), "tp", None),
+        spmd.PartitionSpec("tp", ("dp", "cp")),
+    )
+)
+def varlen_attn_with_lse(*args, **kwargs):
+    return varlen_attn(*args, return_aux=VarlenAuxRequest(lse=True), **kwargs)
+
+
 def local_head_split(
     t: torch.Tensor,
     head_dim: int,
@@ -175,47 +189,30 @@ def forward(
         if kwargs.get("enable_gqa", False):
             varlen_kwargs["enable_gqa"] = True
 
-        if out_transform is not None:
-            varlen_kwargs["return_aux"] = VarlenAuxRequest(lse=True)
+        run_varlen_attn = (
+            varlen_attn_fn if out_transform is None else varlen_attn_with_lse
+        )
 
-        # FA3 varlen attention takes rank-local metadata tensors.
-        # TODO(pianpwk): Move this op contract into pytorch/spmd_types.
-        with spmd.no_typecheck():
-            # Some operators can upcast under AMP, but varlen attention currently only
-            # supports bf16/fp16 inputs. If this changes, or fp16 training support
-            # is added, this may need to be revisited.
-            result = varlen_attn(
-                q_TNH.to(torch.bfloat16),
-                k_TNH.to(torch.bfloat16),
-                v_TNH.to(torch.bfloat16),
-                cu_seq_q,
-                cu_seq_k,
-                max_q,
-                max_k,
-                scale=scale,
-                window_size=self.window_size,
-                **varlen_kwargs,
-            )
+        result = run_varlen_attn(
+            q_TNH.to(torch.bfloat16),
+            k_TNH.to(torch.bfloat16),
+            v_TNH.to(torch.bfloat16),
+            cu_seq_q,
+            cu_seq_k,
+            max_q,
+            max_k,
+            scale=scale,
+            window_size=self.window_size,
+            **varlen_kwargs,
+        )
 
         # varlen_attn returns the packed output (T, N, H), plus the LSE when an
         # out_transform epilogue was requested.
         if out_transform is None:
             assert isinstance(result, torch.Tensor)
-            if get_spmd_backend() == "spmd_types" and spmd.is_type_checking():
-                q_local = spmd.get_local_type(q_TNH)
-                q_ps = get_partition_spec(q_TNH)
-                spmd.assert_type(result, q_local, q_ps)
             return result.to(q_TNH.dtype)
 
         out_TNH, lse_NT = result
-        if get_spmd_backend() == "spmd_types" and spmd.is_type_checking():
-            q_local = spmd.get_local_type(q_TNH)
-            q_ps = get_partition_spec(q_TNH)
-            spmd.assert_type(out_TNH, q_local, q_ps)
-            # The current implementation returns LSE as (N, T).
-            lse_ps = None if q_ps is None else spmd.PartitionSpec(q_ps[1], q_ps[0])
-            spmd.assert_type(lse_NT, q_local, lse_ps)
-
         out_TNH = out_TNH.to(q_TNH.dtype)
         lse_TN = lse_NT.transpose(0, 1)
         return out_transform(out_TNH, lse_TN)
@@ -302,7 +299,7 @@ def compiled_flex_attn(
             )
         if get_spmd_backend() == "spmd_types" and spmd.is_type_checking():
             q_local = spmd.get_local_type(q)
-            q_ps = get_partition_spec(q)
+            q_ps = spmd.get_partition_spec(q)
             spmd.assert_type(out, q_local, q_ps)
             if return_aux.lse:
                 # lse is (B, N, L) = q minus the trailing (unsharded) head dim.
diff --git a/torchtitan/models/common/decoder_sharding.py b/torchtitan/models/common/decoder_sharding.py
index 57d8564236..4ced0df688 100644
--- a/torchtitan/models/common/decoder_sharding.py
+++ b/torchtitan/models/common/decoder_sharding.py
@@ -103,7 +103,7 @@ def dense_sequence_parallel_placement() -> SpmdLayout:
             CP: spmd.V,
             TP: spmd.V,
         },
-        partition_spec=((DP, CP, TP), None),
+        partition_spec=spmd.PartitionSpec((DP, CP, TP), None),
     )
 
 
diff --git a/torchtitan/models/common/token_dispatcher.py b/torchtitan/models/common/token_dispatcher.py
index 8130dfacd9..b5edd242f0 100644
--- a/torchtitan/models/common/token_dispatcher.py
+++ b/torchtitan/models/common/token_dispatcher.py
@@ -20,7 +20,7 @@
     init_buffer as minimal_async_ep_init_buffer,
     MinimalAsyncEPDispatchMetadata,
 )
-from torchtitan.distributed.spmd_types import current_spmd_mesh, maybe_set_sparse_mesh
+from torchtitan.distributed.spmd_types import maybe_set_sparse_mesh
 from torchtitan.distributed.utils import get_spmd_backend
 from torchtitan.ops.scatter_add import deterministic_scatter_add
 from torchtitan.tools.utils import device_module, device_type
@@ -421,22 +421,16 @@ def dispatch(
         if (
             get_spmd_backend() == "spmd_types" and spmd.is_type_checking()
         ):  # sparse mesh reinterpret
-            for axis in ["dp", "cp", "tp"]:
-                spmd.mutate_type(
-                    num_local_tokens_per_expert_E,
-                    axis,
-                    src=spmd.P,
-                    dst=spmd.V,
-                )
+            spmd.mutate_type(
+                num_local_tokens_per_expert_E,
+                src=spmd.P,
+                dst={"dp": spmd.V, "cp": spmd.V, "tp": spmd.V},
+            )
 
         # generate the input splits and output splits for all-to-all
         with maybe_set_sparse_mesh():
             pg = (
-                current_spmd_mesh().get_group(  # pyrefly: ignore [missing-attribute]
-                    "ep"
-                )
-                if get_spmd_backend() == "spmd_types"
-                else self.ep_mesh.get_group()
+                "ep" if get_spmd_backend() == "spmd_types" else self.ep_mesh.get_group()
             )
             if get_spmd_backend() == "spmd_types" and spmd.is_type_checking():
                 num_local_tokens_per_expert_E = spmd.reinterpret_mesh(
@@ -586,11 +580,7 @@ def combine(
 
         with maybe_set_sparse_mesh():
             pg = (
-                current_spmd_mesh().get_group(  # pyrefly: ignore [missing-attribute]
-                    "ep"
-                )
-                if get_spmd_backend() == "spmd_types"
-                else self.ep_mesh.get_group()
+                "ep" if get_spmd_backend() == "spmd_types" else self.ep_mesh.get_group()
             )
             # Reverse expert-major reordering
             routed_output_RD = self._unpermute(
diff --git a/torchtitan/models/deepseek_v3/mtp.py b/torchtitan/models/deepseek_v3/mtp.py
index 90068b1d26..4ed417b9b6 100644
--- a/torchtitan/models/deepseek_v3/mtp.py
+++ b/torchtitan/models/deepseek_v3/mtp.py
@@ -390,17 +390,8 @@ def __call__(
             mtp_loss = depth_loss if mtp_loss is None else mtp_loss + depth_loss
         assert mtp_loss is not None
         if num_mtp_layers > 1:
-            # TODO: Teach spmd_types that V / scalar preserves the scalar
-            # loss placement. This mirrors the base loss normalization.
-            with spmd.no_typecheck():
-                mtp_loss = mtp_loss / num_mtp_layers
-        # TODO: Teach spmd_types that scalar loss composition preserves
-        # the loss placement across auxiliary weighted losses.
-        with spmd.no_typecheck():
-            loss = main_loss + mtp_loss * self.mtp_scale
+            mtp_loss = mtp_loss / num_mtp_layers
+        loss = main_loss + mtp_loss * self.mtp_scale
         if global_valid_tokens is not None:
-            # TODO: Teach spmd_types that scalar loss normalization preserves
-            # the loss placement.
-            with spmd.no_typecheck():
-                loss = loss / global_valid_tokens
+            loss = loss / global_valid_tokens
         return loss, {}
diff --git a/torchtitan/models/gpt_oss/moe.py b/torchtitan/models/gpt_oss/moe.py
index db156afe2f..c6a274c171 100644
--- a/torchtitan/models/gpt_oss/moe.py
+++ b/torchtitan/models/gpt_oss/moe.py
@@ -20,7 +20,6 @@
 from torchtitan.protocols.module import Module
 
 
-@spmd.register_autograd_function
 class ScaleBiasForward(torch.autograd.Function):
     """
     Custom autograd function that scales bias in forward pass but not in backward.
@@ -38,7 +37,7 @@ def forward(ctx, bias, tp_degree, dtype):
         return bias.to(dtype)
 
     @staticmethod
-    def typecheck_forward(bias, tp_degree, dtype):
+    def spmd_typecheck(out, *, bias):
         """
         Typecheck for bias scaling, already interleaved to num tokens shape.
         If EP enabled, V on all axes. If disabled, TP axis: R->V.
@@ -53,9 +52,7 @@ def typecheck_forward(bias, tp_degree, dtype):
             in_type = {"dp": spmd.V, "cp": spmd.V, "tp": spmd.R}
             out_type = {"dp": spmd.V, "cp": spmd.V, "tp": spmd.V}
         spmd.assert_type(bias, in_type)
-        out = ScaleBiasForward.apply(bias, tp_degree, dtype)
         spmd.assert_type(out, out_type)
-        return out
 
     @staticmethod
     # pyrefly: ignore [bad-override]
@@ -142,8 +139,11 @@ def forward(
             and spmd.is_type_checking()
             and spmd_mesh_size("ep") == 1
         ):
-            for axis in ("dp", "cp"):
-                spmd.mutate_type(num_tokens_per_expert_E, axis, src=spmd.P, dst=spmd.V)
+            spmd.mutate_type(
+                num_tokens_per_expert_E,
+                src=spmd.P,
+                dst={"dp": spmd.V, "cp": spmd.V},
+            )
 
         offsets_E = torch.cumsum(num_tokens_per_expert_E, dim=0, dtype=torch.int32)
         # Pad num_tokens_per_expert_E with tail slack so that repeat_interleave
diff --git a/torchtitan/overrides/fused_mla.py b/torchtitan/overrides/fused_mla.py
index 3b7e96179d..a687171362 100644
--- a/torchtitan/overrides/fused_mla.py
+++ b/torchtitan/overrides/fused_mla.py
@@ -609,23 +609,21 @@ def _fused_mla_kv_backward_op_fake(
     )
 
 
-@spmd.register_autograd_function
 class _FusedMLAQ(torch.autograd.Function):
     @staticmethod
-    def typecheck_forward(
+    def spmd_typecheck(
+        output: torch.Tensor,
+        *,
         q: torch.Tensor,
         rope_cache_real: torch.Tensor,
         positions: torch.Tensor,
-        q_nope_dim: int,
-    ) -> torch.Tensor:
+    ) -> None:
         q_type = (spmd.V, spmd.PartitionSpec(None, ("dp", "cp"), "tp", None))
         positions_type = (spmd.V, spmd.PartitionSpec(None, ("dp", "cp")))
         spmd.assert_type(q, *q_type)
         spmd.assert_type(rope_cache_real, spmd.R)
         spmd.assert_type(positions, *positions_type)
-        output = _FusedMLAQ.apply(q, rope_cache_real, positions, q_nope_dim)
         spmd.assert_type(output, *q_type)
-        return output
 
     @staticmethod
     def forward(
@@ -667,16 +665,16 @@ def backward(ctx, grad_q: torch.Tensor):
         return grad_q, None, None, None
 
 
-@spmd.register_autograd_function
 class _FusedMLAKV(torch.autograd.Function):
     @staticmethod
-    def typecheck_forward(
+    def spmd_typecheck(
+        outputs: tuple[torch.Tensor, torch.Tensor],
+        *,
         kv: torch.Tensor,
         k_pe: torch.Tensor,
         rope_cache_real: torch.Tensor,
         positions: torch.Tensor,
-        q_nope_dim: int,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> None:
         kv_partition = spmd.PartitionSpec(None, ("dp", "cp"), "tp", None)
         k_pe_partition = spmd.PartitionSpec(None, ("dp", "cp"), None)
         positions_partition = spmd.PartitionSpec(None, ("dp", "cp"))
@@ -684,16 +682,9 @@ def typecheck_forward(
         spmd.assert_type(k_pe, spmd.V, k_pe_partition)
         spmd.assert_type(rope_cache_real, spmd.R)
         spmd.assert_type(positions, spmd.V, positions_partition)
-        k, v = _FusedMLAKV.apply(
-            kv,
-            k_pe,
-            rope_cache_real,
-            positions,
-            q_nope_dim,
-        )
+        k, v = outputs
         spmd.assert_type(k, spmd.V, kv_partition)
         spmd.assert_type(v, spmd.V, kv_partition)
-        return k, v
 
     @staticmethod
     def forward(
diff --git a/torchtitan/protocols/module.py b/torchtitan/protocols/module.py
index 42e5ad939e..f63f4129cf 100644
--- a/torchtitan/protocols/module.py
+++ b/torchtitan/protocols/module.py
@@ -15,14 +15,13 @@
 import spmd_types as spmd
 import torch
 import torch.nn as nn
-from spmd_types.runtime import get_local_type, get_partition_spec, has_local_type
 from torch.distributed.tensor import distribute_tensor, DTensor
 from torch.distributed.tensor.experimental import local_map
 from torch.distributed.tensor.placement_types import Placement
 from torch.utils._pytree import tree_map
 
 from torchtitan.config import Configurable
-from torchtitan.distributed.parallel_dims import ParallelDims, SpmdLayout
+from torchtitan.distributed.parallel_dims import layout_axes, ParallelDims, SpmdLayout
 from torchtitan.distributed.spmd_types import (
     current_spmd_mesh,
     set_current_spmd_mesh,
@@ -138,24 +137,18 @@ def _preserve_buffer_spmd_types(self) -> Iterator[None]:
 
         saved = {
             fqn: SpmdLayout(
-                # pyrefly: ignore [bad-argument-type]
-                axis_types=get_local_type(buf),
-                partition_spec=get_partition_spec(buf),
+                dict(spmd.get_local_type(buf)),
+                spmd.get_partition_spec(buf),
             )
             for fqn, buf in self.named_buffers()
-            if has_local_type(buf)
+            if spmd.has_local_type(buf)
         }
         try:
             yield
         finally:
             for fqn, buf in self.named_buffers():
-                if fqn in saved and not has_local_type(buf):
-                    layout = saved[fqn]
-                    spmd.assert_type(
-                        buf,
-                        layout.axis_types,
-                        partition_spec=layout.partition_spec,
-                    )
+                if fqn in saved and not spmd.has_local_type(buf):
+                    spmd.assert_type(buf, saved[fqn])
 
     def _init_self_parameters(self) -> None:
         """Initialize this module's own direct parameters.
@@ -305,7 +298,8 @@ def _spmd_distribute_state(
         # Call get_optional_mesh with include_singleton_axes=True, so we're able to call assert_type()
         # using all axes, and defer size-1 axis filtering to spmd_types internals.
         mesh = parallel_dims.get_optional_mesh(
-            [axis.value for axis in layout.axes()], include_singleton_axes=True
+            [axis.value for axis in layout_axes(layout)],
+            include_singleton_axes=True,
         )
         assert mesh is not None
         assert mesh.mesh_dim_names is not None, "DeviceMesh must have named axes"
@@ -319,14 +313,10 @@ def _spmd_distribute_state(
             self.register_buffer(name, tensor, persistent=persistent)
             registered = self._buffers[name]
 
-        # assert_type resolves SpmdLayout's string mesh axis names to concrete
+        # assert_type resolves SpmdType's string mesh axis names to concrete
         # runtime mesh-axis objects, so a mesh context is required here.
         with set_current_spmd_mesh(mesh):
-            spmd.assert_type(
-                registered,
-                layout.axis_types,
-                layout.partition_spec,
-            )
+            spmd.assert_type(registered, layout)
 
     def _distribute_states(self, parallel_dims: ParallelDims) -> None:
         """Distribute params and buffers per ``state_shardings``.
@@ -355,7 +345,7 @@ def _distribute_states(self, parallel_dims: ParallelDims) -> None:
                     is_param=True,
                 )
                 continue
-            axes = spmd_layout.axes()
+            axes = layout_axes(spmd_layout)
             mesh = parallel_dims.resolve_mesh(axes)
             if mesh is None:
                 continue
@@ -398,7 +388,7 @@ def _distribute_states(self, parallel_dims: ParallelDims) -> None:
                     is_param=False,
                 )
                 continue
-            axes = spmd_layout.axes()
+            axes = layout_axes(spmd_layout)
             mesh = parallel_dims.resolve_mesh(axes)
             if mesh is None:
                 continue
@@ -503,14 +493,14 @@ def _spmd_apply_local_map(
     ) -> Callable:
         """Apply spmd_types local_map for a local-tensor compute region."""
         in_types = tuple(
-            (layout.axis_types, layout.partition_spec) for layout in in_named
+            (layout.local_type, layout.partition_spec) for layout in in_named
         )
         out_types = tree_map(
-            lambda layout: (layout.axis_types, layout.partition_spec),
+            lambda layout: (layout.local_type, layout.partition_spec),
             out_src,
-            is_leaf=lambda x: isinstance(x, SpmdLayout),
+            is_leaf=lambda x: isinstance(x, spmd.SpmdType),
         )
-        return spmd.local_map(
+        return spmd.no_typecheck(
             in_types=in_types,
             out_types=out_types,
         )(fn)
@@ -566,11 +556,7 @@ def _redistribute_inputs(
                 # before redistributing so typechecking catches placement mismatch.
                 # Gate assertion so compile doesn't error.
                 if spmd.is_type_checking():
-                    spmd.assert_type(
-                        value,
-                        src_spmd_layout.axis_types,
-                        src_spmd_layout.partition_spec,
-                    )
+                    spmd.assert_type(value, src_spmd_layout)
 
                 if dst_spmd_layout is None:
                     new_kwargs[name] = value
@@ -578,10 +564,8 @@ def _redistribute_inputs(
                 value = spmd_redistribute_per_axis(
                     value,
                     current_spmd_mesh(),
-                    # pyrefly: ignore [bad-argument-type]
-                    src_spmd_layout.per_axis_spmd_types(),
-                    # pyrefly: ignore [bad-argument-type]
-                    dst_spmd_layout.per_axis_spmd_types(),
+                    src_spmd_layout,
+                    dst_spmd_layout,
                 )
                 new_kwargs[name] = value
                 continue
@@ -656,21 +640,15 @@ def _redistribute_outputs(self, parallel_dims: ParallelDims, outputs: Any) -> An
             # before redistributing so typechecking catches placement mismatch.
             # Gate assertion so compile doesn't error.
             if spmd.is_type_checking():
-                spmd.assert_type(
-                    outputs,
-                    out_src.axis_types,
-                    out_src.partition_spec,
-                )
+                spmd.assert_type(outputs, out_src)
 
             if out_dst is None:
                 return outputs
             return spmd_redistribute_per_axis(
                 outputs,
                 current_spmd_mesh(),
-                # pyrefly: ignore [bad-argument-type]
-                out_src.per_axis_spmd_types(),
-                # pyrefly: ignore [bad-argument-type]
-                out_dst.per_axis_spmd_types(),
+                out_src,
+                out_dst,
             )
 
         if isinstance(out_src, tuple):
diff --git a/torchtitan/protocols/sharding.py b/torchtitan/protocols/sharding.py
index 4f3af43cfd..661b24d4da 100644
--- a/torchtitan/protocols/sharding.py
+++ b/torchtitan/protocols/sharding.py
@@ -19,6 +19,7 @@
 from torch.distributed.tensor import Partial, Placement, Replicate, Shard
 
 from torchtitan.distributed.parallel_dims import (
+    layout_axes,
     MeshAxisName,
     SpmdLayout,
     unfold_dp_axis,
@@ -32,12 +33,6 @@
     "resolve_placements",
 ]
 
-# Shard order: we implicitly assume the trivial outer -> inner order matching
-# the mesh axis order. The only non-trivial case is FSDP + TP both sharding on
-# tensor dim 0, but it doesn't need to be annotated today.
-# TODO: integrate with global spmd types (e.g., ``TP: V`` + ``PartitionSpec``
-# carrying explicit shard-order info) once that lands.
-
 
 @dataclass(kw_only=True, slots=True)
 class LocalMapConfig:
@@ -146,22 +141,30 @@ 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"
-    axis_types = {}
-    for axis_name, axis_type in layout.per_axis_spmd_types().items():
+    axis_types = dict(layout.local_type)
+    if layout.partition_spec is not None:
+        for dim, entry in enumerate(layout.partition_spec):
+            for axis_name in (
+                () if entry is None else entry if isinstance(entry, tuple) else (entry,)
+            ):
+                axis_types[axis_name] = spmd.S(dim)
+    concrete_axis_types = {}
+    for axis_name in layout_axes(layout):
+        axis_type = axis_types[axis_name]
         for concrete_axis_name in unfold_dp_axis(axis_name):
-            axis_types[concrete_axis_name] = axis_type
+            concrete_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 axis_types:
+        if key not in concrete_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"{sorted(k.value for k in layout_axes(layout))}; "
                 f"required: {list(mesh.mesh_dim_names)}."
             )
-        p = spmd.spmd_type_to_dtensor_placement(axis_types[key])
+        p = spmd.spmd_type_to_dtensor_placement(concrete_axis_types[key])
         if isinstance(p, (Shard, Partial)) and mesh.size(i) == 1:
             p = Replicate()
         result.append(p)