[DCP] Modify tensor saving logic in DCP (#106415)

Currently, DCP treats tensors as duplicates and only saves them on rank0. This won't work for PiPPy as PiPPy does have unique tensors across different ranks. With the current setup, we would only be saving the tensors on rank0 (coordinator rank).

In this PR, we are changing to letting each rank create its own WriteItem for tensors. For the ones that does replicate across different ranks, we are handling it thru dedup_tensors(), which will dedup the replicate WriteItem so we only do the actual writing once.
Pull Request resolved: #106415
Approved by: https://github.com/wz337

Author: eddogola

test/distributed/checkpoint/test_planner.py

@@ -38,6 +38,7 @@
 )
 
 from torch.distributed.checkpoint.planner_helpers import create_read_items_for_chunk_list
+from torch.distributed.checkpoint._dedup_tensors import dedup_tensors
 
 
 if TEST_WITH_DEV_DBG_ASAN:
@@ -86,14 +87,22 @@ def test_local_plan(self):
             "st": st
         }
         plan = create_default_local_save_plan(state_dict, False)
-        self.assertEqual(1, len(plan.items))
+        self.assertEqual(2, len(plan.items))
         wi = plan.items[0]
-        self.assertEqual(wi.index, MetadataIndex("st", [8]))
-        self.assertEqual(wi.type, WriteItemType.SHARD)
-        self.assertEqual(wi.tensor_data.size, st.size())
+        self.assertEqual(wi.index, MetadataIndex("tensor", [0]))
+        self.assertEqual(wi.type, WriteItemType.TENSOR)
+        self.assertEqual(wi.tensor_data.size, tensor.size())
         self.assertEqual(wi.tensor_data.properties, TensorProperties.create_from_tensor(torch.zeros(1)))
-        self.assertEqual(wi.tensor_data.chunk.offsets, torch.Size([8]))
-        self.assertEqual(wi.tensor_data.chunk.sizes, torch.Size([8]))
+        self.assertEqual(wi.tensor_data.chunk.offsets, torch.Size([0]))
+        self.assertEqual(wi.tensor_data.chunk.sizes, torch.Size([10]))
+
+        st_wi = plan.items[1]
+        self.assertEqual(st_wi.index, MetadataIndex("st", [8]))
+        self.assertEqual(st_wi.type, WriteItemType.SHARD)
+        self.assertEqual(st_wi.tensor_data.size, st.size())
+        self.assertEqual(st_wi.tensor_data.properties, TensorProperties.create_from_tensor(torch.zeros(1)))
+        self.assertEqual(st_wi.tensor_data.chunk.offsets, torch.Size([8]))
+        self.assertEqual(st_wi.tensor_data.chunk.sizes, torch.Size([8]))
 
         # Coordinator rank, should include replicated items as well
         plan = create_default_local_save_plan(state_dict, True)
@@ -124,6 +133,7 @@ def create_data(rank):
                 return create_default_local_save_plan(state_dict, rank == 0)
 
         all_plans = [create_data(0), create_data(1), create_data(2), create_data(3)]
+        all_plans = dedup_tensors(all_plans)
         final_plans, metadata = create_default_global_save_plan(all_plans=all_plans)
 
         # The default global plan updates all indexes to include hints

torch/distributed/checkpoint/default_planner.py

@@ -11,7 +11,6 @@
 import torch
 
 from torch.distributed._shard._utils import narrow_tensor_by_index
-from torch.distributed._shard.sharded_tensor import ShardedTensor
 from torch.distributed._tensor import DTensor
 
 
@@ -294,7 +293,7 @@ def create_default_local_save_plan(
         if isinstance(obj, DTensor):
             if obj.device_mesh.get_coordinate() is not None:
                 requests += _create_write_items(fqn, obj)
-        elif isinstance(obj, (ShardedTensor)) or is_coordinator:
+        elif isinstance(obj, (torch.Tensor)) or is_coordinator:
             requests += _create_write_items(fqn, obj)
 
     return SavePlan(requests)