[dist.cp] Introduce LoadPlanner and SavePlanner extensibility API. (pytorch#83419)

The planners come with default implementations in default_planner.py.

The default planners expose their core functionality as separate functions
to make it easy for other checkpoint implementations to use this functionality.

Pull Request resolved: pytorch#83419
Approved by: https://github.com/wanchaol

Author: Rodrigo Kumpera

Diff for: test/distributed/_shard/checkpoint/test_planner.py

@@ -0,0 +1,268 @@
+# Owner(s): ["oncall: distributed"]
+
+import sys
+
+import torch
+from torch.distributed._shard.checkpoint.planner import LoadItemType, WriteItemType
+
+from torch.distributed._shard.sharded_tensor import (
+    Shard,
+    ShardMetadata,
+    ShardedTensor,
+    ShardedTensorMetadata,
+)
+from torch.distributed._shard.sharded_tensor.metadata import TensorProperties
+
+from torch.testing._internal.common_utils import (
+    TestCase,
+    TEST_WITH_DEV_DBG_ASAN,
+    run_tests,
+)
+from torch.distributed._shard.checkpoint.metadata import BytesStorageMetadata, MetadataIndex, TensorStorageMetadata
+from torch.testing._internal.distributed.distributed_utils import (
+    with_fake_comms,
+    with_dist
+)
+
+from torch.distributed._shard.checkpoint.default_planner import (
+    create_default_global_save_plan,
+    create_default_local_save_plan,
+    create_default_local_load_plan,
+    _create_default_local_metadata
+)
+
+if TEST_WITH_DEV_DBG_ASAN:
+    print(
+        "Skip dev-asan as torch + multiprocessing spawn have known issues",
+        file=sys.stderr,
+    )
+    sys.exit(0)
+
+def create_sharded_tensor(rank, world_size, shards_per_rank, shard_size=8):
+    shards_metadata = []
+    local_shards = []
+    for idx in range(0, world_size * shards_per_rank):
+        shard_rank = idx // shards_per_rank
+        shard_md = ShardMetadata(shard_offsets=[idx * shard_size], shard_sizes=[shard_size], placement=f"rank:{shard_rank}/cpu")
+        shards_metadata.append(shard_md)
+        if shard_rank == rank:
+            shard = Shard.from_tensor_and_offsets(
+                torch.rand(*shard_md.shard_sizes),
+                shard_offsets=shard_md.shard_offsets,
+                rank=rank
+            )
+            local_shards.append(shard)
+
+    sharded_tensor_md = ShardedTensorMetadata(
+        shards_metadata=shards_metadata,
+        size=torch.Size([shard_size * len(shards_metadata)]),
+        tensor_properties=TensorProperties.create_from_tensor(torch.zeros(1))
+    )
+
+    return ShardedTensor._init_from_local_shards_and_global_metadata(
+        local_shards=local_shards,
+        sharded_tensor_metadata=sharded_tensor_md
+    )
+
+
+class TestSavePlan(TestCase):
+    @with_fake_comms(rank=1, world_size=4)
+    def test_local_plan(self):
+        tensor = torch.rand(10)
+        val = [1, 2, 3]
+        st = create_sharded_tensor(rank=1, world_size=4, shards_per_rank=1)
+        state_dict = {
+            "tensor": tensor,
+            "value": val,
+            "st": st
+        }
+        plan = create_default_local_save_plan(state_dict, False)
+        self.assertEqual(1, 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.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]))
+
+        # Coordinator rank, should include replicated items as well
+        plan = create_default_local_save_plan(state_dict, True)
+        self.assertEqual(3, len(plan.items))
+
+        tensor_wi = next(wi for wi in plan.items if wi.type == WriteItemType.TENSOR)
+        self.assertEqual(tensor_wi.index, MetadataIndex("tensor", [0]))
+        self.assertEqual(tensor_wi.tensor_data.size, tensor.size())
+        self.assertEqual(tensor_wi.tensor_data.properties, TensorProperties.create_from_tensor(tensor))
+        self.assertEqual(tensor_wi.tensor_data.chunk.offsets, torch.Size([0]))
+        self.assertEqual(tensor_wi.tensor_data.chunk.sizes, torch.Size([10]))
+
+        bytes_wi = next(wi for wi in plan.items if wi.type == WriteItemType.BYTE_IO)
+        self.assertEqual(bytes_wi.index, MetadataIndex("value"))
+        self.assertIsNone(bytes_wi.tensor_data)
+
+    def test_global_plan(self):
+        def create_data(rank):
+            with with_dist(rank=rank, world_size=4):
+                tensor = torch.rand(10)
+                val = [1, 2, 3]
+                st = create_sharded_tensor(rank=rank, world_size=4, shards_per_rank=1)
+                state_dict = {
+                    "tensor": tensor,
+                    "value": val,
+                    "st": st
+                }
+                return create_default_local_save_plan(state_dict, rank == 0)
+
+        all_plans = [create_data(0), create_data(1), create_data(2), create_data(3)]
+        final_plans, metadata = create_default_global_save_plan(all_plans=all_plans)
+        # The default global plan updates all indexes to include hints
+        for new_plan, old_plan in zip(final_plans, all_plans):
+            for new_item, old_item in zip(new_plan.items, old_plan.items):
+                self.assertEqual(new_item.index, old_item.index)
+                self.assertEqual(new_item.type, old_item.type)
+                self.assertEqual(new_item.tensor_data, old_item.tensor_data)
+                self.assertIn(new_item.index.fqn, metadata.state_dict_metadata)
+
+                item_md = metadata.state_dict_metadata[new_item.index.fqn]
+                if new_item.type == WriteItemType.BYTE_IO:
+                    self.assertTrue(isinstance(item_md, BytesStorageMetadata))
+                else:
+                    self.assertTrue(isinstance(item_md, TensorStorageMetadata))
+                    self.assertEqual(item_md.size, old_item.tensor_data.size)
+                    self.assertEqual(item_md.properties, old_item.tensor_data.properties)
+
+                    self.assertIsNotNone(new_item.index.index)
+                    # Make sure the hint is correct
+                    self.assertEqual(item_md.chunks[new_item.index.index], old_item.tensor_data.chunk)
+
+    def test_local_load_plan(self):
+        def create_state_dict(rank):
+            with with_dist(rank=rank, world_size=4):
+                tensor = torch.rand(10)
+                val = [1, 2, 3]
+                st = create_sharded_tensor(rank=rank, world_size=4, shards_per_rank=1)
+                return {
+                    "tensor": tensor,
+                    "value": val,
+                    "st": st
+                }
+
+        state_dict = create_state_dict(1)
+        metadata = _create_default_local_metadata(state_dict)
+
+        load_plan = create_default_local_load_plan(state_dict, metadata)
+        # This will create 3 entries
+        self.assertEqual(3, len(load_plan.items))
+        st_item = next(ri for ri in load_plan.items if ri.dest_index.fqn == "st")
+        tensor_item = next(ri for ri in load_plan.items if ri.dest_index.fqn == "tensor")
+        bytes_item = next(ri for ri in load_plan.items if ri.dest_index.fqn == "value")
+
+        self.assertEqual(st_item.type, LoadItemType.TENSOR)
+        # This is an exact copy
+        self.assertEqual(st_item.dest_index, MetadataIndex("st", [8]))
+        self.assertEqual(st_item.dest_offsets, torch.Size([0]))
+        self.assertEqual(st_item.storage_index, MetadataIndex("st", [8]))
+        self.assertEqual(st_item.storage_offsets, torch.Size([0]))
+        self.assertEqual(st_item.lengths, torch.Size([8]))
+
+        self.assertEqual(tensor_item.type, LoadItemType.TENSOR)
+        self.assertEqual(tensor_item.dest_index, MetadataIndex("tensor", [0]))
+        self.assertEqual(tensor_item.dest_offsets, torch.Size([0]))
+        self.assertEqual(tensor_item.storage_index, MetadataIndex("tensor", [0]))
+        self.assertEqual(tensor_item.storage_offsets, torch.Size([0]))
+        self.assertEqual(tensor_item.lengths, torch.Size([10]))
+
+        self.assertEqual(bytes_item.type, LoadItemType.BYTE_IO)
+        self.assertEqual(bytes_item.dest_index, MetadataIndex("value"))
+
+    def test_load_with_resharding(self):
+        def create_state_dict(rank, world_size):
+            with with_dist(rank=rank, world_size=world_size):
+                return {
+                    "st": create_sharded_tensor(
+                        rank=rank,
+                        world_size=world_size,
+                        shards_per_rank=1,
+                        shard_size=128 // world_size,
+                    )
+                }
+
+
+        # Rank 1 has a 16 bytes shard from [16, 32[
+        world8_state_dict = create_state_dict(rank=1, world_size=8)
+        world8_metadata = _create_default_local_metadata(world8_state_dict)
+
+        # Rank 1 has a 32 bytes shard from [32, 64[
+        world4_state_dict = create_state_dict(rank=1, world_size=4)
+        world4_metadata = _create_default_local_metadata(world4_state_dict)
+
+        # First scenario, going from world=8 to world=4, need to load 2 shards
+        # Each 4-world shard has 32 elements, so it needs to load 2 shards
+        load_plan = create_default_local_load_plan(world4_state_dict, world8_metadata)
+        self.assertEqual(2, len(load_plan.items))
+        low_ri = next(ri for ri in load_plan.items if ri.dest_offsets == torch.Size([0]))
+        high_ri = next(ri for ri in load_plan.items if ri.dest_offsets == torch.Size([16]))
+
+        self.assertEqual(low_ri.storage_index, MetadataIndex("st", [32]))
+        self.assertEqual(low_ri.storage_offsets, torch.Size([0]))
+        self.assertEqual(low_ri.dest_index, MetadataIndex("st", [32]))
+        self.assertEqual(low_ri.dest_offsets, torch.Size([0]))
+        self.assertEqual(low_ri.lengths, torch.Size([16]))
+
+        self.assertEqual(high_ri.storage_index, MetadataIndex("st", [48]))
+        self.assertEqual(high_ri.storage_offsets, torch.Size([0]))
+        self.assertEqual(high_ri.dest_index, MetadataIndex("st", [32]))
+        self.assertEqual(high_ri.dest_offsets, torch.Size([16]))
+        self.assertEqual(high_ri.lengths, torch.Size([16]))
+
+        # Second scenario, going from world=4 to world=8, need to load half of 1 shard
+        # rank1 on 8-world needs to load the upper half of the rank0 4-world shard
+        load_plan = create_default_local_load_plan(world8_state_dict, world4_metadata)
+        self.assertEqual(1, len(load_plan.items))
+        ri = load_plan.items[0]
+        self.assertEqual(ri.storage_index, MetadataIndex("st", [0]))
+        self.assertEqual(ri.storage_offsets, torch.Size([16]))
+        self.assertEqual(ri.dest_index, MetadataIndex("st", [16]))
+        self.assertEqual(ri.dest_offsets, torch.Size([0]))
+        self.assertEqual(ri.lengths, torch.Size([16]))
+
+    def test_load_with_world_size_diff_by_one(self):
+        def create_state_dict(rank, world_size):
+            with with_dist(rank=rank, world_size=world_size):
+                return {
+                    "st": create_sharded_tensor(
+                        rank=rank,
+                        world_size=world_size,
+                        shards_per_rank=1,
+                        shard_size=120 // world_size,
+                    )
+                }
+        # rank 1 has a 30 bytes shard from [30, 60[
+        world4_state_dict = create_state_dict(rank=1, world_size=4)
+        world4_metadata = _create_default_local_metadata(world4_state_dict)
+
+        # rank 1 has a 40 bytes shard from [40, 80[
+        world3_state_dict = create_state_dict(rank=1, world_size=3)
+
+        load_plan = create_default_local_load_plan(world3_state_dict, world4_metadata)
+        self.assertEqual(2, len(load_plan.items))
+        # this is [30, 60] to load [40, 60]
+        low_ri = next(ri for ri in load_plan.items if ri.dest_offsets == torch.Size([0]))
+        # this is [60, 90] to load [60, 80]
+        high_ri = next(ri for ri in load_plan.items if ri.dest_offsets == torch.Size([20]))
+
+        self.assertEqual(low_ri.storage_index, MetadataIndex("st", [30]))
+        self.assertEqual(low_ri.storage_offsets, torch.Size([10]))
+        self.assertEqual(low_ri.dest_index, MetadataIndex("st", [40]))
+        self.assertEqual(low_ri.dest_offsets, torch.Size([0]))
+        self.assertEqual(low_ri.lengths, torch.Size([20]))
+
+        self.assertEqual(high_ri.storage_index, MetadataIndex("st", [60]))
+        self.assertEqual(high_ri.storage_offsets, torch.Size([0]))
+        self.assertEqual(high_ri.dest_index, MetadataIndex("st", [40]))
+        self.assertEqual(high_ri.dest_offsets, torch.Size([20]))
+        self.assertEqual(high_ri.lengths, torch.Size([20]))
+
+if __name__ == "__main__":
+    run_tests()

Diff for: test/distributed/_shard/checkpoint/test_utils.py

@@ -81,6 +81,8 @@ def test_flat_data(self):
 
         a = find_state_dict_object(state_dict, MetadataIndex("a"))
         self.assertEqual(a, state_dict["a"])
+        a = find_state_dict_object(state_dict, MetadataIndex("a", [0]))
+        self.assertEqual(a, state_dict["a"])
         a = find_state_dict_object(state_dict, MetadataIndex("a", index=99))
         self.assertEqual(a, state_dict["a"])
 
@@ -91,8 +93,6 @@ def test_flat_data(self):
 
         with self.assertRaisesRegex(ValueError, "FQN"):
             find_state_dict_object(state_dict, MetadataIndex("c"))
-        with self.assertRaisesRegex(ValueError, "ShardedTensor"):
-            find_state_dict_object(state_dict, MetadataIndex("a", [0]))
         with self.assertRaisesRegex(ValueError, "ShardedTensor"):
             find_state_dict_object(state_dict, MetadataIndex("b", [1]))
 

Diff for: torch/distributed/_shard/checkpoint/__init__.py

@@ -13,3 +13,13 @@
 from .storage import StorageReader, StorageWriter
 from .filesystem import FileSystemReader, FileSystemWriter
 from .api import CheckpointException
+
+
+from .planner import (
+    SavePlanner,
+    LoadPlanner,
+    SavePlan,
+    LoadPlan,
+    ReadItem,
+    WriteItem,
+)