已合并
fix d2h pinned memeory bug #31778
zzhongmin创建于 3月13日
fix d2h pinned memeory bug #31778
已合并
共 2 个文件变更+51-0
| @@ -2,6 +2,7 @@ import io | |||
| 2 | import os | 2 | import os |
| 3 | import tempfile | 3 | import tempfile |
| 4 | import argparse | 4 | import argparse |
| 5 | +from unittest.mock import patch | ||
| 5 | 6 | ||
| 6 | import torch | 7 | import torch |
| 7 | import torch.nn as nn | 8 | import torch.nn as nn |
| @@ -254,6 +255,34 @@ class TestSerialization(TestCase): | |||
| 254 | self.assertRtolEqual(before_save['fc1.bias'].cpu(), after_load['fc1.bias'].cpu()) | 255 | self.assertRtolEqual(before_save['fc1.bias'].cpu(), after_load['fc1.bias'].cpu()) |
| 255 | self.assertRtolEqual(before_save['fc2.bias'].cpu(), after_load['fc2.bias'].cpu()) | 256 | self.assertRtolEqual(before_save['fc2.bias'].cpu(), after_load['fc2.bias'].cpu()) |
| 256 | 257 | ||
| 258 | + def test_use_pinned_memory_for_d2h(self): | ||
| 259 | + if not torch.accelerator.is_available(): | ||
| 260 | + self.skipTest("accelerator is not available") | ||
| 261 | + if torch.accelerator.current_accelerator().type != "npu": | ||
| 262 | + self.skipTest("current accelerator is not npu") | ||
| 263 | + | ||
| 264 | + from torch.utils.serialization import config as serialization_config | ||
| 265 | + | ||
| 266 | + def patched_write_record(self, filename, data, nbytes): | ||
| 267 | + if isinstance(data, (torch.TypedStorage, torch.UntypedStorage)): | ||
| 268 | + if not data.is_pinned(device="npu"): | ||
| 269 | + raise RuntimeError("Expected storage to be in pinned memory") | ||
| 270 | + | ||
| 271 | + sd = torch.nn.Linear(3, 5, device="npu").state_dict() | ||
| 272 | + | ||
| 273 | + with patch('torch._C.PyTorchFileWriter.write_record', patched_write_record): | ||
| 274 | + with tempfile.NamedTemporaryFile() as f: | ||
| 275 | + with self.assertRaisesRegex(RuntimeError, "Expected storage to be in pinned memory"): | ||
| 276 | + torch.save(sd, f) | ||
| 277 | + | ||
| 278 | + with tempfile.NamedTemporaryFile() as f: | ||
| 279 | + pinned_before = serialization_config.save.use_pinned_memory_for_d2h | ||
| 280 | + try: | ||
| 281 | + serialization_config.save.use_pinned_memory_for_d2h = True | ||
| 282 | + torch.save(sd, f) | ||
| 283 | + finally: | ||
| 284 | + serialization_config.save.use_pinned_memory_for_d2h = pinned_before | ||
| 285 | + | ||
| 257 | def test_save_different_dtype_unallocated(self): | 286 | def test_save_different_dtype_unallocated(self): |
| 258 | 287 | ||
| 259 | def save_load_check(): | 288 | def save_load_check(): |
| @@ -416,6 +416,28 @@ def _npu_save( | |||
| 416 | storage = new_storage | 416 | storage = new_storage |
| 417 | else: | 417 | else: |
| 418 | storage = storage.cpu() | 418 | storage = storage.cpu() |
| 419 | + | ||
| 420 | + # Fallback for cases where tensor reduce path has already materialized | ||
| 421 | + # NPU tensors as CPU storages before reaching _npu_save. | ||
| 422 | + else: | ||
| 423 | + from torch.utils.serialization import config | ||
| 424 | + | ||
| 425 | + if ( | ||
| 426 | + config.save.use_pinned_memory_for_d2h | ||
| 427 | + and ( | ||
| 428 | + acc := torch.accelerator.current_accelerator( | ||
| 429 | + check_available=True | ||
| 430 | + ) | ||
| 431 | + ) | ||
| 432 | + is not None | ||
| 433 | + and acc.type == "npu" | ||
| 434 | + ): | ||
| 435 | + new_storage = torch.empty( | ||
| 436 | + num_bytes, dtype=torch.uint8, device="cpu", pin_memory=True | ||
| 437 | + ).untyped_storage() | ||
| 438 | + new_storage.copy_(storage) | ||
| 439 | + torch.accelerator.current_stream(storage.device.index).synchronize() | ||
| 440 | + storage = new_storage | ||
| 419 | # Now that it is on the CPU we can directly copy it into the zip file | 441 | # Now that it is on the CPU we can directly copy it into the zip file |
| 420 | zip_file.write_record(name, storage, num_bytes) | 442 | zip_file.write_record(name, storage, num_bytes) |
| 421 | 443 | ||