已合并
test(distributed): add test for WriteItem.tensor_storage_size #35400
Flipped创建于 5月12日
test(distributed): add test for WriteItem.tensor_storage_size #35400
已合并
共 1 个文件变更+73-0
| @@ -0,0 +1,73 @@ | |||
| 1 | +""" | ||
| 2 | +Add validation cases for torch.distributed.checkpoint.planner APIs on NPU: | ||
| 3 | +1. PyTorch community tests lack sufficient validation for WriteItem.tensor_storage_size, so this file is added. | ||
| 4 | +2. This file validates torch.distributed.checkpoint.planner.WriteItem.tensor_storage_size. | ||
| 5 | +""" | ||
| 6 | + | ||
| 7 | +import torch | ||
| 8 | +from torch.distributed.checkpoint.metadata import ( | ||
| 9 | + ChunkStorageMetadata, | ||
| 10 | + MetadataIndex, | ||
| 11 | + TensorProperties, | ||
| 12 | +) | ||
| 13 | +from torch.distributed.checkpoint.planner import ( | ||
| 14 | + TensorWriteData, | ||
| 15 | + WriteItem, | ||
| 16 | + WriteItemType, | ||
| 17 | +) | ||
| 18 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 22 | + | ||
| 23 | + | ||
| 24 | +class TestPlannerAPI(TestCase): | ||
| 25 | + | ||
| 26 | + def _make_tensor_write_item(self, tensor, write_item_type): | ||
| 27 | + tensor_data = TensorWriteData( | ||
| 28 | + chunk=ChunkStorageMetadata( | ||
| 29 | + offsets=torch.Size([0] * tensor.dim()), | ||
| 30 | + sizes=tensor.size(), | ||
| 31 | + ), | ||
| 32 | + properties=TensorProperties.create_from_tensor(tensor), | ||
| 33 | + size=tensor.size(), | ||
| 34 | + ) | ||
| 35 | + | ||
| 36 | + return WriteItem( | ||
| 37 | + index=MetadataIndex("tensor"), | ||
| 38 | + type=write_item_type, | ||
| 39 | + tensor_data=tensor_data, | ||
| 40 | + ) | ||
| 41 | + | ||
| 42 | + def test_write_item_tensor_storage_size_for_tensor(self): | ||
| 43 | + for dtype in (torch.float32, torch.float16, torch.int8): | ||
| 44 | + tensor = torch.empty((2, 3), dtype=dtype).to(device_type) | ||
| 45 | + write_item = self._make_tensor_write_item( | ||
| 46 | + tensor, | ||
| 47 | + WriteItemType.TENSOR, | ||
| 48 | + ) | ||
| 49 | + | ||
| 50 | + expected_size = tensor.numel() * tensor.element_size() | ||
| 51 | + self.assertEqual(write_item.tensor_storage_size(), expected_size) | ||
| 52 | + | ||
| 53 | + def test_write_item_tensor_storage_size_for_shard(self): | ||
| 54 | + tensor = torch.empty((2, 3), dtype=torch.float32).to(device_type) | ||
| 55 | + write_item = self._make_tensor_write_item( | ||
| 56 | + tensor, | ||
| 57 | + WriteItemType.SHARD, | ||
| 58 | + ) | ||
| 59 | + | ||
| 60 | + expected_size = tensor.numel() * tensor.element_size() | ||
| 61 | + self.assertEqual(write_item.tensor_storage_size(), expected_size) | ||
| 62 | + | ||
| 63 | + def test_write_item_tensor_storage_size_for_non_tensor(self): | ||
| 64 | + write_item = WriteItem( | ||
| 65 | + index=MetadataIndex("bytes"), | ||
| 66 | + type=WriteItemType.BYTE_IO, | ||
| 67 | + ) | ||
| 68 | + | ||
| 69 | + self.assertIsNone(write_item.tensor_storage_size()) | ||
| 70 | + | ||
| 71 | + | ||
| 72 | +if __name__ == "__main__": | ||
| 73 | + run_tests() | ||