[DTensor] Raise an RuntimeError when checkpointing APIs are used with Partial placement (#163941)

fegin · pytorchmergebot · commit dfda239ccee4 · 2025-09-27T19:50:16.000Z
A DTensor that contains partial placement shouldn't be checkpointed (DCP.save) -- the result is not correct and DCP doesn't know how to handle it. There are several APIs that are only used by checkpointing, e.g.,`__create_write_items__`. These APIs should raise an exception if the DTensor, `self`, has Partial placement. Ideally, we want to add the following test: ``` with self.assertRaisesRegex( RuntimeError, "Any checkpointing related operations are not supported for" ): dcp.save({"dtensor": dtensor}, checkpoint_id=tempfile.gettempdir()) ``` While we do see the RuntimeError is raised, the error was raised in another thread due to DTensor checkpoint APIs are called by DCP in a separate thread, which assertRaisesRegex cannot capture. Pull Request resolved: #163941 Approved by: https://github.com/tianyu-l
diff --git a/test/distributed/tensor/test_api.py b/test/distributed/tensor/test_api.py
@@ -1,13 +1,17 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # Owner(s): ["oncall: distributed"]
 
+import tempfile
+
 import torch
+import torch.distributed.checkpoint as dcp
 import torch.nn as nn
 from torch.distributed.tensor import (
     DeviceMesh,
     distribute_module,
     distribute_tensor,
     DTensor,
+    Partial,
     Replicate,
     Shard,
 )
@@ -356,6 +360,33 @@ def shard_fn(name, module, device_mesh):
             self.assertFalse(param.is_meta)
             self.assertTrue(param.device.type == device_mesh.device_type)
 
+    @with_comms
+    def test_checkpoint_apis_check_partial_placement(self):
+        device_mesh = self.build_device_mesh()
+        tensor = torch.randn(5, 5, device=self.device_type)
+        dtensor = DTensor.from_local(tensor, device_mesh, [Partial()])
+        with self.assertRaisesRegex(
+            ValueError, "Any checkpointing related operations are not supported for"
+        ):
+            dtensor.__create_write_items__("fqn", None)
+
+        with self.assertRaisesRegex(
+            ValueError, "Any checkpointing related operations are not supported for"
+        ):
+            dtensor.__create_chunk_list__()
+
+        with self.assertRaisesRegex(
+            ValueError, "Any checkpointing related operations are not supported for"
+        ):
+            dtensor.__get_tensor_shard__(0)
+
+        # Ideally we should not allow checkpointing related operations for DTensor
+        with self.assertRaisesRegex(
+            dcp.api.CheckpointException,
+            "Any checkpointing related operations are not supported for",
+        ):
+            dcp.save({"fqn": dtensor}, checkpoint_id=tempfile.mkdtemp())
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/torch/distributed/tensor/_api.py b/torch/distributed/tensor/_api.py
@@ -592,7 +592,21 @@ def placements(self) -> tuple[Placement, ...]:
         """
         return self._spec.placements
 
+    def _raise_if_contains_partial_placements(self) -> None:
+        """
+        Raise an error if the DTensor contains partial placements.
+        """
+        for placement in self._spec.placements:
+            if not isinstance(placement, Partial):
+                continue
+
+            raise ValueError(
+                "Any checkpointing related operations are not supported for "
+                "DTensor with partial placements!"
+            )
+
     def __create_write_items__(self, fqn: str, object: Any):
+        self._raise_if_contains_partial_placements()
         from torch.distributed.checkpoint.planner_helpers import (
             _create_write_items_for_dtensor,
         )
@@ -615,6 +629,7 @@ def __create_chunk_list__(self):
         Returns:
             A List[:class:`ChunkStorageMetadata`] object that represents the shard size/offset on the current rank.
         """
+        self._raise_if_contains_partial_placements()
         from torch.distributed.checkpoint.planner_helpers import (
             _create_chunk_from_dtensor,
         )
@@ -627,6 +642,7 @@ def __create_chunk_list__(self):
             raise RuntimeError("Unsupported tensor type!")
 
     def __get_tensor_shard__(self, index):
+        self._raise_if_contains_partial_placements()
         if hasattr(self._local_tensor, "__get_tensor_shard__"):
             return self._local_tensor.__get_tensor_shard__(index)  # type: ignore[attr-defined]
         elif isinstance(self._local_tensor, torch.Tensor):
