已合并
test: add staging api npu tests #35469
zjucn创建于 5月13日
test: add staging api npu tests #35469
已合并
共 1 个文件变更+206-0
| @@ -0,0 +1,206 @@ | |||
| 1 | +""" | ||
| 2 | +1. PyTorch community lacks direct validation cases for some | ||
| 3 | + torch.distributed.checkpoint.staging APIs, so this file is added. | ||
| 4 | + | ||
| 5 | +2. This file validates the following APIs: | ||
| 6 | + torch.distributed.checkpoint.staging.AsyncStager | ||
| 7 | + torch.distributed.checkpoint.staging.AsyncStager.should_synchronize_after_execute | ||
| 8 | + torch.distributed.checkpoint.staging.AsyncStager.stage | ||
| 9 | + torch.distributed.checkpoint.staging.AsyncStager.synchronize_staging | ||
| 10 | + torch.distributed.checkpoint.staging.BlockingAsyncStager | ||
| 11 | + torch.distributed.checkpoint.staging.BlockingAsyncStager.stage | ||
| 12 | + torch.distributed.checkpoint.staging.BlockingAsyncStager.synchronize_staging | ||
| 13 | + (extendable) | ||
| 14 | +""" | ||
| 15 | + | ||
| 16 | +import tempfile | ||
| 17 | + | ||
| 18 | +import torch | ||
| 19 | +import torch.distributed.checkpoint as dcp | ||
| 20 | +from torch.distributed.checkpoint import FileSystemWriter | ||
| 21 | +from torch.distributed.checkpoint.staging import AsyncStager, BlockingAsyncStager | ||
| 22 | +from torch.distributed.checkpoint.state_dict_saver import AsyncCheckpointerType | ||
| 23 | + | ||
| 24 | +from torch_npu.testing.testcase import run_tests, TestCase | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 28 | + | ||
| 29 | + | ||
| 30 | +class TestStatefulObj: | ||
| 31 | + def __init__(self, tensor): | ||
| 32 | + self.tensor = tensor | ||
| 33 | + | ||
| 34 | + def state_dict(self): | ||
| 35 | + return {"tensor": self.tensor} | ||
| 36 | + | ||
| 37 | + def load_state_dict(self, state_dict): | ||
| 38 | + self.tensor.copy_(state_dict["tensor"]) | ||
| 39 | + | ||
| 40 | + | ||
| 41 | +class SyncTrackingFileSystemWriter(FileSystemWriter): | ||
| 42 | + _synchronize_after_execute = True | ||
| 43 | + | ||
| 44 | + def __init__(self, *args, **kwargs): | ||
| 45 | + super().__init__(*args, **kwargs) | ||
| 46 | + self.synchronize_staging_called = False | ||
| 47 | + | ||
| 48 | + def synchronize_staging(self): | ||
| 49 | + self.synchronize_staging_called = True | ||
| 50 | + | ||
| 51 | + | ||
| 52 | +class TestCheckpointStagingApi(TestCase): | ||
| 53 | + def test_async_stager_protocol_defaults(self): | ||
| 54 | + class UnimplementedStager: | ||
| 55 | + _synchronize_after_execute = True | ||
| 56 | + | ||
| 57 | + | ||
| 58 | + def should_synchronize_after_execute(self): | ||
| 59 | + return AsyncStager.should_synchronize_after_execute.fget(self) | ||
| 60 | + | ||
| 61 | + def stage(self, state_dict): | ||
| 62 | + return AsyncStager.stage(self, state_dict) | ||
| 63 | + | ||
| 64 | + def synchronize_staging(self): | ||
| 65 | + return AsyncStager.synchronize_staging(self) | ||
| 66 | + | ||
| 67 | + def close(self): | ||
| 68 | + return None | ||
| 69 | + | ||
| 70 | + stager = UnimplementedStager() | ||
| 71 | + | ||
| 72 | + self.assertIsInstance(stager, AsyncStager) | ||
| 73 | + self.assertTrue(stager.should_synchronize_after_execute) | ||
| 74 | + with self.assertRaisesRegex(NotImplementedError, "must implement stage method"): | ||
| 75 | + stager.stage({"tensor": torch.ones(2).to(device_type)}) | ||
| 76 | + self.assertIsNone(stager.synchronize_staging()) | ||
| 77 | + | ||
| 78 | + def test_blocking_async_stager_stage_npu_tensor_to_cpu_snapshot(self): | ||
| 79 | + state_dict = { | ||
| 80 | + "weight": torch.arange(12, dtype=torch.float32) | ||
| 81 | + .reshape(3, 4) | ||
| 82 | + .to(device_type), | ||
| 83 | + "nested": { | ||
| 84 | + "bias": torch.ones(4, dtype=torch.float32).to(device_type), | ||
| 85 | + }, | ||
| 86 | + } | ||
| 87 | + original_weight = state_dict["weight"].cpu().clone() | ||
| 88 | + | ||
| 89 | + stager = BlockingAsyncStager(cache_staged_state_dict=False) | ||
| 90 | + staged_state_dict = stager.stage(state_dict) | ||
| 91 | + | ||
| 92 | + self.assertEqual("cpu", staged_state_dict["weight"].device.type) | ||
| 93 | + self.assertEqual("cpu", staged_state_dict["nested"]["bias"].device.type) | ||
| 94 | + self.assertEqual(original_weight, staged_state_dict["weight"]) | ||
| 95 | + | ||
| 96 | + state_dict["weight"].add_(100) | ||
| 97 | + self.assertEqual(original_weight, staged_state_dict["weight"]) | ||
| 98 | + self.assertNotEqual(state_dict["weight"].cpu(), staged_state_dict["weight"]) | ||
| 99 | + | ||
| 100 | + def test_blocking_async_stager_cached_reuses_cpu_buffer(self): | ||
| 101 | + stager = BlockingAsyncStager(cache_staged_state_dict=True) | ||
| 102 | + | ||
| 103 | + first_state_dict = { | ||
| 104 | + "tensor": torch.ones(8, dtype=torch.float32).to(device_type) | ||
| 105 | + } | ||
| 106 | + first_staged = stager.stage(first_state_dict) | ||
| 107 | + first_data_ptr = first_staged["tensor"].data_ptr() | ||
| 108 | + | ||
| 109 | + second_state_dict = { | ||
| 110 | + "tensor": torch.full((8,), 3.0, dtype=torch.float32).to(device_type), | ||
| 111 | + } | ||
| 112 | + second_staged = stager.stage(second_state_dict) | ||
| 113 | + second_data_ptr = second_staged["tensor"].data_ptr() | ||
| 114 | + | ||
| 115 | + self.assertEqual(first_data_ptr, second_data_ptr) | ||
| 116 | + self.assertEqual(second_state_dict["tensor"].cpu(), second_staged["tensor"]) | ||
| 117 | + | ||
| 118 | + def test_blocking_async_stager_sync_noop_and_property(self): | ||
| 119 | + stager = BlockingAsyncStager() | ||
| 120 | + | ||
| 121 | + self.assertIsInstance(stager, AsyncStager) | ||
| 122 | + self.assertFalse(stager.should_synchronize_after_execute) | ||
| 123 | + self.assertIsNone(stager.synchronize_staging()) | ||
| 124 | + | ||
| 125 | + def test_filesystem_writer_stage_sets_copy_ahead_zero(self): | ||
| 126 | + with tempfile.TemporaryDirectory() as checkpoint_dir: | ||
| 127 | + writer = FileSystemWriter( | ||
| 128 | + checkpoint_dir, | ||
| 129 | + per_thread_copy_ahead=1024, | ||
| 130 | + ) | ||
| 131 | + staged_state_dict = writer.stage( | ||
| 132 | + {"tensor": torch.ones(4, dtype=torch.float32).to(device_type)} | ||
| 133 | + ) | ||
| 134 | + | ||
| 135 | + self.assertEqual(0, writer.per_thread_copy_ahead) | ||
| 136 | + self.assertEqual("cpu", staged_state_dict["tensor"].device.type) | ||
| 137 | + self.assertEqual( | ||
| 138 | + torch.ones(4, dtype=torch.float32), | ||
| 139 | + staged_state_dict["tensor"], | ||
| 140 | + ) | ||
| 141 | + | ||
| 142 | + def test_async_save_npu_state_dict_with_filesystem_writer(self): | ||
| 143 | + for cache_staged_state_dict in (False, True): | ||
| 144 | + with self.subTest(cache_staged_state_dict=cache_staged_state_dict): | ||
| 145 | + with tempfile.TemporaryDirectory() as checkpoint_dir: | ||
| 146 | + source_state_dict = { | ||
| 147 | + "tensor": torch.arange(8, dtype=torch.float32).to(device_type), | ||
| 148 | + } | ||
| 149 | + writer = FileSystemWriter( | ||
| 150 | + checkpoint_dir, | ||
| 151 | + cache_staged_state_dict=cache_staged_state_dict, | ||
| 152 | + ) | ||
| 153 | + | ||
| 154 | + future = dcp.async_save( | ||
| 155 | + source_state_dict, | ||
| 156 | + storage_writer=writer, | ||
| 157 | + async_checkpointer_type=AsyncCheckpointerType.THREAD, | ||
| 158 | + ) | ||
| 159 | + future.result() | ||
| 160 | + | ||
| 161 | + loaded_state_dict = { | ||
| 162 | + "tensor": torch.zeros(8, dtype=torch.float32).to(device_type), | ||
| 163 | + } | ||
| 164 | + dcp.load(loaded_state_dict, checkpoint_id=checkpoint_dir) | ||
| 165 | + | ||
| 166 | + self.assertEqual( | ||
| 167 | + source_state_dict["tensor"].cpu(), | ||
| 168 | + loaded_state_dict["tensor"].cpu(), | ||
| 169 | + ) | ||
| 170 | + | ||
| 171 | + def test_async_save_converts_stateful_object(self): | ||
| 172 | + with tempfile.TemporaryDirectory() as checkpoint_dir: | ||
| 173 | + source_obj = TestStatefulObj( | ||
| 174 | + torch.arange(4, dtype=torch.float32).to(device_type) | ||
| 175 | + ) | ||
| 176 | + | ||
| 177 | + future = dcp.async_save( | ||
| 178 | + {"obj": source_obj}, | ||
| 179 | + storage_writer=FileSystemWriter(checkpoint_dir), | ||
| 180 | + async_checkpointer_type=AsyncCheckpointerType.THREAD, | ||
| 181 | + ) | ||
| 182 | + future.result() | ||
| 183 | + | ||
| 184 | + loaded_obj = TestStatefulObj( | ||
| 185 | + torch.zeros(4, dtype=torch.float32).to(device_type) | ||
| 186 | + ) | ||
| 187 | + dcp.load({"obj": loaded_obj}, checkpoint_id=checkpoint_dir) | ||
| 188 | + | ||
| 189 | + self.assertEqual(source_obj.tensor.cpu(), loaded_obj.tensor.cpu()) | ||
| 190 | + | ||
| 191 | + def test_async_save_calls_synchronize_when_stager_requests_it(self): | ||
| 192 | + with tempfile.TemporaryDirectory() as checkpoint_dir: | ||
| 193 | + writer = SyncTrackingFileSystemWriter(checkpoint_dir) | ||
| 194 | + | ||
| 195 | + future = dcp.async_save( | ||
| 196 | + {"tensor": torch.ones(4, dtype=torch.float32).to(device_type)}, | ||
| 197 | + storage_writer=writer, | ||
| 198 | + async_checkpointer_type=AsyncCheckpointerType.THREAD, | ||
| 199 | + ) | ||
| 200 | + | ||
| 201 | + self.assertTrue(writer.synchronize_staging_called) | ||
| 202 | + future.result() | ||
| 203 | + | ||
| 204 | + | ||
| 205 | +if __name__ == "__main__": | ||
| 206 | + run_tests() | ||