diff --git a/tests/unit_tests/test_distributed_utils.py b/tests/unit_tests/test_distributed_utils.py
index 4614ad1bed..4a40aac9b5 100644
--- a/tests/unit_tests/test_distributed_utils.py
+++ b/tests/unit_tests/test_distributed_utils.py
@@ -8,11 +8,22 @@
 from unittest.mock import patch
 
 import pytest
+import spmd_types as spmd
 import torch
-from torch.distributed.device_mesh import DeviceMesh
+import torch.nn as nn
+from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
+from torch.testing._internal.distributed._tensor.common_dtensor import (
+    DTensorTestBase,
+    with_comms,
+)
 
 from torchtitan.config import CommConfig
 from torchtitan.distributed import utils as dist_utils
+from torchtitan.distributed.spmd_types import (
+    set_current_spmd_mesh,
+    set_spmd_meshes,
+    spmd_distribute_tensor,
+)
 from torchtitan.distributed.utils import init_distributed
 
 
@@ -61,3 +72,107 @@ def test_dist_sum_tensor_waits_for_distributed_result():
     assert result is reduced
     reduce.assert_called_once_with(value, reduceOp="SUM", group=mesh)
     wait.assert_called_once_with(reduced)
+
+
+class TestSpmdLocalGradNorm(DTensorTestBase):
+    @property
+    def world_size(self) -> int:
+        return 8
+
+    @property
+    def device_type(self) -> str:
+        return "cpu"
+
+    @with_comms
+    def test_multiple_meshes_and_placements(self) -> None:
+        dense_mesh = init_device_mesh(
+            "cpu",
+            (2, 4),
+            mesh_dim_names=("dense_0", "dense_1"),
+        )
+        sparse_mesh = init_device_mesh(
+            "cpu",
+            (4, 2),
+            mesh_dim_names=("sparse_0", "sparse_1"),
+        )
+        set_spmd_meshes(
+            dense_mesh=dense_mesh,
+            dense_storage_mesh=dense_mesh,
+            sparse_mesh=sparse_mesh,
+            sparse_storage_mesh=sparse_mesh,
+        )
+
+        # Construct eight parameters across two meshes. Each mesh has two R,R
+        # parameters and two R,S(0) parameters.
+        model = nn.Module()
+        model.dense_params = nn.ParameterList()
+        model.sparse_params = nn.ParameterList()
+        global_grads = []
+        torch.manual_seed(42)
+        second_axis_sharded_placements = (False, True)
+        num_params_per_placement = 2
+
+        for mesh, params in (
+            (dense_mesh, model.dense_params),
+            (sparse_mesh, model.sparse_params),
+        ):
+            assert mesh.mesh_dim_names is not None
+            mesh_axes = tuple(
+                spmd.MeshAxis.of(mesh.get_group(axis_name))
+                for axis_name in mesh.mesh_dim_names
+            )
+            for second_axis_sharded in second_axis_sharded_placements:
+                for _ in range(num_params_per_placement):
+                    global_grad = torch.randn(8, 8, dtype=torch.float32)
+                    axis_types = {
+                        mesh_axes[0]: spmd.R,
+                        mesh_axes[1]: (spmd.S(0) if second_axis_sharded else spmd.R),
+                    }
+                    local_grad = spmd_distribute_tensor(
+                        global_grad.clone(), mesh, spmd.SpmdType(axis_types)
+                    )
+
+                    parameter = nn.Parameter(local_grad.clone())
+                    with set_current_spmd_mesh(mesh):
+                        spmd.assert_type(
+                            parameter,
+                            {
+                                mesh.mesh_dim_names[0]: spmd.R,
+                                mesh.mesh_dim_names[1]: (
+                                    spmd.S(0) if second_axis_sharded else spmd.R
+                                ),
+                            },
+                        )
+                    parameter.grad = local_grad.clone()
+                    params.append(parameter)
+                    global_grads.append(global_grad)
+
+        # compare against globally computed grad norm
+        expected_norm = (
+            torch.stack([global_grad.square().sum() for global_grad in global_grads])
+            .sum()
+            .sqrt()
+        )
+        # clip_grad_norm_spmd_ should only issue 1 all-reduce per mesh.
+        expected_groups = [
+            dense_mesh.get_group("dense_1"),
+            sparse_mesh.get_group("sparse_1"),
+        ]
+
+        with patch.object(
+            dist_utils.dist,
+            "all_reduce",
+            wraps=dist_utils.dist.all_reduce,
+        ) as all_reduce:
+            actual_norm = dist_utils.clip_grad_norm_spmd_(
+                list(model.parameters()),
+                max_norm=float("inf"),
+                foreach=True,
+            )
+
+        self.assertEqual(actual_norm, expected_norm, rtol=1e-5, atol=1e-5)
+        self.assertEqual(all_reduce.call_count, 2)
+        self.assertEqual(
+            [call.kwargs["group"] for call in all_reduce.call_args_list],
+            expected_groups,
+        )
diff --git a/tests/unit_tests/test_state_dict_adapter.py b/tests/unit_tests/test_state_dict_adapter.py
index 457f27fb17..0667132441 100644
--- a/tests/unit_tests/test_state_dict_adapter.py
+++ b/tests/unit_tests/test_state_dict_adapter.py
@@ -7,13 +7,20 @@
 import tempfile
 import unittest
 
+import spmd_types as spmd
 import torch
 import torch.distributed as dist
 from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.tensor import DTensor, Replicate, Shard
+from torch.testing._internal.distributed._tensor.common_dtensor import (
+    DTensorTestBase,
+    with_comms,
+)
 
+from torchtitan.distributed.spmd_types import set_spmd_meshes, spmd_distribute_tensor
 from torchtitan.models.deepseek_v3 import deepseekv3_configs
 from torchtitan.models.deepseek_v3.state_dict_adapter import DeepSeekV3StateDictAdapter
+from torchtitan.protocols.state_dict_adapter import PlainToDTensorStateDictAdapter
 
 
 class DeepSeekV3StateDictAdapterTest(unittest.TestCase):
@@ -70,3 +77,38 @@ def test_to_hf_handles_replicated_grouped_experts(self) -> None:
                 hf_state_dict[key].to_local(),
                 local_weight[expert],
             )
+
+
+class PlainToDTensorStateDictAdapterTest(DTensorTestBase):
+    @property
+    def world_size(self) -> int:
+        return 4
+
+    @property
+    def device_type(self) -> str:
+        return "cpu"
+
+    @with_comms
+    def test_spmd_type_partition_order(self) -> None:
+        global_weight = torch.arange(16, dtype=torch.float32).reshape(8, 2)
+        mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("dp_shard", "tp"))
+        set_spmd_meshes(
+            dense_mesh=mesh,
+            dense_storage_mesh=mesh,
+            sparse_mesh=None,
+            sparse_storage_mesh=None,
+        )
+        dp_axis = spmd.MeshAxis.of(mesh.get_group("dp_shard"))
+        tp_axis = spmd.MeshAxis.of(mesh.get_group("tp"))
+        layout = spmd.SpmdType(
+            {dp_axis: spmd.V, tp_axis: spmd.V},
+            spmd.PartitionSpec((tp_axis, dp_axis), None),
+        )
+        local_weight = spmd_distribute_tensor(global_weight, mesh, layout)
+        adapter = PlainToDTensorStateDictAdapter({"weight": layout})
+        state_dict = adapter.convert_save_state_dict({"weight": local_weight})
+
+        torch.testing.assert_close(state_dict["weight"].full_tensor(), global_weight)
+        torch.testing.assert_close(
+            adapter.convert_load_state_dict(state_dict)["weight"], local_weight
+        )
diff --git a/torchtitan/components/loss.py b/torchtitan/components/loss.py
index 971a1ad49f..58db7bc2d9 100644
--- a/torchtitan/components/loss.py
+++ b/torchtitan/components/loss.py
@@ -19,8 +19,13 @@
 from torch.distributed.tensor.experimental import local_map
 
 from torchtitan.config import CompileConfig, Configurable
-from torchtitan.distributed.spmd_types import current_spmd_mesh, spmd_mesh_size
+from torchtitan.distributed.spmd_types import (
+    current_spmd_mesh,
+    get_mesh_pg,
+    spmd_mesh_size,
+)
 from torchtitan.distributed.utils import get_spmd_backend
+from torchtitan.models.common.embedding import get_tp_rank
 from torchtitan.tools.logging import logger
 
 # PyTorch's default ignore index for cross-entropy loss
@@ -35,28 +40,29 @@ def cross_entropy_loss(
     *,
     global_vocab_size: int | None = None,
 ) -> torch.Tensor:
-    """Cross-entropy over ``pred[T, V]`` and ``labels[T]`` with sum reduction."""
+    """Cross-entropy loss with sum reduction for token-based normalization."""
     if isinstance(pred, DTensor):
         assert get_spmd_backend() == "partial_dtensor"
-        if pred.placements == (Shard(1),):
+        if pred.placements == (Shard(pred.ndim - 1),):
             return _LossParallelCrossEntropy.apply(
-                pred.to_local().float(),
-                labels,
+                pred.to_local().flatten(0, -2).float(),
+                labels.flatten(),
                 pred.device_mesh.get_group("tp"),
                 pred.shape[-1],
                 "sum",
             )
     elif get_spmd_backend() == "spmd_types" and spmd_mesh_size("tp") > 1:
         return _LossParallelCrossEntropy.apply(
-            pred.float(),
-            labels,
-            current_spmd_mesh().get_group("tp"),  # pyrefly: ignore[missing-attribute]
+            pred.flatten(0, -2).float(),
+            labels.flatten(),
+            get_mesh_pg("tp"),
             global_vocab_size,
+            "sum",
         )
 
     return torch.nn.functional.cross_entropy(
-        pred.float(),
-        labels,
+        pred.flatten(0, -2).float(),
+        labels.flatten(),
         reduction="sum",
         ignore_index=IGNORE_INDEX,
     )
@@ -64,7 +70,7 @@ def cross_entropy_loss(
 
 class _LossParallelCrossEntropy(torch.autograd.Function):
     """
-    Vocab-parallel cross-entropy on local ``[T, V_local]`` logits.
+    Vocab-parallel cross-entropy on plain (non-DTensor) local tensors.
 
     Replaces ``torch.distributed.tensor.parallel.loss_parallel()`` with an
     explicit autograd Function so that SPMD code can operate on local tensors
@@ -109,38 +115,38 @@ def forward(
         """Compute exact CE from local vocab shards via TP all-reduces.
 
         ``reduction="sum"`` returns the scalar summed loss (SFT/CE).
-        ``reduction="none"`` returns the per-token NLL ``[T]``, which GRPO
+        ``reduction="none"`` returns the per-token NLL ``[N]``, which GRPO
         negates to get per-token logprobs without all-gathering the vocab.
         """
+        logits_shape = logits.shape
         logits_dtype = logits.dtype
-        logits = logits.float()
+        logits_2d = logits.flatten(0, -2).float()
+        labels_1d = labels.flatten()
 
         # Compute this rank's vocab shard bounds for the local logits.
         tp_world_size = dist.get_world_size(tp_group)
-        tp_rank = dist.get_rank(tp_group)
+        tp_rank = get_tp_rank(tp_group)
         chunk_size = (global_vocab_size + tp_world_size - 1) // tp_world_size
-        vocab_start = min(global_vocab_size, chunk_size * tp_rank)
-        vocab_end = min(global_vocab_size, vocab_start + chunk_size)
-        local_vocab_size = max(0, vocab_end - vocab_start)
-        if logits.shape[-1] != local_vocab_size:
-            raise ValueError(
-                "_LossParallelCrossEntropy expected local vocab size "
-                f"{local_vocab_size} for global vocab size {global_vocab_size}, "
-                f"got {logits.shape[-1]}."
-            )
-        if local_vocab_size == 0:
-            raise ValueError(
-                "_LossParallelCrossEntropy does not support empty vocab shards."
-            )
+        vocab_start = torch.sym_min(global_vocab_size, chunk_size * tp_rank)
+        vocab_end = torch.sym_min(global_vocab_size, vocab_start + chunk_size)
+        local_vocab_size = torch.sym_max(0, vocab_end - vocab_start)
+        torch._check(
+            logits_2d.shape[-1] == local_vocab_size,
+            lambda: "_LossParallelCrossEntropy local vocab size mismatch.",
+        )
+        torch._check(
+            local_vocab_size > 0,
+            lambda: "_LossParallelCrossEntropy does not support empty vocab shards.",
+        )
 
         # All-reduce max for numerically stable distributed log-softmax.
-        local_max = torch.amax(logits, dim=-1, keepdim=True)
+        local_max = torch.amax(logits_2d, dim=-1, keepdim=True)
         local_max = funcol.all_reduce(
             local_max, reduceOp=dist.ReduceOp.MAX.name, group=tp_group
         )
 
         # All-reduce sum over shifted logits for the global softmax denominator.
-        shifted = logits - local_max
+        shifted = logits_2d - local_max
         shifted_sumexp = torch.sum(torch.exp(shifted), dim=-1, keepdim=True)
         shifted_sumexp = funcol.all_reduce(
             shifted_sumexp, reduceOp=dist.ReduceOp.SUM.name, group=tp_group
@@ -149,7 +155,7 @@ def forward(
 
         # Mask labels outside this vocab shard; the TP all-reduce below selects
         # the owner rank's log probability for each target token.
-        safe_labels = torch.where(labels != IGNORE_INDEX, labels, 0)
+        safe_labels = torch.where(labels_1d != IGNORE_INDEX, labels_1d, 0)
         out_of_range = (safe_labels < vocab_start) | (
             safe_labels >= vocab_start + local_vocab_size
         )
@@ -164,11 +170,12 @@ def forward(
 
         # Per-token NLL, dropping ignored labels (logprob 0 for ignored).
         result = -local_result.squeeze(-1)
-        result = torch.where(labels != IGNORE_INDEX, result, 0)
+        result = torch.where(labels_1d != IGNORE_INDEX, result, 0)
 
         # Save local-shard log probabilities for the fused CE backward.
-        ctx.save_for_backward(log_probs, labels)
-        ctx.logits_dtype = logits_dtype
+        ctx.save_for_backward(log_probs, labels_1d)
+        ctx.logits_shape = logits_shape
+        ctx.logits_dtype = logits.dtype
         ctx.vocab_start = vocab_start
         ctx.local_vocab_size = local_vocab_size
         ctx.reduction = reduction
@@ -181,8 +188,8 @@ def backward(  # pyrefly: ignore[bad-override]
         ctx,
         grad_output: torch.Tensor,
     ) -> tuple[torch.Tensor, None, None, None, None]:
-        log_probs, labels = ctx.saved_tensors
-        safe_labels = torch.where(labels != IGNORE_INDEX, labels, 0)
+        log_probs, labels_1d = ctx.saved_tensors
+        safe_labels = torch.where(labels_1d != IGNORE_INDEX, labels_1d, 0)
         out_of_range = (safe_labels < ctx.vocab_start) | (
             safe_labels >= ctx.vocab_start + ctx.local_vocab_size
         )
@@ -194,16 +201,16 @@ def backward(  # pyrefly: ignore[bad-override]
         grad_update = out_of_range.to(grad_input.dtype) - 1.0
         grad_input[row_idx, local_labels] = grad_update
 
-        # reduction="none" gives a per-token ``[T]`` upstream grad; unsqueeze to
-        # ``[T, 1]`` to broadcast over the local vocab. "sum" gives the scalar
+        # reduction="none" gives a per-token ``[N]`` upstream grad; reshape to
+        # ``[N, 1]`` to broadcast over the local vocab. "sum" gives the scalar
         # loss grad, which broadcasts as-is.
         if ctx.reduction == "none":
-            grad_output = grad_output.unsqueeze(-1)
+            grad_output = grad_output.reshape(-1, 1)
         grad_output = torch.where(
-            (labels != IGNORE_INDEX).unsqueeze(-1), grad_output, 0
+            (labels_1d != IGNORE_INDEX).unsqueeze(-1), grad_output, 0
         )
         grad_logits = (grad_input + torch.exp(log_probs)) * grad_output
-        grad_logits = grad_logits.to(ctx.logits_dtype)
+        grad_logits = grad_logits.reshape(ctx.logits_shape).to(ctx.logits_dtype)
         return grad_logits, None, None, None, None
 
 
diff --git a/torchtitan/distributed/parallel_dims.py b/torchtitan/distributed/parallel_dims.py
index 5c12be71d4..ded530f0e9 100644
--- a/torchtitan/distributed/parallel_dims.py
+++ b/torchtitan/distributed/parallel_dims.py
@@ -240,6 +240,7 @@ def unflatten_mesh(
         )
         loss_mesh = dataloading_mesh["batch", "cp"]._flatten("loss_mesh")
         spmd_dense_mesh_for_fwdbwd = None
+        spmd_dense_mesh_for_storage = None
         if self.spmd_backend == "spmd_types":
             # Two mesh views over the same devices:
             #
@@ -258,6 +259,16 @@ def unflatten_mesh(
                 ("pp", "dp_replicate", "dp_shard", "cp", "tp"),
                 (self.pp, self.dp_replicate, self.dp_shard, self.cp, self.tp),
             )
+            dense_fsdp_mesh = full_dense_mesh_for_fsdp["dp_shard", "cp"]._flatten(
+                "fsdp"
+            )
+            spmd_dense_mesh_for_storage = DeviceMesh._concatenate(
+                [
+                    full_dense_mesh_for_fsdp["dp_replicate"],
+                    dense_fsdp_mesh,
+                    full_dense_mesh_for_fsdp["tp"],
+                ]
+            )
             full_dense_mesh_for_fwdbwd = unflatten_mesh(
                 self._world_mesh,
                 ("pp", "dp", "cp", "tp"),
@@ -286,6 +297,8 @@ def unflatten_mesh(
         }
         if spmd_dense_mesh_for_fwdbwd is not None:
             self._global_meshes["spmd_dense_for_fwdbwd"] = spmd_dense_mesh_for_fwdbwd
+        if spmd_dense_mesh_for_storage is not None:
+            self._global_meshes["spmd_dense_for_storage"] = spmd_dense_mesh_for_storage
         if self.spmd_backend == "spmd_types" and self.ep > 1:
             self._global_meshes["spmd_sparse_for_fwdbwd"] = full_sparse_mesh[
                 "dp_replicate", "efsdp", "ep"
@@ -302,8 +315,10 @@ def unflatten_mesh(
         }
         if self.spmd_backend == "spmd_types":
             assert spmd_dense_mesh_for_fwdbwd is not None
+            assert spmd_dense_mesh_for_storage is not None
             self._single_axis_meshes["dp"] = spmd_dense_mesh_for_fwdbwd["dp"]
             self._single_axis_meshes["dp_shard"] = full_dense_mesh_for_fsdp["dp_shard"]
+            self._single_axis_meshes["fsdp"] = spmd_dense_mesh_for_storage["fsdp"]
         else:
             self._single_axis_meshes["fsdp"] = full_dense_mesh_for_fsdp["fsdp"]
 
@@ -331,6 +346,7 @@ def _validate_meshes(self):
         if self.spmd_backend == "spmd_types":
             expected_sizes["dp"] = self.dp_replicate * self.dp_shard
             expected_sizes["dp_shard"] = self.dp_shard
+            expected_sizes["fsdp"] = self.dp_shard * self.cp
         else:
             expected_sizes["fsdp"] = self.dp_shard * self.cp
 
@@ -441,12 +457,26 @@ def spmd_dense_mesh(self) -> DeviceMesh:
             self.build_mesh()
         return self._global_meshes["spmd_dense_for_fwdbwd"]
 
+    def spmd_dense_storage_mesh(self) -> DeviceMesh:
+        """Dense SPMD mesh used for parameter storage and initialization."""
+        if not self._single_axis_meshes:
+            self.build_mesh()
+        assert self.spmd_backend == "spmd_types"
+        return self._global_meshes["spmd_dense_for_storage"]
+
     def spmd_sparse_mesh(self) -> DeviceMesh | None:
         """Sparse SPMD mesh used inside expert dispatch."""
         if not self._single_axis_meshes:
             self.build_mesh()
         return self._global_meshes.get("spmd_sparse_for_fwdbwd")
 
+    def spmd_sparse_storage_mesh(self) -> DeviceMesh | None:
+        """Sparse SPMD mesh used for routed-expert parameter storage."""
+        if not self._single_axis_meshes:
+            self.build_mesh()
+        assert self.spmd_backend == "spmd_types"
+        return self._global_meshes.get("spmd_sparse_for_fwdbwd")
+
     def get_dense_tp_mesh(self) -> DeviceMesh:
         """Return the TP-axis mesh used by dense forward/backward computation."""
         if self.spmd_backend == "spmd_types":
diff --git a/torchtitan/distributed/spmd_types.py b/torchtitan/distributed/spmd_types.py
index f11bb4dc69..8076fbb5d1 100644
--- a/torchtitan/distributed/spmd_types.py
+++ b/torchtitan/distributed/spmd_types.py
@@ -9,12 +9,14 @@
 from __future__ import annotations
 
 import contextlib
-from collections.abc import Iterator, Mapping
+from collections.abc import Iterable, Iterator, Mapping
 from threading import local
 from typing import Any, TYPE_CHECKING
 
 import spmd_types as spmd
 import torch
+import torch.distributed as dist
+from spmd_types.types import DeviceMeshAxis, partition_spec_get_shard
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor import DTensor
 
@@ -23,7 +25,7 @@
 # Avoid circular import: protocols.__init__ imports module.py, which imports us.
 if TYPE_CHECKING:
     from torchtitan.distributed.parallel_dims import (
-        MeshAxisName,
+        MeshAxisName as _MeshAxisName,
         ParallelDims,
         SpmdLayout,
     )
@@ -35,13 +37,17 @@
     "maybe_set_sparse_mesh",
     "plain_tensor_to_dtensor_state_dict",
     "spmd_dense_mesh",
+    "spmd_dense_storage_mesh",
     "spmd_sparse_mesh",
+    "spmd_sparse_storage_mesh",
     "spmd_mesh_size",
     "spmd_distribute_tensor",
     "spmd_redistribute_per_axis",
     "spmd_validate_redistributions",
     "set_current_spmd_mesh",
     "set_spmd_meshes",
+    "device_mesh_from_spmd_axes",
+    "get_mesh_pg",
 ]
 
 
@@ -51,8 +57,8 @@
 def plain_tensor_to_dtensor_state_dict(
     state_dict: dict[str, Any],
     *,
-    state_dict_layouts: Mapping[str, spmd.SpmdType],
-    parallel_dims: "ParallelDims",
+    state_dict_layouts: Mapping[str, spmd.SpmdType] | None = None,
+    parallel_dims: "ParallelDims | None" = None,
 ) -> dict[str, Any]:
     """Represent plain local state tensors as DTensors for state transfer."""
     from torchtitan.distributed.parallel_dims import layout_axes, unfold_dp_axes
@@ -64,20 +70,74 @@ def plain_tensor_to_dtensor_state_dict(
             if not isinstance(target, torch.Tensor) or isinstance(target, DTensor):
                 continue
 
-            layout = state_dict_layouts.get(name)
+            layout = state_dict_layouts.get(name) if state_dict_layouts else None
             if layout is None:
-                raise KeyError(f"{name} is missing SPMD layout metadata")
+                if not spmd.has_local_type(target):
+                    raise KeyError(f"{name} is missing SPMD layout metadata")
+                layout = spmd.SpmdType(
+                    dict(spmd.get_local_type(target)),
+                    spmd.get_partition_spec(target),
+                )
 
-            mesh = parallel_dims.get_activated_mesh(
-                unfold_dp_axes(layout_axes(layout))
-            )
-            if mesh is None:
-                continue
+            if parallel_dims is None:
+                mesh = device_mesh_from_spmd_axes(layout.local_type, storage_mesh=True)
+                if mesh is None:
+                    raise ValueError(f"No storage mesh for {name}")
+                assert mesh.mesh_dim_names is not None
+                mesh_axis_names = mesh.mesh_dim_names
+                if layout.partition_spec is not None:
+                    with spmd.set_current_mesh(mesh):
+                        spec = spmd.normalize_partition_spec(layout.partition_spec)
+                    shard_axes = [
+                        spmd.normalize_axis(axis)
+                        for entry in spec
+                        if entry is not None
+                        for axis in (entry if isinstance(entry, tuple) else (entry,))
+                    ]
+                    axis_names = {
+                        spmd.MeshAxis.of(get_mesh_pg(axis_name, mesh=mesh)): axis_name
+                        for axis_name in mesh_axis_names
+                    }
+                    ordered_names = [axis_names[axis] for axis in shard_axes]
+                    ordered_names.extend(
+                        axis_name
+                        for axis_name in mesh_axis_names
+                        if axis_name not in ordered_names
+                    )
+                    if tuple(ordered_names) != mesh_axis_names:
+                        permutation = [
+                            mesh_axis_names.index(axis_name)
+                            for axis_name in ordered_names
+                        ]
+                        mesh = DeviceMesh.from_group(
+                            [mesh.get_group(axis_name) for axis_name in ordered_names],
+                            mesh.device_type,
+                            mesh.mesh.permute(permutation).contiguous(),
+                            mesh_dim_names=tuple(ordered_names),
+                        )
+                        mesh_axis_names = tuple(ordered_names)
+                placements = []
+                for axis_name in mesh_axis_names:
+                    axis = get_mesh_pg(axis_name, mesh=mesh)
+                    shard = partition_spec_get_shard(layout.partition_spec, axis)
+                    axis_type = (
+                        spmd.S(shard.dim)
+                        if shard is not None
+                        else layout.local_type.get(spmd.MeshAxis.of(axis), spmd.R)
+                    )
+                    placements.append(spmd.spmd_type_to_dtensor_placement(axis_type))
+            else:
+                mesh = parallel_dims.get_activated_mesh(
+                    unfold_dp_axes(layout_axes(layout))
+                )
+                if mesh is None:
+                    continue
+                placements = resolve_placements(layout, mesh)
 
             dtensor_state_dict[name] = DTensor.from_local(
                 target,
                 mesh,
-                resolve_placements(layout, mesh),
+                tuple(placements),
                 run_check=False,
             )
     return dtensor_state_dict
@@ -93,14 +153,38 @@ def dtensor_to_plain_tensor_state_dict(
     }
 
 
+def get_mesh_pg(
+    axis_name: str,
+    *,
+    mesh: DeviceMesh | None = None,
+) -> dist.ProcessGroup:
+    """Return an axis PG through traceable DeviceMesh operations."""
+    mesh = mesh if mesh is not None else current_spmd_mesh()
+    if mesh is None:
+        raise ValueError("No SPMD mesh is active.")
+    mesh_axis_names = mesh.mesh_dim_names
+    assert mesh_axis_names is not None, "DeviceMesh must have named axes"
+    if axis_name not in mesh_axis_names:
+        raise ValueError(f"SPMD mesh axis {axis_name!r} is not active.")
+    if torch.compiler.config.compile_on_one_rank:
+        axis_index = mesh_axis_names.index(axis_name)
+        axis_mesh = torch.ops.device_mesh._get_submesh(mesh, [axis_index])
+        return torch.ops._dtensor.mesh_get_process_group.default(axis_mesh, 0)
+    return mesh.get_group(axis_name)
+
+
 def set_spmd_meshes(
     *,
     dense_mesh: DeviceMesh,
+    dense_storage_mesh: DeviceMesh,
     sparse_mesh: DeviceMesh | None,
+    sparse_storage_mesh: DeviceMesh | None,
 ) -> None:
-    """Register the SPMD meshes for dense and sparse runtime regions."""
+    """Register SPMD compute and storage meshes."""
     _MESH_TLS.dense_mesh = dense_mesh
+    _MESH_TLS.dense_storage_mesh = dense_storage_mesh
     _MESH_TLS.sparse_mesh = sparse_mesh
+    _MESH_TLS.sparse_storage_mesh = sparse_storage_mesh
 
 
 def spmd_dense_mesh() -> DeviceMesh:
@@ -110,11 +194,52 @@ def spmd_dense_mesh() -> DeviceMesh:
     return mesh
 
 
+def spmd_dense_storage_mesh() -> DeviceMesh:
+    """Return the registered dense parameter-storage mesh."""
+    mesh = getattr(_MESH_TLS, "dense_storage_mesh", None)
+    assert mesh is not None, "SPMD dense storage mesh has not been registered"
+    return mesh
+
+
 def spmd_sparse_mesh() -> DeviceMesh | None:
     """Return the registered sparse SPMD mesh, if EP is enabled."""
     return getattr(_MESH_TLS, "sparse_mesh", None)
 
 
+def spmd_sparse_storage_mesh() -> DeviceMesh | None:
+    """Return the registered sparse parameter-storage mesh, if EP is enabled."""
+    return getattr(_MESH_TLS, "sparse_storage_mesh", None)
+
+
+def device_mesh_from_spmd_axes(
+    axes: Iterable[DeviceMeshAxis],
+    *,
+    storage_mesh: bool,
+) -> DeviceMesh | None:
+    """Return the registered compute or storage mesh matching the SPMD axes."""
+    normalized_axes = spmd.normalize_mesh(
+        frozenset(spmd.normalize_axis(axis) for axis in axes)
+    )
+    meshes = (
+        (spmd_dense_storage_mesh(), spmd_sparse_storage_mesh())
+        if storage_mesh
+        else (spmd_dense_mesh(), spmd_sparse_mesh())
+    )
+    for mesh in meshes:
+        if mesh is None:
+            continue
+        assert mesh.mesh_dim_names is not None
+        mesh_axes = spmd.normalize_mesh(
+            frozenset(
+                spmd.MeshAxis.of(mesh.get_group(axis_name))
+                for axis_name in mesh.mesh_dim_names
+            )
+        )
+        if normalized_axes == mesh_axes:
+            return mesh
+    return None
+
+
 def _spmd_mesh_stack() -> list[DeviceMesh | None]:
     stack = getattr(_MESH_TLS, "mesh_stack", None)
     if stack is None:
@@ -251,12 +376,12 @@ def spmd_validate_redistributions(sharding_config: Any) -> None:
     from torchtitan.distributed.parallel_dims import MeshAxisName
 
     def _normalize_partition_spec(
-        axis_types: Mapping["MeshAxisName", spmd.PerMeshAxisSpmdType],
+        axis_types: Mapping[DeviceMeshAxis, spmd.PerMeshAxisSpmdType],
         *,
         ndim: int,
-    ) -> tuple[tuple["MeshAxisName", ...], ...]:
+    ) -> tuple[tuple["_MeshAxisName", ...], ...]:
         """Normalize per-axis-types w/ S(dim) -> PartitionSpec-style tuple."""
-        entries: list[tuple["MeshAxisName", ...]] = [()] * ndim
+        entries: list[tuple["_MeshAxisName", ...]] = [()] * ndim
         for axis_name, axis_type in axis_types.items():
             if not isinstance(axis_type, spmd.Shard):
                 continue
@@ -411,7 +536,7 @@ def spmd_redistribute_per_axis(
             continue
         x = spmd.redistribute(
             x,
-            mesh.get_group(axis),
+            get_mesh_pg(axis, mesh=mesh),
             src=src_t,
             dst=dst_t,
             backward_options={"op_dtype": x.dtype},
@@ -429,8 +554,11 @@ def spmd_distribute_tensor(
     Direct ``S(dim)`` layouts are applied per axis. For ``V + PartitionSpec``
     layouts, raw PartitionSpec tuple order controls repeated sharding of the
     same tensor dim, e.g. ``(DP, CP)`` means shard by DP, then shard each DP
-    slice by CP.
+    slice by CP. Axes may be logical ``MeshAxisName`` values or concrete
+    ``MeshAxis`` objects from an SPMD annotation.
     """
+    from torchtitan.distributed.parallel_dims import MeshAxisName
+
     if layout.partition_spec is None:
         axis_shard_dims = [
             (axis_name, axis_type.dim)
@@ -450,17 +578,31 @@ def spmd_distribute_tensor(
                 axis_shard_dims.append((axis_name, dim))
 
     assert mesh.mesh_dim_names is not None, "DeviceMesh must have named axes"
-    for axis_name, dim in axis_shard_dims:
-        axis = axis_name.value
-        axis_size = (
-            mesh.size(mesh.mesh_dim_names.index(axis))
-            if axis in mesh.mesh_dim_names
-            else 1
-        )
+    mesh_axis_groups = {
+        spmd.MeshAxis.of(mesh.get_group(axis_name)): mesh.get_group(axis_name)
+        for axis_name in mesh.mesh_dim_names
+    }
+    for axis, dim in axis_shard_dims:
+        if isinstance(axis, spmd.MeshAxis):
+            if axis not in mesh_axis_groups:
+                raise ValueError(f"SPMD annotation axis {axis!r} is not in {mesh=}.")
+            group = mesh_axis_groups[axis]
+            axis_size = axis.size()
+        elif isinstance(axis, MeshAxisName):
+            axis_name = axis.value
+            if axis_name not in mesh.mesh_dim_names:
+                raise ValueError(f"SPMD layout axis {axis_name!r} is not in {mesh=}.")
+            group = mesh.get_group(axis_name)
+            axis_size = mesh.size(mesh.mesh_dim_names.index(axis_name))
+        else:
+            raise TypeError(
+                "SPMD distribution axes must be MeshAxisName or MeshAxis, "
+                f"got {axis!r}."
+            )
         if axis_size > 1:
             tensor = spmd.shard(
                 tensor,
-                mesh.get_group(axis),
+                group,
                 src=spmd.I,
                 dst=spmd.S(dim),
             )
diff --git a/torchtitan/distributed/utils.py b/torchtitan/distributed/utils.py
index b18d5da206..efa3518807 100644
--- a/torchtitan/distributed/utils.py
+++ b/torchtitan/distributed/utils.py
@@ -413,7 +413,9 @@ def context():
 
                 set_spmd_meshes(
                     dense_mesh=parallel_dims.spmd_dense_mesh(),
+                    dense_storage_mesh=parallel_dims.spmd_dense_storage_mesh(),
                     sparse_mesh=parallel_dims.spmd_sparse_mesh(),
+                    sparse_storage_mesh=parallel_dims.spmd_sparse_storage_mesh(),
                 )
 
                 stack.enter_context(set_current_spmd_mesh(spmd_dense_mesh()))
@@ -469,7 +471,7 @@ def init_distributed(
 
     # disable autograd multithreading, to enable TLS DeviceMesh stack for spmd_types backend.
     # this is needed for AC functionality; multi-threaded autograd means BWD threads performing recompute,
-    # cannot access PGs, e.g. current_spmd_mesh().get_group("tp") to perform the collectives they need.
+    # cannot access the SPMD process groups needed for collectives.
     torch.autograd.set_multithreading_enabled(False)
 
     if comm_config.mode in ("fake_backend", "local_tensor"):
@@ -591,6 +593,193 @@ def set_pg_timeouts(
         torch.distributed.set_timeout(timeout, group)
 
 
+def _param_spmd_mesh_and_shard_axes(
+    parameter: torch.Tensor,
+) -> tuple[DeviceMesh, tuple[str, ...]]:
+    """
+    Given a parameter with SPMD type annotations, returns the storage DeviceMesh
+    and mesh axis names that shard it. Used for grouping params in grad norm calculation.
+    """
+    import spmd_types as spmd
+    from spmd_types.types import partition_spec_get_shard
+
+    from torchtitan.distributed.spmd_types import (
+        device_mesh_from_spmd_axes,
+        get_mesh_pg,
+    )
+
+    if not spmd.has_local_type(parameter):
+        raise ValueError("SPMD local parameters must have type annotations.")
+
+    axis_types = spmd.get_local_type(parameter)
+    mesh = device_mesh_from_spmd_axes(axis_types, storage_mesh=True)
+    if mesh is None:
+        raise ValueError(
+            "SPMD parameter annotations do not match a registered storage mesh: "
+            f"{axis_types}"
+        )
+
+    assert mesh.mesh_dim_names is not None
+    partition_spec = spmd.get_partition_spec(parameter)
+    sharded_axes = []
+    for axis_name in mesh.mesh_dim_names:
+        axis = get_mesh_pg(axis_name, mesh=mesh)
+        shard = partition_spec_get_shard(partition_spec, axis)
+        axis_type = spmd.maybe_get_axis_local_type(parameter, axis)
+        if shard is not None:
+            sharded_axes.append(axis_name)
+        elif axis_type is spmd.V or axis_type is spmd.P:
+            raise ValueError(
+                f"Gradient norm does not support {axis_type!r} placement on "
+                f"mesh axis {axis_name!r}."
+            )
+
+    return mesh, tuple(sharded_axes)
+
+
+def _get_spmd_grad_norm_squared(
+    parameters: list[torch.Tensor],
+    foreach: bool | None,
+) -> torch.Tensor:
+    """
+    Compute the global squared L2 norm of local SPMD parameter gradients.
+
+    1) Group gradients by storage (dense/sparse) mesh, then by device, dtype,
+       and placement (sharded axes). For each grad group, the total sum-of-squares
+       (squared L2 norm) is computed per-param, then summed over the group.
+       This produces a scalar for each group, which is either replicated or Partial(sum)
+       on each axis.
+
+       In the general case, this results in each mesh (dense/sparse) having multiple
+       groups with their sum-of-squares scalar, with differing placements, e.g.
+       {fsdp: P, tp: R} for TP-replicated gradient, {fsdp: P, tp: P} for TP-sharded.
+
+    2) For each mesh (dense/sparse) and mesh axis, if any group exists that is Partial
+       on that axis, groups with Replicated sum-of-squares on that axis are normalized
+       to also be Partial for collective coalescing; dividing by mesh_size.
+
+       This mimics the Replicate -> NormPartial implicit conversion that happens in the
+       DTensor path; when `torch.stack()` is called over groups, a 1.0 / sqrt(mesh_size)
+       factor is contributed when necessary.
+
+       This results in all group sum-of-squares being in a "common partial placement".
+
+    3) For each mesh, the group sum-of-squares can be stacked, summed, and all-reduced
+       over partial axes. The results are summed over meshes and returned so the caller
+       can reduce over pipeline stages before applying the final square root.
+    """
+    from torch.utils._foreach_utils import (
+        _device_has_foreach_support,
+        _has_foreach_support,
+    )
+
+    from torchtitan.distributed.spmd_types import get_mesh_pg
+
+    # grad groups by mesh -> (device, dtype, sharded axes) -> list[grad]
+    mesh_groups: dict[
+        DeviceMesh,
+        dict[tuple[torch.device, torch.dtype, tuple[str, ...]], list[torch.Tensor]],
+    ] = {}
+    first_device: torch.device | None = None
+    for parameter in parameters:
+        grad = parameter.grad
+        if grad is None:
+            continue
+        mesh, sharded_axes = _param_spmd_mesh_and_shard_axes(parameter)
+        if first_device is None:
+            first_device = grad.device
+        per_mesh_per_placement_grouped_grads = mesh_groups.setdefault(mesh, {})
+        per_mesh_per_placement_grouped_grads.setdefault(
+            (grad.device, grad.dtype, sharded_axes), []
+        ).append(grad)
+
+    if first_device is None:
+        return torch.tensor(0.0)
+
+    per_mesh_powsums = []
+    for storage_mesh, per_mesh_per_placement_grouped_grads in mesh_groups.items():
+        assert storage_mesh.mesh_dim_names is not None
+        # mesh axes that shard any of the param groups, we need to normalize R->P to avoid multiple allreduces
+        partial_sum_axes = {
+            axis_name
+            for _, _, sharded_axes in per_mesh_per_placement_grouped_grads
+            for axis_name in sharded_axes
+        }
+        group_powsums = []  # sum-of-squares for each grad group
+        for (
+            device,
+            _,
+            sharded_axes,
+        ), grads in per_mesh_per_placement_grouped_grads.items():
+            use_foreach = (foreach is None and _has_foreach_support(grads, device)) or (
+                foreach is True and _device_has_foreach_support(device)
+            )
+            if use_foreach:
+                powsums = torch._foreach_powsum(grads, 2.0)
+            elif foreach:
+                raise RuntimeError(
+                    f"foreach=True was passed, but can't use the foreach API on "
+                    f"{device.type} tensors"
+                )
+            else:
+                powsums = [torch.linalg._powsum(grad, 2.0) for grad in grads]
+
+            # one sum-of-squares scalar for whole group, R->P normalize
+            group_powsum = torch.stack(powsums).sum()
+            for axis_name in partial_sum_axes:
+                if axis_name not in sharded_axes:
+                    group_powsum /= storage_mesh[axis_name].size()
+            group_powsums.append(group_powsum.to(first_device))
+
+        # stack across groups, produce mesh-wide sum-of-squares
+        per_mesh_powsum = torch.stack(group_powsums).sum()
+        for axis_name in storage_mesh.mesh_dim_names:
+            if axis_name in partial_sum_axes:
+                dist.all_reduce(
+                    per_mesh_powsum,
+                    op=dist.ReduceOp.SUM,
+                    group=get_mesh_pg(axis_name, mesh=storage_mesh),
+                )
+        per_mesh_powsums.append(per_mesh_powsum)
+
+    return torch.stack(per_mesh_powsums).sum()
+
+
+@torch.no_grad()
+def clip_grad_norm_spmd_(
+    parameters: torch.Tensor | Iterable[torch.Tensor],
+    max_norm: float,
+    error_if_nonfinite: bool = False,
+    foreach: bool | None = None,
+    pp_mesh: DeviceMesh | None = None,
+) -> torch.Tensor:
+    """Clip local SPMD gradients using their annotated storage layouts."""
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    else:
+        parameters = list(parameters)
+
+    grad_norm_squared = _get_spmd_grad_norm_squared(parameters, foreach)
+    if pp_mesh is not None:
+        from torchtitan.distributed.spmd_types import get_mesh_pg
+
+        dist.all_reduce(
+            grad_norm_squared,
+            op=dist.ReduceOp.SUM,
+            group=get_mesh_pg("pp", mesh=pp_mesh),
+        )
+    total_norm = grad_norm_squared.sqrt()
+
+    if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()):
+        raise RuntimeError(
+            "The total norm of order 2.0 for gradients from `parameters` is "
+            "non-finite."
+        )
+
+    torch.nn.utils.clip_grads_with_norm_(parameters, max_norm, total_norm, foreach)
+    return total_norm
+
+
 @torch.no_grad()
 def clip_grad_norm_(
     parameters: torch.Tensor | Iterable[torch.Tensor],
diff --git a/torchtitan/experiments/graph_trainer/common_utils.py b/torchtitan/experiments/graph_trainer/common_utils.py
index 6ff9f3494f..5a989ba032 100644
--- a/torchtitan/experiments/graph_trainer/common_utils.py
+++ b/torchtitan/experiments/graph_trainer/common_utils.py
@@ -19,10 +19,14 @@
 
 from torchtitan.config import TORCH_DTYPE_MAP, TrainingConfig
 from torchtitan.distributed import ParallelDims
+from torchtitan.distributed.utils import get_spmd_backend
 from torchtitan.experiments.graph_trainer.simple_fsdp import (
-    data_parallel,
+    data_parallel as dtensor_data_parallel,
     MixedPrecisionPolicy,
 )
+from torchtitan.experiments.graph_trainer.simple_fsdp_spmd import (
+    data_parallel as spmd_data_parallel,
+)
 from torchtitan.models.common.attention import ScaledDotProductAttention
 from torchtitan.models.common.decoder import Decoder, TransformerBlock
 from torchtitan.tools.logging import logger
@@ -430,15 +434,20 @@ def apply_simple_fsdp(
     (the routed-expert weights) are separately wrapped on the EDP mesh when expert
     parallelism is enabled.
     """
+    use_spmd_types = get_spmd_backend() == "spmd_types"
+    fsdp_axis = "fsdp"
+    dense_non_dp_mesh = (
+        parallel_dims.get_activated_mesh(["tp"]) if use_spmd_types else None
+    )
     if parallel_dims.dp_replicate_enabled:
         if parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled:
-            dp_mesh_dim_names = ["dp_replicate", "fsdp"]
+            dp_mesh_dim_names = ["dp_replicate", fsdp_axis]
             dp_mode = "hybrid_shard"
         else:
             dp_mesh_dim_names = ["dp_replicate"]
             dp_mode = "replicate"
     else:
-        dp_mesh_dim_names = ["fsdp"]
+        dp_mesh_dim_names = [fsdp_axis]
         dp_mode = "fully_shard"
 
     dp_mesh = parallel_dims.get_mesh(dp_mesh_dim_names)
@@ -455,6 +464,16 @@ def apply_simple_fsdp(
         )
         edp_mesh = parallel_dims.get_optional_mesh(edp_mesh_names)
         assert edp_mesh is not None
+        efsdp_degree = edp_mesh["efsdp"].size()
+        if parallel_dims.dp_replicate_enabled and efsdp_degree > 1:
+            expert_dp_mode = "hybrid_shard"
+        elif parallel_dims.dp_replicate_enabled:
+            edp_mesh = edp_mesh["dp_replicate"]
+            expert_dp_mode = "replicate"
+        else:
+            expert_dp_mode = "fully_shard"
+
+        expert_non_dp_mesh = parallel_dims.get_mesh("ep")
 
         for _, transformer_block in model.layers.items():
             if not isinstance(transformer_block, TransformerBlock):
@@ -464,23 +483,48 @@ def apply_simple_fsdp(
                 continue
             inner_experts = moe.routed_experts.inner_experts
             experts_shard_dim = 0
-            if edp_mesh["efsdp"].size() * parallel_dims.ep > inner_experts.num_experts:
+            if efsdp_degree * parallel_dims.ep > inner_experts.num_experts:
                 experts_shard_dim = 1
 
-            moe.routed_experts.inner_experts = data_parallel(
-                inner_experts,
-                edp_mesh,
-                dp_mode,
-                mp_policy=mp_policy,
-                shard_dim=experts_shard_dim,
-            )
-
-    model = data_parallel(
-        model,
-        dp_mesh,
-        dp_mode,
-        mp_policy=mp_policy,
-    )
+            if use_spmd_types:
+                sparse_storage_mesh = parallel_dims.spmd_sparse_storage_mesh()
+                assert sparse_storage_mesh is not None
+                moe.routed_experts.inner_experts = spmd_data_parallel(
+                    inner_experts,
+                    edp_mesh,
+                    expert_dp_mode,
+                    mp_policy=mp_policy,
+                    shard_dim=experts_shard_dim,
+                    non_dp_mesh=expert_non_dp_mesh,
+                    storage_mesh=sparse_storage_mesh,
+                )
+            else:
+                moe.routed_experts.inner_experts = dtensor_data_parallel(
+                    inner_experts,
+                    edp_mesh,
+                    expert_dp_mode,
+                    mp_policy=mp_policy,
+                    shard_dim=experts_shard_dim,
+                    non_dp_mesh=expert_non_dp_mesh,
+                )
+
+    if use_spmd_types:
+        model = spmd_data_parallel(
+            model,
+            dp_mesh,
+            dp_mode,
+            mp_policy=mp_policy,
+            non_dp_mesh=dense_non_dp_mesh,
+            storage_mesh=parallel_dims.spmd_dense_storage_mesh(),
+        )
+    else:
+        model = dtensor_data_parallel(
+            model,
+            dp_mesh,
+            dp_mode,
+            mp_policy=mp_policy,
+            non_dp_mesh=dense_non_dp_mesh,
+        )
     logger.info(
         "Applied Data Parallel (simple_fsdp) (dp mode=%s) to the model", dp_mode
     )
diff --git a/torchtitan/experiments/graph_trainer/deepseek_v3/parallelize.py b/torchtitan/experiments/graph_trainer/deepseek_v3/parallelize.py
index 4037f74f7b..2017f2b455 100644
--- a/torchtitan/experiments/graph_trainer/deepseek_v3/parallelize.py
+++ b/torchtitan/experiments/graph_trainer/deepseek_v3/parallelize.py
@@ -58,7 +58,11 @@ def parallelize_deepseekv3(
 
     annotate_deepseekv3(model)
 
-    if parallel_dims.tp_enabled or parallel_dims.ep_enabled:
+    if (
+        parallelism.spmd_backend == "spmd_types"
+        or parallel_dims.tp_enabled
+        or parallel_dims.ep_enabled
+    ):
         model.parallelize(parallel_dims)
 
     # Apply simple_fsdp unconditionally. The `fsdp` mesh always exists with a
diff --git a/torchtitan/experiments/graph_trainer/llama3/parallelize.py b/torchtitan/experiments/graph_trainer/llama3/parallelize.py
index e3c015e41d..d0b50f932c 100644
--- a/torchtitan/experiments/graph_trainer/llama3/parallelize.py
+++ b/torchtitan/experiments/graph_trainer/llama3/parallelize.py
@@ -52,8 +52,13 @@ def parallelize_llama(
 
     annotate_llama(model)
 
-    if parallel_dims.tp_enabled:
+    if parallelism.spmd_backend == "spmd_types":
         model.parallelize(parallel_dims)
+    else:
+        if parallel_dims.cp_enabled:
+            apply_cp_to_attention(model, parallel_dims)
+        if parallel_dims.tp_enabled:
+            model.parallelize(parallel_dims)
 
     # Apply simple_fsdp unconditionally. The `fsdp` mesh always exists with a
     # real backend (see ParallelDims._mesh_exist), even at degree 1, so that
diff --git a/torchtitan/experiments/graph_trainer/make_fx_tracer.py b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
index d70b9b7b6c..a6e6c653b0 100644
--- a/torchtitan/experiments/graph_trainer/make_fx_tracer.py
+++ b/torchtitan/experiments/graph_trainer/make_fx_tracer.py
@@ -16,6 +16,7 @@
 import torch.utils._pytree as pytree
 from torch._guards import tracing, TracingContext
 from torch._subclasses import FakeTensorMode
+from torch.distributed.device_mesh import DeviceMesh
 from torch.fx.experimental.proxy_tensor import make_fx
 from torch.fx.traceback import preserve_node_meta
 from torch.nn.utils import stateless
@@ -30,7 +31,15 @@
 # Tensors and make_fx-safe primitives are allowed as pytree leaves in args.
 # Everything else (callables, custom objects) should be registered as pytree
 # nodes/constants or captured in fn's closure.
-_ALLOWED_LEAF_TYPES = (torch.Tensor, int, float, bool, str, type(None))
+_ALLOWED_LEAF_TYPES = (
+    torch.Tensor,
+    DeviceMesh,
+    int,
+    float,
+    bool,
+    str,
+    type(None),
+)
 
 
 @contextmanager
@@ -336,6 +345,7 @@ def minimal_fx_tracer(
     module: nn.Module | None = None,
     optimizer: "torch.optim.Optimizer | None" = None,
     *,
+    precompile_meshes: list[DeviceMesh] | None = None,
     prepare_inputs: Callable[[tuple[Any, ...], dict[str, Any]], None] | None = None,
     prepare_call_inputs: Callable[
         [tuple[Any, ...], dict[str, Any]],
@@ -384,6 +394,9 @@ def minimal_fx_tracer(
     ``record_stack_traces`` controls whether make_fx records Python stack traces
     in node metadata. It defaults to on to preserve the existing debugging
     behavior.
+
+    ``precompile_meshes`` contains DeviceMeshes that must be explicit graph
+    inputs for CooR precompilation.
     """
     _check_optimizer_has_module(module, optimizer)
 
@@ -393,15 +406,17 @@ def _trace_with_args(*args: Any, **kwargs: Any) -> TracedResult:
 
         model_state, optim_state = extract_train_state(module, optimizer)
         state_fqns = list(model_state.keys())
+        trace_meshes = precompile_meshes or []
 
         state_tree = {"model": model_state, "optim": optim_state}
         state_flat, state_spec = pytree.tree_flatten(state_tree)
         num_state_inputs = len(state_flat)
+        num_mesh_inputs = len(trace_meshes)
 
         user_inputs_flat, user_inputs_spec = pytree.tree_flatten((args, kwargs))
 
         # Validate leaves.
-        for leaf in [*state_flat, *user_inputs_flat]:
+        for leaf in [*state_flat, *trace_meshes, *user_inputs_flat]:
             if isinstance(leaf, nn.Module):
                 raise ValueError(
                     "minimal_fx_tracer requires explicit tensor state, not nn.Module "
@@ -411,14 +426,19 @@ def _trace_with_args(*args: Any, **kwargs: Any) -> TracedResult:
             if not isinstance(leaf, _ALLOWED_LEAF_TYPES):
                 raise ValueError(
                     "minimal_fx_tracer requires all pytree leaves in state/args to "
-                    f"be tensors or primitives (int/float/bool/str), got "
+                    f"be tensors, DeviceMeshes, or primitives "
+                    f"(int/float/bool/str), got "
                     f"{type(leaf).__name__}. Non-primitive values should either be "
                     "registered as pytree nodes (register_pytree_node) or constants "
                     f"(pytree.register_constant), or captured in fn's closure."
                 )
 
-        # Combined flat input: [*state, *user_args] with subclasses unwrapped.
-        full_args = list(state_flat) + list(user_inputs_flat)
+        # Combined flat input: train state, precompile meshes, then user inputs.
+        full_args = [
+            *state_flat,
+            *trace_meshes,
+            *user_inputs_flat,
+        ]
         num_full_args = len(full_args)
         for arg in full_args:
             if not isinstance(arg, torch.Tensor):
@@ -451,7 +471,8 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
 
             wrapped = _wrap_subclasses(plain_args, num_full_args, input_layouts)
             state_wrapped = wrapped[:num_state_inputs]
-            user_flat = wrapped[num_state_inputs:]
+            mesh_end = num_state_inputs + num_mesh_inputs
+            user_flat = wrapped[mesh_end:]
 
             state_t = pytree.tree_unflatten(list(state_wrapped), state_spec)
             model_state_t = state_t["model"]
@@ -464,9 +485,12 @@ def fn_with_subclass_handling(*plain_args: Any) -> list:
                 if prepared is not None:
                     user_args, user_kwargs = prepared
 
-            with _reparametrize_train_state(
-                module, optimizer, model_state_t, optim_state_t
-            ), torch.compiler._patch_engine_backward():
+            with (
+                _reparametrize_train_state(
+                    module, optimizer, model_state_t, optim_state_t
+                ),
+                torch.compiler._patch_engine_backward(),
+            ):
                 result = fn(*user_args, **user_kwargs)
 
             flat_outs, output_spec = pytree.tree_flatten(result)
@@ -534,6 +558,7 @@ def run_traced(
     *,
     module: nn.Module | None = None,
     optimizer: "torch.optim.Optimizer | None" = None,
+    precompile_meshes: list[DeviceMesh] | None = None,
     _validate_runtime: bool = False,
     interpreter_cls: type | None = None,
 ) -> Callable[..., Any]:
@@ -571,9 +596,9 @@ def _run(*args: Any, **kwargs: Any) -> Any:
                 f"  Traced: {traced_result.state_fqns}\n"
                 f"  Got:    {list(model_state.keys())}"
             )
+        runtime_meshes = precompile_meshes or []
         state_tree = {"model": model_state, "optim": optim_state}
         state_flat, _ = pytree.tree_flatten(state_tree)
-
         user_inputs_flat, runtime_spec = pytree.tree_flatten((args, kwargs))
         # TODO: pytree's dict flatten preserves insertion order, so kwargs in a
         # different order than trace produce a different spec even though they
@@ -586,13 +611,22 @@ def _run(*args: Any, **kwargs: Any) -> Any:
                 f"trace-time {traced_result.user_inputs_spec}"
             )
         if any(
-            isinstance(leaf, nn.Module) for leaf in [*state_flat, *user_inputs_flat]
+            isinstance(leaf, nn.Module)
+            for leaf in [
+                *state_flat,
+                *runtime_meshes,
+                *user_inputs_flat,
+            ]
         ):
             raise ValueError(
                 "run_traced requires explicit tensor state, not nn.Module instances. "
                 "Capture nn.Modules in fn's closure or pass them via the 'module' kwarg."
             )
-        all_args = list(state_flat) + list(user_inputs_flat)
+        all_args = [
+            *state_flat,
+            *runtime_meshes,
+            *user_inputs_flat,
+        ]
         flat_inputs, _ = _unwrap_subclasses(all_args)
 
         with torch.no_grad():
diff --git a/torchtitan/experiments/graph_trainer/muse_glimmer/parallelize.py b/torchtitan/experiments/graph_trainer/muse_glimmer/parallelize.py
index 1ce7a3c597..23c4d478b9 100644
--- a/torchtitan/experiments/graph_trainer/muse_glimmer/parallelize.py
+++ b/torchtitan/experiments/graph_trainer/muse_glimmer/parallelize.py
@@ -34,7 +34,7 @@ def parallelize_muse_glimmer(
 
     annotate_module_fqns(model)
 
-    if parallel_dims.tp_enabled:
+    if parallelism.spmd_backend == "spmd_types" or parallel_dims.tp_enabled:
         model.parallelize(parallel_dims)
 
     parallelized_model = apply_simple_fsdp(
diff --git a/torchtitan/experiments/graph_trainer/precompile.py b/torchtitan/experiments/graph_trainer/precompile.py
index 3ac445b98c..7117e1a8d9 100644
--- a/torchtitan/experiments/graph_trainer/precompile.py
+++ b/torchtitan/experiments/graph_trainer/precompile.py
@@ -19,6 +19,7 @@
 
 import torch
 import torch.utils._pytree as pytree
+from torch.distributed.device_mesh import DeviceMesh
 
 from torchtitan.experiments.graph_trainer.make_fx_tracer import (
     SubclassLayout,
@@ -30,6 +31,22 @@
 ConfigFingerprint = NewType("ConfigFingerprint", str)
 
 
+def get_spmd_precompile_meshes(parallel_dims: ParallelDims) -> list[DeviceMesh]:
+    """Return the SPMD meshes that must be DeviceMesh graph inputs."""
+    candidates = [
+        parallel_dims.spmd_dense_mesh(),
+        parallel_dims.spmd_dense_storage_mesh(),
+        parallel_dims.spmd_sparse_mesh(),
+        parallel_dims.spmd_sparse_storage_mesh(),
+        parallel_dims.get_optional_mesh("pp"),
+    ]
+    meshes: list[DeviceMesh] = []
+    for mesh in candidates:
+        if mesh is not None and all(mesh is not other for other in meshes):
+            meshes.append(mesh)
+    return meshes
+
+
 def compute_config_fingerprint(
     model: torch.nn.Module,
     compile_config: GraphTrainerCompileConfig,
@@ -86,19 +103,20 @@ def compute_config_fingerprint(
 
 
 def _register_coor_ops() -> None:
-    """Register CooR custom ops required for deserialization.
+    """Register CooR custom ops required for tracing and deserialization.
 
     CooR-compiled artifacts reference custom ops (e.g.
     device_mesh._runtime_compute_coordinate_on_dim) that are lazily
     registered. The ops module uses @torch.library.custom_op with
     DeviceMesh, which requires DeviceMesh to be registered as an
-    opaque type first. Must be called before deserializing any
+    opaque type first. Must be called before tracing or deserializing a
     CooR-compiled artifact.
     """
     from torch.distributed.device_mesh import _register_distributed_opaque_types
 
     _register_distributed_opaque_types()
     from torch.distributed._ops import device_mesh as _dm_ops  # noqa: F401
+    from torch.distributed.tensor import _collective_utils  # noqa: F401
 
 
 def _validate_config_fingerprint(
diff --git a/torchtitan/experiments/graph_trainer/precompile_main.py b/torchtitan/experiments/graph_trainer/precompile_main.py
index e5383b754a..9507e58e8a 100644
--- a/torchtitan/experiments/graph_trainer/precompile_main.py
+++ b/torchtitan/experiments/graph_trainer/precompile_main.py
@@ -36,7 +36,10 @@
 from torchtitan.experiments.graph_trainer.memory_policy import (
     validate_memory_policy_config,
 )
-from torchtitan.experiments.graph_trainer.precompile import _FX_TRACE_ARTIFACT_KEY
+from torchtitan.experiments.graph_trainer.precompile import (
+    _FX_TRACE_ARTIFACT_KEY,
+    _register_coor_ops,
+)
 from torchtitan.experiments.graph_trainer.storage import DiskStorageAdapter
 from torchtitan.models.common.attention import FlexAttention, VarlenAttention
 from torchtitan.models.common.decoder import Decoder
@@ -86,6 +89,7 @@ def _common_setup(config):
     import torch.distributed.config as dist_config
 
     dist_config.compile_on_one_rank = True
+    _register_coor_ops()
 
     # Match the deterministic mode that the training loop will use.
     # The backward graph captures use_deterministic_algorithms() at
@@ -150,7 +154,12 @@ def _common_setup(config):
     model.to_empty(device=device_type)
     dist_config.compile_on_one_rank = False
     try:
-        with torch.no_grad():
+        with (
+            torch.no_grad(),
+            dist_utils.get_spmd_context(
+                parallel_dims=parallel_dims,
+            )(),
+        ):
             model.init_weights(buffer_device=None)
     finally:
         dist_config.compile_on_one_rank = True
@@ -198,6 +207,7 @@ def _precompile_aot_fx_trace(
     from torchtitan.experiments.graph_trainer.make_fx_tracer import minimal_fx_tracer
     from torchtitan.experiments.graph_trainer.precompile import (
         compute_config_fingerprint,
+        get_spmd_precompile_meshes,
         precompile_fx_trace_save,
     )
     from torchtitan.experiments.graph_trainer.trainer import make_fwd_bwd_step
@@ -282,10 +292,22 @@ def prepare_trace_call_inputs(
         return args, kwargs
 
     logger.info("Tracing fwd+loss+bwd via make_fx...")
-    with loss_parallel_ctx:
+    spmd_context = dist_utils.get_spmd_context(
+        parallel_dims=parallel_dims,
+        spmd_typechecking=(
+            config.parallelism.spmd_backend == "spmd_types"
+            and config.debug.spmd_typechecking
+        ),
+    )
+    with loss_parallel_ctx, spmd_context():
         traced_result = minimal_fx_tracer(
             fwd_bwd_fn,
             module=model,
+            precompile_meshes=(
+                get_spmd_precompile_meshes(parallel_dims)
+                if config.parallelism.spmd_backend == "spmd_types"
+                else None
+            ),
             prepare_inputs=prepare_trace_inputs,
             prepare_call_inputs=prepare_trace_call_inputs,
         )(dummy_inputs, dummy_labels, dummy_global_valid_tokens, extra_kwargs)
diff --git a/torchtitan/experiments/graph_trainer/qwen3/parallelize.py b/torchtitan/experiments/graph_trainer/qwen3/parallelize.py
index 72ee7f6b9b..2a9a9043cb 100644
--- a/torchtitan/experiments/graph_trainer/qwen3/parallelize.py
+++ b/torchtitan/experiments/graph_trainer/qwen3/parallelize.py
@@ -63,7 +63,11 @@ def parallelize_qwen3(
 
     annotate_qwen3(model)
 
-    if parallel_dims.tp_enabled or parallel_dims.ep_enabled:
+    if (
+        parallelism.spmd_backend == "spmd_types"
+        or parallel_dims.tp_enabled
+        or parallel_dims.ep_enabled
+    ):
         model.parallelize(parallel_dims)
 
     # Apply simple_fsdp unconditionally. The `fsdp` mesh always exists with a
diff --git a/torchtitan/experiments/graph_trainer/simple_fsdp.py b/torchtitan/experiments/graph_trainer/simple_fsdp.py
index c108da4f25..5b52ccd622 100644
--- a/torchtitan/experiments/graph_trainer/simple_fsdp.py
+++ b/torchtitan/experiments/graph_trainer/simple_fsdp.py
@@ -9,9 +9,11 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 
+import spmd_types as spmd
 import torch
 import torch.nn as nn
 
+from spmd_types.types import partition_spec_get_shard
 from torch.distributed._tensor import (
     distribute_tensor,
     DTensor,
@@ -24,6 +26,7 @@
 from torch.distributed.tensor._redistribute import redistribute_local_tensor
 from torch.distributed.tensor.placement_types import _StridedShard, Placement
 
+from torchtitan.distributed.utils import get_spmd_backend
 from torchtitan.protocols.module import Module
 
 _active_parametrization = True
@@ -39,12 +42,47 @@ def disable_active_parametrization() -> Generator[None, None, None]:
         _active_parametrization = True
 
 
+def is_active_parametrization() -> bool:
+    return _active_parametrization
+
+
 @dataclass(frozen=True)
 class MixedPrecisionPolicy:
     param_dtype: torch.dtype | None = None
     reduce_dtype: torch.dtype | None = None
 
 
+def _spmd_local_tensor_to_dtensor(
+    tensor: torch.Tensor,
+    non_dp_mesh: DeviceMesh | None,
+) -> torch.Tensor:
+    """Reconstruct model-parallel DTensor metadata from an SPMD local tensor."""
+    if (
+        get_spmd_backend() != "spmd_types"
+        or not spmd.has_local_type(tensor)
+        or non_dp_mesh is None
+    ):
+        return tensor
+
+    assert non_dp_mesh.mesh_dim_names is not None
+    partition_spec = spmd.get_partition_spec(tensor)
+    with spmd.set_current_mesh(non_dp_mesh):
+        placements = tuple(
+            spmd.spmd_type_to_dtensor_placement(
+                partition_spec_get_shard(partition_spec, axis_name)
+                or spmd.get_axis_local_type(tensor, axis_name)
+            )
+            for axis_name in non_dp_mesh.mesh_dim_names
+        )
+
+    return DTensor.from_local(
+        tensor,
+        non_dp_mesh,
+        placements,
+        run_check=False,
+    )
+
+
 def _distribute_dtensor(
     tensor: DTensor,
     device_mesh: DeviceMesh,
@@ -142,9 +180,12 @@ def _register_parametrization(
     TODO: In checkpoint saving/loading, avoid parametrization calls when calling
     get_model_state_dict func in torchtitan/components/checkpointer/dcp.py.
     """
+    object.__setattr__(module, "_simple_fsdp_parametrization", parametrization)
     param_name_to_property = {
         param_name: property(
-            lambda self, pn=param_name: parametrization(self._parameters[pn])
+            lambda self, pn=param_name: self._simple_fsdp_parametrization(
+                self._parameters[pn], pn
+            )
         )
         for param_name in param_names
     }
@@ -219,8 +260,12 @@ def replicate_compute(self, x: DTensor) -> torch.Tensor:
             )
             non_dp_mesh = x._spec.mesh[non_dp_mesh_dim_names]
 
-            output = DTensor.from_local(
-                replicated_local_tensor, non_dp_mesh, non_dp_placements
+            output = (
+                replicated_local_tensor
+                if get_spmd_backend() == "spmd_types"
+                else DTensor.from_local(
+                    replicated_local_tensor, non_dp_mesh, non_dp_placements
+                )
             )
         elif non_dp_mesh_dims == 0:
             output = x.redistribute(
@@ -236,18 +281,16 @@ def replicate_compute(self, x: DTensor) -> torch.Tensor:
 
         return output
 
-    def forward(self, x: DTensor) -> torch.Tensor:
-        global _active_parametrization
+    def forward(self, x: DTensor, _param_name: str) -> torch.Tensor:
         # This should never be set to true during forward, only outside for model
         # inspection / debugging / initialization
         # model initialization can be done now through
         # with disable_active_parametrization():
         #     model.init_states()
-        if not _active_parametrization:
+        if not is_active_parametrization():
             return x
 
-        output = self.replicate_compute(x)
-        return output
+        return self.replicate_compute(x)
 
 
 def data_parallel(
@@ -256,6 +299,7 @@ def data_parallel(
     mode: str = "replicate",
     mp_policy: MixedPrecisionPolicy | None = None,
     shard_dim: int = 0,
+    non_dp_mesh: DeviceMesh | None = None,
 ) -> nn.Module:
     param_sharding: tuple[Placement, ...]
     if mode == "replicate":
@@ -272,7 +316,6 @@ def data_parallel(
         raise ValueError(f"Unsupported mode {mode}")
 
     modules = list(model.modules())
-
     for mod in modules:
         params_dict = dict(mod.named_parameters(recurse=False))
         # we shouldn't apply data parallel to the modules that are already
@@ -282,6 +325,7 @@ def data_parallel(
 
         for p_name, p in params_dict.items():
             if p is not None and p.numel() > 0:
+                p = _spmd_local_tensor_to_dtensor(p, non_dp_mesh)
                 distribute_tensor_func = (
                     _distribute_dtensor if isinstance(p, DTensor) else distribute_tensor
                 )
diff --git a/torchtitan/experiments/graph_trainer/simple_fsdp_spmd.py b/torchtitan/experiments/graph_trainer/simple_fsdp_spmd.py
new file mode 100644
index 0000000000..5e7138c27d
--- /dev/null
+++ b/torchtitan/experiments/graph_trainer/simple_fsdp_spmd.py
@@ -0,0 +1,231 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import TypeAlias
+
+import spmd_types as spmd
+import torch
+import torch.nn as nn
+
+from spmd_types.checker import typecheck
+from torch.distributed.device_mesh import DeviceMesh
+
+from torchtitan.distributed.spmd_types import get_mesh_pg
+from torchtitan.experiments.graph_trainer.simple_fsdp import (
+    _register_parametrization,
+    is_active_parametrization,
+    MixedPrecisionPolicy,
+)
+from torchtitan.protocols.module import Module
+
+
+StorageTimePlacement: TypeAlias = tuple[
+    DeviceMesh,
+    str,
+    spmd.PerMeshAxisSpmdType,
+]
+
+
+def _shard_param_for_fsdp_storage(
+    param: nn.Parameter,
+    fsdp_mesh: DeviceMesh,
+    storage_mesh: DeviceMesh,
+    shard_axis_name: str,
+    shard_dim: int,
+) -> nn.Parameter:
+    if not -param.ndim <= shard_dim < param.ndim:
+        raise ValueError(
+            f"Cannot shard parameter with {param.ndim} dimensions on tensor "
+            f"dim {shard_dim}."
+        )
+    canonical_shard_dim = shard_dim % param.ndim
+    dim_size = param.size(canonical_shard_dim)
+    shard_degree = fsdp_mesh[shard_axis_name].size()
+    if dim_size % shard_degree != 0:
+        raise ValueError(
+            f"Cannot evenly shard parameter shape {tuple(param.shape)} on tensor "
+            f"dim {canonical_shard_dim} across {shard_degree} ranks."
+        )
+
+    # track GSPMD annotation on param, so weight init can read & feed correct shard
+    with typecheck(local=False):
+        fsdp_group = get_mesh_pg(shard_axis_name, mesh=storage_mesh)
+        fsdp_axis = spmd.MeshAxis.of(fsdp_group)
+        param_for_storage = param.detach()
+        spmd.mutate_type(
+            param_for_storage,
+            fsdp_axis,
+            src=spmd.R,
+            dst=spmd.I,
+        )
+        local_tensor = spmd.redistribute(
+            param_for_storage,
+            fsdp_group,
+            src=spmd.I,
+            dst=spmd.S(canonical_shard_dim),
+        )
+
+    local_param = nn.Parameter(local_tensor, requires_grad=param.requires_grad)
+    spmd.assert_type_like(local_param, local_tensor)
+    return local_param
+
+
+def _get_non_dp_storage_time_placements(
+    param: nn.Parameter,
+    non_dp_mesh: DeviceMesh | None,
+    storage_mesh: DeviceMesh,
+) -> tuple[StorageTimePlacement, ...]:
+    """
+    Return list of non-DP mesh axes requiring I->R convert in FWD,
+    translating to P->I all-reduce in BWD.
+    """
+    if non_dp_mesh is None:
+        return ()
+    if not spmd.has_local_type(param):
+        raise ValueError(
+            "Parameters must have SPMD layouts before applying SimpleFSDP "
+            "with a non_dp_mesh."
+        )
+
+    assert non_dp_mesh.mesh_dim_names is not None
+    return tuple(
+        (storage_mesh, axis_name, spmd.I)
+        for axis, axis_name in enumerate(non_dp_mesh.mesh_dim_names)
+        if non_dp_mesh.size(axis) > 1
+        and spmd.get_axis_local_type(
+            param,
+            get_mesh_pg(axis_name, mesh=non_dp_mesh),
+        )
+        is spmd.R
+    )
+
+
+class ReplicateComputation(Module):
+    """Materialize parameter storage for replicated compute."""
+
+    def __init__(
+        self,
+        storage_time_placements_by_param: dict[str, tuple[StorageTimePlacement, ...]],
+        mp_policy: MixedPrecisionPolicy | None,
+    ) -> None:
+        super().__init__()
+        self.storage_time_placements_by_param = storage_time_placements_by_param
+        mp_policy = mp_policy or MixedPrecisionPolicy()
+        self.param_dtype = mp_policy.param_dtype
+        self.reduce_dtype = mp_policy.reduce_dtype
+
+    def forward(self, param: torch.Tensor, param_name: str) -> torch.Tensor:
+        if not is_active_parametrization():
+            return param
+
+        result = param
+        for mesh, axis_name, storage_type in self.storage_time_placements_by_param[
+            param_name
+        ]:
+            result = spmd.redistribute(
+                result,
+                get_mesh_pg(axis_name, mesh=mesh),
+                src=storage_type,
+                dst=spmd.R,
+                op_dtype=self.param_dtype,
+                backward_options={"op_dtype": self.reduce_dtype},
+            )
+        return result
+
+
+def data_parallel(
+    model: nn.Module,
+    fsdp_mesh: DeviceMesh,
+    mode: str = "replicate",
+    mp_policy: MixedPrecisionPolicy | None = None,
+    shard_dim: int = 0,
+    non_dp_mesh: DeviceMesh | None = None,
+    *,
+    storage_mesh: DeviceMesh,
+) -> nn.Module:
+    """Apply local-tensor SPMD data parallelism to ``model``."""
+    if fsdp_mesh.mesh_dim_names is None:
+        raise ValueError("fsdp_mesh must have named axes.")
+    fsdp_axis_names = fsdp_mesh.mesh_dim_names
+    if non_dp_mesh is not None and non_dp_mesh.mesh_dim_names is None:
+        raise ValueError("non_dp_mesh must have named axes.")
+
+    # configure FSDP storage mesh axes, placements
+    if mode == "replicate":
+        storage_shard_axis_name = None
+        fsdp_storage_types = (spmd.I,) * fsdp_mesh.ndim
+    elif mode == "fully_shard":
+        if fsdp_mesh.ndim != 1:
+            raise ValueError("fully_shard requires a one-dimensional fsdp_mesh.")
+        shard_axis_name = fsdp_axis_names[0]
+        storage_shard_axis_name = shard_axis_name
+        fsdp_storage_types = (spmd.S(shard_dim),)
+    elif mode == "hybrid_shard":
+        if fsdp_mesh.ndim != 2:
+            raise ValueError("hybrid_shard requires a two-dimensional fsdp_mesh.")
+        shard_axis_name = fsdp_axis_names[1]
+        storage_shard_axis_name = shard_axis_name
+        fsdp_storage_types = (spmd.I, spmd.S(shard_dim))
+    else:
+        raise ValueError(f"Unsupported mode {mode!r}.")
+
+    fsdp_storage_time_placements = tuple(
+        (storage_mesh, axis_name, storage_type)
+        for axis_name, storage_type in zip(
+            fsdp_axis_names,
+            fsdp_storage_types,
+            strict=True,
+        )
+    )
+
+    modules = list(model.modules())
+    for module in modules:
+        params = dict(module.named_parameters(recurse=False))
+        if "SimpleFSDP" in module.__class__.__name__:
+            continue
+
+        # for compute-time / pre-forward: besides FSDP axis unshard, collect params requiring
+        # model-parallel axis I->R convert calls, so FSDP can handle BWD all-reduces.
+        storage_time_placements_by_param = {}
+        for param_name, param in params.items():
+            if param is None:
+                continue
+            storage_axes = [
+                fsdp_mesh.get_group(axis_name) for axis_name in fsdp_axis_names
+            ]
+            if non_dp_mesh is not None:
+                storage_axes.extend(
+                    non_dp_mesh.get_group(axis_name)
+                    for axis_name in non_dp_mesh.mesh_dim_names
+                )
+            spmd.reinterpret_mesh(param, storage_axes, inplace=True)
+            non_dp_storage_time_placements = _get_non_dp_storage_time_placements(
+                param, non_dp_mesh, storage_mesh
+            )
+            storage_time_placements_by_param[param_name] = (
+                fsdp_storage_time_placements + non_dp_storage_time_placements
+            )
+            if storage_shard_axis_name is not None and param.numel() > 0:
+                module.register_parameter(
+                    param_name,
+                    _shard_param_for_fsdp_storage(
+                        param,
+                        fsdp_mesh,
+                        storage_mesh,
+                        storage_shard_axis_name,
+                        shard_dim,
+                    ),
+                )
+
+        _register_parametrization(
+            module,
+            list(params),
+            ReplicateComputation(
+                storage_time_placements_by_param,
+                mp_policy,
+            ),
+        )
+    return model
diff --git a/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py b/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py
index dd9379580c..ee65a6e978 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py
@@ -13,6 +13,7 @@
 
 from torchtitan.config.configs import TrainingConfig
 from torchtitan.distributed import ParallelDims
+from torchtitan.distributed.utils import set_spmd_backend
 from torchtitan.experiments.graph_trainer.common_utils import apply_simple_fsdp
 
 
@@ -40,6 +41,7 @@ def test_param_cast_to_bf16_at_ngpu_1(self):
         simple_fsdp wrap, parameters silently stay in fp32 on a single GPU and
         any downstream bf16-only kernel (e.g. MXFP8) breaks.
         """
+        set_spmd_backend("partial_dtensor")
         parallel_dims = ParallelDims(
             dp_replicate=1,
             dp_shard=1,
diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py
index 63afaab6a4..5a43672a97 100644
--- a/torchtitan/experiments/graph_trainer/trainer.py
+++ b/torchtitan/experiments/graph_trainer/trainer.py
@@ -8,9 +8,13 @@
 from dataclasses import dataclass, field
 from typing import Any
 
+import spmd_types as spmd
 import torch
 import torch.nn as nn
 
+from torchtitan.components.checkpointer import CheckpointManager
+from torchtitan.components.checkpointer.utils import canonical_fqn
+from torchtitan.distributed import utils as dist_utils
 from torchtitan.experiments.graph_trainer.common_utils import (
     accumulate_param_grads_,
     compute_annotated_loss,
@@ -41,10 +45,42 @@
     TRACE_CALL_INPUT_PREPARERS,
     TRACE_INPUT_PREPARERS,
 )
+from torchtitan.protocols.state_dict_adapter import PlainToDTensorStateDictAdapter
 from torchtitan.tools.logging import logger
 from torchtitan.trainer import Trainer
 
 
+def _get_state_dict_layouts(model_parts: list[nn.Module]):
+    layouts = {}
+    exposed_keys = {key for model in model_parts for key in model.state_dict()}
+    for model in model_parts:
+        for name, tensor in (
+            *model.named_parameters(remove_duplicate=False),
+            *model.named_buffers(remove_duplicate=False),
+        ):
+            if not spmd.has_local_type(tensor):
+                continue
+            name = canonical_fqn(name)
+            layout = spmd.SpmdType(
+                dict(spmd.get_local_type(tensor)), spmd.get_partition_spec(tensor)
+            )
+            names = [name]
+            if name.endswith(("wqkv.weight", "wqkv.bias")):
+                prefix, kind = name.rsplit("wqkv.", 1)
+                names = [f"{prefix}{proj}.{kind}" for proj in ("wq", "wk", "wv")]
+            elif name.endswith("w13"):
+                prefix = name[: -len("w13")]
+                names = [
+                    f"{prefix}{key}"
+                    for key in ("w1.weight", "w3.weight", "w1_EFD", "w3_EFD")
+                ]
+            layouts.update(dict.fromkeys(names, layout))
+    missing = exposed_keys - layouts.keys()
+    if missing:
+        raise ValueError(f"Missing SPMD checkpoint layouts: {sorted(missing)}")
+    return layouts
+
+
 def _maybe_apply_numa_binding(device_index: int, device_type: str) -> None:
     """Pin this process to the NUMA node of its GPU for local memory bandwidth.
 
@@ -111,6 +147,13 @@ class Config(Trainer.Config):
     def __init__(self, config):
         super().__init__(config)
 
+        if config.parallelism.spmd_backend == "spmd_types" and self.checkpointer.enable:
+            if not isinstance(self.checkpointer, CheckpointManager):
+                raise ValueError("GraphTrainer SPMD checkpoints require DCP")
+            self.checkpointer.state_dict_adapter = PlainToDTensorStateDictAdapter(
+                _get_state_dict_layouts(self.model_parts)
+            )
+
         validate_memory_policy_config(self.config.compile)
 
         _maybe_apply_numa_binding(self.device.index, self.device.type)
@@ -131,6 +174,16 @@ def __init__(self, config):
         # Run post-init hook for the active pass pipeline
         POST_INIT_HOOKS.get(self.config.compile.pass_pipeline, lambda _: None)(self)
 
+    def _clip_grad_norm(self, parameters: list[torch.Tensor]) -> torch.Tensor:
+        if self.config.parallelism.spmd_backend != "spmd_types":
+            return super()._clip_grad_norm(parameters)
+        return dist_utils.clip_grad_norm_spmd_(
+            parameters,
+            self.config.training.max_norm,
+            foreach=True,
+            pp_mesh=self.parallel_dims.get_optional_mesh("pp"),
+        )
+
     def forward_backward_step(
         self,
         *,
@@ -241,7 +294,21 @@ def _make_fx_forward_backward_step(
                     compile_config=self.config.compile,
                 )
         with self.train_context():
-            outputs = run_traced(self._traced_step, module=model)(
+            precompile_meshes = None
+            if (
+                self.config.compile.precompile_artifact_dir
+                and self.config.parallelism.spmd_backend == "spmd_types"
+            ):
+                from torchtitan.experiments.graph_trainer.precompile import (
+                    get_spmd_precompile_meshes,
+                )
+
+                precompile_meshes = get_spmd_precompile_meshes(self.parallel_dims)
+            outputs = run_traced(
+                self._traced_step,
+                module=model,
+                precompile_meshes=precompile_meshes,
+            )(
                 inputs,
                 labels,
                 global_valid_tokens,
diff --git a/torchtitan/experiments/rl/actors/trainer.py b/torchtitan/experiments/rl/actors/trainer.py
index a76e07d610..92fef990a4 100644
--- a/torchtitan/experiments/rl/actors/trainer.py
+++ b/torchtitan/experiments/rl/actors/trainer.py
@@ -309,7 +309,12 @@ def _build_model(
         )
 
         model.to_empty(device=device_type)
-        with torch.no_grad():
+        with (
+            torch.no_grad(),
+            dist_utils.get_spmd_context(
+                parallel_dims=self.parallel_dims,
+            )(),
+        ):
             model.init_weights(buffer_device=None)
 
         return model
diff --git a/torchtitan/experiments/rl/models/vllm_wrapper.py b/torchtitan/experiments/rl/models/vllm_wrapper.py
index c69280ba63..5f58176f2e 100644
--- a/torchtitan/experiments/rl/models/vllm_wrapper.py
+++ b/torchtitan/experiments/rl/models/vllm_wrapper.py
@@ -30,7 +30,6 @@
 )
 from torchtitan.distributed import utils as dist_utils
 from torchtitan.distributed.parallel_dims import ParallelDims
-from torchtitan.distributed.spmd_types import current_spmd_mesh
 from torchtitan.distributed.utils import is_in_batch_invariant_mode
 from torchtitan.experiments.rl.models.vllm_registry import InferenceParallelismConfig
 from torchtitan.models.common.attention import FusedQKVLinear
@@ -470,11 +469,9 @@ def compute_logits(
             # full local logits tensor that vLLM expects.
             if self.parallel_dims.tp_enabled:
                 if self.parallel_dims.spmd_backend == "spmd_types":
-                    mesh = current_spmd_mesh()
-                    assert mesh is not None
                     logits = spmd.redistribute(
                         logits,
-                        mesh.get_group("tp"),
+                        "tp",
                         src=spmd.S(-1),
                         dst=spmd.R,
                         backward_options={"op_dtype": logits.dtype},
diff --git a/torchtitan/models/common/config_utils.py b/torchtitan/models/common/config_utils.py
index 1cbc09a02d..daeeb7c494 100644
--- a/torchtitan/models/common/config_utils.py
+++ b/torchtitan/models/common/config_utils.py
@@ -15,9 +15,13 @@
 from typing import Literal
 
 import torch
+from spmd_types import get_local_type, get_partition_spec, has_local_type, SpmdType
 from torch.distributed.tensor import DTensor
 
-from torchtitan.distributed.spmd_types import current_spmd_mesh, spmd_mesh_size
+from torchtitan.distributed.spmd_types import (
+    device_mesh_from_spmd_axes,
+    spmd_distribute_tensor,
+)
 from torchtitan.models.common.attention import (
     FlexAttention,
     FusedQKVLinear,
@@ -163,15 +167,26 @@ def _init(t):
                 # avoids the "unflatten unevenly sharded" error when dp_shard*tp
                 # does not divide n_kv_heads (e.g. dp_shard=8, n_kv_heads=4); no
                 # gather, since fused is already replicated.
-                if not isinstance(t, DTensor) and (tp_size := spmd_mesh_size("tp")) > 1:
+                if not isinstance(t, DTensor):
                     # RL generator init_weights() only needs non-persistent
                     # buffers; weights come from trainer state dict. Until it has
                     # a DTensor static state dict path, copy this TP shard here.
                     # TODO: Remove once RL can init buffers without weight init.
-                    mesh = current_spmd_mesh()
-                    assert mesh is not None
-                    tp_rank = mesh.get_local_rank("tp")
-                    fused = fused.chunk(tp_size, dim=0)[tp_rank]
+                    annotation = get_local_type(t) if has_local_type(t) else None
+                    if annotation is not None:
+                        mesh = device_mesh_from_spmd_axes(
+                            annotation.keys(),
+                            storage_mesh=True,
+                        )
+                        assert mesh is not None, (
+                            "Fused QKV parameter annotation does not match a "
+                            f"registered storage mesh: {annotation}"
+                        )
+                        fused = spmd_distribute_tensor(
+                            fused,
+                            mesh,
+                            SpmdType(dict(get_local_type(t)), get_partition_spec(t)),
+                        )
                 t.copy_(fused)
 
         return _init
diff --git a/torchtitan/models/common/dist_gemm.py b/torchtitan/models/common/dist_gemm.py
index bb4824b3af..17717ef409 100644
--- a/torchtitan/models/common/dist_gemm.py
+++ b/torchtitan/models/common/dist_gemm.py
@@ -38,7 +38,7 @@
     LinearReduceScatter,
 )
 
-from torchtitan.distributed.spmd_types import current_spmd_mesh
+from torchtitan.distributed.spmd_types import current_spmd_mesh, get_mesh_pg
 from torchtitan.distributed.utils import get_spmd_backend
 
 from torchtitan.models.common.attention import FusedQKVLinear
@@ -81,7 +81,7 @@ def _tp_group_from_context() -> dist.ProcessGroup | None:
     mesh = current_spmd_mesh()
     if mesh is None or "tp" not in (mesh.mesh_dim_names or ()):
         return None
-    tp_group = mesh.get_group("tp")
+    tp_group = get_mesh_pg("tp")
     return tp_group if tp_group.size() > 1 else None
 
 
diff --git a/torchtitan/models/common/embedding.py b/torchtitan/models/common/embedding.py
index 39cea66711..d35ad95e50 100644
--- a/torchtitan/models/common/embedding.py
+++ b/torchtitan/models/common/embedding.py
@@ -15,12 +15,33 @@
 import torch.nn.functional as F
 from torch.distributed.tensor import DTensor
 
+from torchtitan.distributed.spmd_types import current_spmd_mesh
 from torchtitan.protocols.module import Module
 
 if TYPE_CHECKING:
     from torchtitan.distributed import ParallelDims
 
 
+def get_tp_rank(tp_group: dist.ProcessGroup) -> int | torch.SymInt:
+    """Return the TP rank, using a runtime symbol only during CooR tracing.
+
+    ``DeviceMesh._sym_get_coordinate`` returns the concrete coordinate in eager
+    execution and emits a runtime coordinate op only under compile-on-one-rank
+    fake tracing.
+    """
+    mesh = current_spmd_mesh()
+    if mesh is None:
+        return dist.get_rank(tp_group)
+
+    mesh_axis_names = mesh.mesh_dim_names
+    assert mesh_axis_names is not None, "DeviceMesh must have named axes"
+    if "tp" not in mesh_axis_names:
+        raise ValueError(
+            f"TP rank requires a 'tp' mesh axis, but got {mesh_axis_names}."
+        )
+    return mesh._sym_get_coordinate(mesh_axis_names.index("tp"))
+
+
 class Embedding(nn.Embedding, Module):
     """
     Configurable embedding with optional local vocab-parallel execution.
@@ -37,7 +58,7 @@ def __init__(self, config: Config):
         self.tp_group: dist.ProcessGroup | None = None
 
     def parallelize(self, parallel_dims: "ParallelDims") -> None:
-        # TODO(pianpwk): delete and rely on `current_spmd_mesh().get_group("tp")`
+        # TODO(pianpwk): delete and rely on `get_mesh_pg("tp")`
         # once the partial_dtensor backend is removed.
         tp_mesh = parallel_dims.get_optional_mesh("tp")
         if tp_mesh is not None:
@@ -49,7 +70,8 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         weight = (
             self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight
         )
-        if self.tp_group is None:
+        tp_group = self.tp_group
+        if tp_group is None:
             return F.embedding(
                 input,
                 weight,
@@ -60,11 +82,11 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
                 self.sparse,
             )
 
-        tp_pg = self.tp_group
-        tp_size = dist.get_world_size(tp_pg)
+        tp_rank = get_tp_rank(tp_group)
+        tp_size = dist.get_world_size(tp_group)
         weight = weight.to_local() if isinstance(weight, DTensor) else weight
         chunk_size = (self.num_embeddings + tp_size - 1) // tp_size
-        offset = dist.get_rank(tp_pg) * chunk_size
+        offset = tp_rank * chunk_size
         mask = (input >= offset) & (input < offset + weight.shape[0])
         local_input = (input - offset).clamp(0, weight.shape[0] - 1)
         out = F.embedding(
diff --git a/torchtitan/protocols/module.py b/torchtitan/protocols/module.py
index f63f4129cf..f097e8fac2 100644
--- a/torchtitan/protocols/module.py
+++ b/torchtitan/protocols/module.py
@@ -104,7 +104,7 @@ def init_states(
             for name, buf in self._buffers.items()
             if isinstance(buf, DTensor)
         }
-        with self._preserve_buffer_spmd_types():
+        with self._preserve_spmd_types():
             self._init_self_buffers(buffer_device=buffer_device)
         for name, (mesh, placements) in dtensor_meta.items():
             new_buf = self._buffers.get(name)
@@ -119,36 +119,38 @@ def init_states(
 
     def _apply(self, fn, recurse=True):
         """Override to preserve annotations across model.to_empty() in trainer.py"""
-        with self._preserve_buffer_spmd_types():
+        with self._preserve_spmd_types():
             return super()._apply(fn, recurse=recurse)
 
     @contextlib.contextmanager
-    def _preserve_buffer_spmd_types(self) -> Iterator[None]:
+    def _preserve_spmd_types(self) -> Iterator[None]:
         """
-        Preserve SPMD type annotations on buffers across reinitialization.
+        Preserve SPMD type annotations on parameters and buffers.
 
-        ``to_empty()`` and ``_init_self_buffers()`` re-materialize buffer data,
-        clobbering over SPMD annotations. Instead of attempting to typecheck over
-        this, we save-restore annotations on their respective mesh axes.
+        ``to_empty()`` re-materializes parameters and buffers, while
+        ``_init_self_buffers()`` can replace buffer data. Save and restore their
+        annotations across both operations.
         """
         if get_spmd_backend() != "spmd_types":
             yield
             return
 
+        states = (*self.named_parameters(), *self.named_buffers())
         saved = {
             fqn: SpmdLayout(
-                dict(spmd.get_local_type(buf)),
-                spmd.get_partition_spec(buf),
+                dict(spmd.get_local_type(state)),
+                spmd.get_partition_spec(state),
             )
-            for fqn, buf in self.named_buffers()
-            if spmd.has_local_type(buf)
+            for fqn, state in states
+            if spmd.has_local_type(state)
         }
         try:
             yield
         finally:
-            for fqn, buf in self.named_buffers():
-                if fqn in saved and not spmd.has_local_type(buf):
-                    spmd.assert_type(buf, saved[fqn])
+            states = (*self.named_parameters(), *self.named_buffers())
+            for fqn, state in states:
+                if fqn in saved and not spmd.has_local_type(state):
+                    spmd.assert_type(state, saved[fqn])
 
     def _init_self_parameters(self) -> None:
         """Initialize this module's own direct parameters.
@@ -304,7 +306,11 @@ def _spmd_distribute_state(
         assert mesh is not None
         assert mesh.mesh_dim_names is not None, "DeviceMesh must have named axes"
 
-        tensor = spmd_distribute_tensor(tensor, mesh, layout)
+        tensor = spmd_distribute_tensor(
+            tensor,
+            mesh,
+            layout,
+        )
         if is_param:
             self.register_parameter(name, nn.Parameter(tensor))
             registered = self._parameters[name]
diff --git a/torchtitan/protocols/state_dict_adapter.py b/torchtitan/protocols/state_dict_adapter.py
index 53d7ef0090..b48e6db92a 100644
--- a/torchtitan/protocols/state_dict_adapter.py
+++ b/torchtitan/protocols/state_dict_adapter.py
@@ -11,13 +11,14 @@
 from collections.abc import Mapping
 from typing import Any
 
+import spmd_types as spmd
+import torch
 from torch.distributed.checkpoint import HuggingFaceStorageReader
+from torch.distributed.tensor import DTensor
+from torch.utils._pytree import tree_map_only
 
-from torchtitan.distributed.parallel_dims import ParallelDims, SpmdLayout
-from torchtitan.distributed.spmd_types import (
-    dtensor_to_plain_tensor_state_dict,
-    plain_tensor_to_dtensor_state_dict,
-)
+from torchtitan.distributed.parallel_dims import ParallelDims
+from torchtitan.distributed.spmd_types import plain_tensor_to_dtensor_state_dict
 from torchtitan.tools.logging import logger
 from .model import BaseModel
 
@@ -161,18 +162,58 @@ def get_hf_storage_reader(
 class PlainToDTensorStateDictAdapter(BaseStateDictAdapter):
     def __init__(
         self,
-        state_dict_layouts: Mapping[str, SpmdLayout],
-        parallel_dims: ParallelDims,
+        state_dict_layouts: Mapping[str, spmd.SpmdType] | None = None,
+        parallel_dims: ParallelDims | None = None,
     ) -> None:
         self.state_dict_layouts = state_dict_layouts
         self.parallel_dims = parallel_dims
+        self.optimizers = None
+
+    @staticmethod
+    def _optimizer_layouts(optimizers, state_dict):
+        params = {
+            fqn: param
+            for optimizer in optimizers.optimizers
+            for group in optimizer.param_groups
+            for fqn, param in zip(group["param_names"], group["params"], strict=True)
+        }
+        return {
+            key: spmd.SpmdType(
+                dict(spmd.get_local_type(param)), spmd.get_partition_spec(param)
+            )
+            for key, value in state_dict.items()
+            for fqn, param in params.items()
+            if key.startswith(f"state.{fqn}.")
+            and isinstance(value, torch.Tensor)
+            and value.shape == param.shape
+            and spmd.has_local_type(param)
+        }
 
     def convert_save_state_dict(self, state_dict: dict[str, Any]) -> dict[str, Any]:
-        return plain_tensor_to_dtensor_state_dict(
+        converted = plain_tensor_to_dtensor_state_dict(
             state_dict,
             state_dict_layouts=self.state_dict_layouts,
             parallel_dims=self.parallel_dims,
         )
+        self.optimizers = state_dict.get("optimizer")
+        if self.optimizers is not None:
+            optimizer_state = self.optimizers.state_dict()
+            layouts = self._optimizer_layouts(self.optimizers, optimizer_state)
+            converted_optimizer = dict(optimizer_state)
+            converted_optimizer.update(
+                plain_tensor_to_dtensor_state_dict(
+                    {key: optimizer_state[key] for key in layouts},
+                    state_dict_layouts=layouts,
+                )
+            )
+            converted["optimizer"] = converted_optimizer
+        return converted
 
     def convert_load_state_dict(self, state_dict: dict[str, Any]) -> dict[str, Any]:
-        return dtensor_to_plain_tensor_state_dict(state_dict)
+        converted = tree_map_only(DTensor, lambda value: value.to_local(), state_dict)
+        if self.optimizers is not None and "optimizer" in converted:
+            # Optimizer state may not exist before loading, so preserve its
+            # normal load_state_dict path after DCP fills the target shards.
+            self.optimizers.load_state_dict(converted["optimizer"])
+            converted["optimizer"] = self.optimizers
+        return converted
diff --git a/torchtitan/trainer.py b/torchtitan/trainer.py
index d74ca2da03..31db779c75 100644
--- a/torchtitan/trainer.py
+++ b/torchtitan/trainer.py
@@ -460,7 +460,12 @@ def __init__(self, config: Config):
 
                 for m in self.model_parts:
                     m.to_empty(device=init_device)
-                    with torch.no_grad():
+                    with (
+                        torch.no_grad(),
+                        dist_utils.get_spmd_context(
+                            parallel_dims=parallel_dims,
+                        )(),
+                    ):
                         # TODO: Change this back to init_weights once
                         # autoparallel contains the wrap_init_states
                         cast(BaseModel, m).init_weights(buffer_device=buffer_device)
@@ -487,7 +492,12 @@ def __init__(self, config: Config):
                     )
 
                 model.to_empty(device=init_device)
-                with torch.no_grad():
+                with (
+                    torch.no_grad(),
+                    dist_utils.get_spmd_context(
+                        parallel_dims=parallel_dims,
+                    )(),
+                ):
                     # TODO: Change this back to init_weights once
                     # autoparallel contains the wrap_init_states
                     cast(BaseModel, model).init_weights(buffer_device=buffer_device)
@@ -839,6 +849,15 @@ def pp_forward_backward_step(
             return torch.sum(torch.stack(losses)).to(self.device)
         return torch.tensor([-1.0], device=self.device)
 
+    def _clip_grad_norm(self, parameters: list[torch.Tensor]) -> torch.Tensor:
+        return dist_utils.clip_grad_norm_(
+            parameters,
+            self.config.training.max_norm,
+            foreach=True,
+            pp_mesh=self.parallel_dims.get_optional_mesh("pp"),
+            ep_enabled=self.parallel_dims.ep_enabled,
+        )
+
     def train_step(
         self, data_iterator: Iterator[tuple[dict[str, torch.Tensor], torch.Tensor]]
     ):
@@ -916,12 +935,8 @@ def train_step(
                     accumulated_loss.add_(detached_loss)
 
         with sl.log_trace_span("optim"):
-            grad_norm = dist_utils.clip_grad_norm_(
-                [p for m in self.model_parts for p in m.parameters()],
-                self.config.training.max_norm,
-                foreach=True,
-                pp_mesh=parallel_dims.get_optional_mesh("pp"),
-                ep_enabled=parallel_dims.ep_enabled,
+            grad_norm = self._clip_grad_norm(
+                [p for m in self.model_parts for p in m.parameters()]
             )
             # Only the last PP stage owns the loss. First combine its DP/CP
             # replicas, then propagate the result across PP. TP replicas have