已合并
add save_async #22465
AtomGit-Bot创建于 2025年7月1日
add save_async #22465
已合并
从refs/pull/22465/head合入到master
共 3 个文件变更+266-5
| @@ -0,0 +1,119 @@ | |||
| 1 | +import os | ||
| 2 | +import time | ||
| 3 | +import copy | ||
| 4 | + | ||
| 5 | +import torch | ||
| 6 | +import torch.nn as nn | ||
| 7 | +import torch.optim as optim | ||
| 8 | + | ||
| 9 | +import torch_npu | ||
| 10 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 11 | +from torch_npu.utils._path_manager import PathManager | ||
| 12 | + | ||
| 13 | + | ||
| 14 | +class TestAsyncSave(TestCase): | ||
| 15 | + test_save_path = os.path.join( | ||
| 16 | + os.path.realpath(os.path.dirname(__file__)), "test_save_async") | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + def setUpClass(cls): | ||
| 20 | + PathManager.make_dir_safety(TestAsyncSave.test_save_path) | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + def tearDownClass(cls): | ||
| 24 | + PathManager.remove_path_safety(TestAsyncSave.test_save_path) | ||
| 25 | + | ||
| 26 | + def wait_for_save_completion(self, file_path, timeout_sec=60, poll_interval_sec=0.5): | ||
| 27 | + start_time = time.time() | ||
| 28 | + | ||
| 29 | + while time.time() - start_time < timeout_sec: | ||
| 30 | + if os.path.exists(file_path): | ||
| 31 | + current_size = os.path.getsize(file_path) | ||
| 32 | + time.sleep(poll_interval_sec) | ||
| 33 | + new_size = os.path.getsize(file_path) | ||
| 34 | + | ||
| 35 | + if current_size == new_size: | ||
| 36 | + return True | ||
| 37 | + else: | ||
| 38 | + time.sleep(poll_interval_sec) | ||
| 39 | + | ||
| 40 | + return False | ||
| 41 | + | ||
| 42 | + def test_save_async_tensor(self): | ||
| 43 | + save_tensor = torch.rand(1024, dtype=torch.float32).npu() | ||
| 44 | + async_save_path = os.path.join(TestAsyncSave.test_save_path, "async_save_tensor.pt") | ||
| 45 | + torch_npu.utils.save_async(save_tensor, async_save_path) | ||
| 46 | + | ||
| 47 | + if self.wait_for_save_completion(async_save_path): | ||
| 48 | + tensor_async = torch.load(async_save_path, weights_only=False) | ||
| 49 | + self.assertEqual(tensor_async, save_tensor) | ||
| 50 | + else: | ||
| 51 | + self.assertTrue(False, f"{async_save_path} is not exist!") | ||
| 52 | + | ||
| 53 | + def test_save_async(self): | ||
| 54 | + loss1 = [1.6099495, 1.6099086, 1.6098710] | ||
| 55 | + loss2 = [] | ||
| 56 | + model_list = [] | ||
| 57 | + checkpoint_list = [] | ||
| 58 | + model_origin = nn.Sequential( | ||
| 59 | + nn.Linear(100, 50), | ||
| 60 | + nn.ReLU(), | ||
| 61 | + nn.Linear(50, 20), | ||
| 62 | + nn.ReLU(), | ||
| 63 | + nn.Linear(20, 5), | ||
| 64 | + nn.ReLU() | ||
| 65 | + ) | ||
| 66 | + | ||
| 67 | + input_data = torch.ones(6400, 100).npu() | ||
| 68 | + labels = torch.arange(5).repeat(1280).npu() | ||
| 69 | + | ||
| 70 | + criterion = nn.CrossEntropyLoss() | ||
| 71 | + model = model_origin.npu() | ||
| 72 | + optimerizer = optim.SGD(model.parameters(), lr=0.1) | ||
| 73 | + for step in range(3): | ||
| 74 | + outputs = model(input_data) | ||
| 75 | + loss = criterion(outputs, labels) | ||
| 76 | + | ||
| 77 | + optimerizer.zero_grad() | ||
| 78 | + loss.backward() | ||
| 79 | + | ||
| 80 | + optimerizer.step() | ||
| 81 | + | ||
| 82 | + loss2.append(loss) | ||
| 83 | + checkpoint = { | ||
| 84 | + "model": model.state_dict(), | ||
| 85 | + "optimizer": optimerizer.state_dict() | ||
| 86 | + } | ||
| 87 | + checkpoint_list.append(copy.deepcopy(checkpoint)) | ||
| 88 | + model_list.append(copy.deepcopy(model)) | ||
| 89 | + checkpoint_async_path = os.path.join(TestAsyncSave.test_save_path, f"checkpoint_async_{step}.path") | ||
| 90 | + model_async_path = os.path.join(TestAsyncSave.test_save_path, f"model_async_{step}.path") | ||
| 91 | + torch_npu.utils.save_async(checkpoint, checkpoint_async_path, model=model) | ||
| 92 | + torch_npu.utils.save_async(model, model_async_path, model=model) | ||
| 93 | + | ||
| 94 | + for i in range(3): | ||
| 95 | + self.assertEqual(loss1[i], loss2[i].item()) | ||
| 96 | + checkpoint_async_path = os.path.join(TestAsyncSave.test_save_path, f"checkpoint_async_{i}.path") | ||
| 97 | + if self.wait_for_save_completion(checkpoint_async_path): | ||
| 98 | + checkpoint_async = torch.load(checkpoint_async_path, weights_only=False) | ||
| 99 | + self.assertEqual(checkpoint_list[i], checkpoint_async, prec=2e-3) | ||
| 100 | + else: | ||
| 101 | + self.assertTrue(False, f"{checkpoint_async_path} is not exist!") | ||
| 102 | + model_async_path = os.path.join(TestAsyncSave.test_save_path, f"model_async_{i}.path") | ||
| 103 | + if self.wait_for_save_completion(model_async_path): | ||
| 104 | + model_async = torch.load(model_async_path, weights_only=False) | ||
| 105 | + else: | ||
| 106 | + self.assertTrue(False, f"{model_async_path} is not exist!") | ||
| 107 | + state_dict_sync = model_list[i].state_dict() | ||
| 108 | + state_dict_async = model_async.state_dict() | ||
| 109 | + | ||
| 110 | + key_sync = sorted(state_dict_sync.keys()) | ||
| 111 | + key_async = sorted(state_dict_async.keys()) | ||
| 112 | + | ||
| 113 | + self.assertEqual(key_sync, key_async) | ||
| 114 | + for key in key_async: | ||
| 115 | + self.assertEqual(state_dict_async[key], state_dict_sync[key], prec=2e-3) | ||
| 116 | + | ||
| 117 | +if __name__ == '__main__': | ||
| 118 | + torch.npu.set_device(0) | ||
| 119 | + run_tests() | ||
| @@ -1,12 +1,12 @@ | |||
| 1 | __all__ = ["npu_combine_tensors", "get_part_combined_tensor", "is_combined_tensor_valid", "FlopsCounter", | 1 | __all__ = ["npu_combine_tensors", "get_part_combined_tensor", "is_combined_tensor_valid", "FlopsCounter", |
| 2 | - "set_thread_affinity", "reset_thread_affinity"] | 2 | + "set_thread_affinity", "reset_thread_affinity", "save_async"] |
| 3 | 3 | ||
| 4 | from torch_npu import _C | 4 | from torch_npu import _C |
| 5 | from ._module import _apply_module_patch | 5 | from ._module import _apply_module_patch |
| 6 | from .tensor_methods import _add_tensor_methods | 6 | from .tensor_methods import _add_tensor_methods |
| 7 | from .storage import _add_storage_methods | 7 | from .storage import _add_storage_methods |
| 8 | from .combine_tensors import npu_combine_tensors, get_part_combined_tensor, is_combined_tensor_valid | 8 | from .combine_tensors import npu_combine_tensors, get_part_combined_tensor, is_combined_tensor_valid |
| 9 | -from .serialization import _add_serialization_methods | 9 | +from .serialization import _add_serialization_methods, save_async |
| 10 | from .npu_intercept import _cann_package_check, _add_intercept_methods | 10 | from .npu_intercept import _cann_package_check, _add_intercept_methods |
| 11 | from .dtensor import _register_ops_under_dtensor_rules | 11 | from .dtensor import _register_ops_under_dtensor_rules |
| 12 | from .collect_env import _add_collect_env_methods | 12 | from .collect_env import _add_collect_env_methods |
| @@ -3,6 +3,7 @@ import io | |||
| 3 | import sys | 3 | import sys |
| 4 | import pickle | 4 | import pickle |
| 5 | import re | 5 | import re |
| 6 | +import threading | ||
| 6 | from typing import Any, Optional | 7 | from typing import Any, Optional |
| 7 | 8 | ||
| 8 | import torch | 9 | import torch |
| @@ -10,16 +11,17 @@ from torch.serialization import _check_dill_version, _open_file_like, _is_zipfil | |||
| 10 | _open_zipfile_reader, _is_torchscript_zip, _weights_only_unpickler, \ | 11 | _open_zipfile_reader, _is_torchscript_zip, _weights_only_unpickler, \ |
| 11 | _legacy_load, _load, FileLike, MAP_LOCATION, DEFAULT_PROTOCOL, \ | 12 | _legacy_load, _load, FileLike, MAP_LOCATION, DEFAULT_PROTOCOL, \ |
| 12 | normalize_storage_type, location_tag, _serialization_tls, _get_storage_alignment | 13 | normalize_storage_type, location_tag, _serialization_tls, _get_storage_alignment |
| 13 | -from torch.serialization import _default_to_weights_only, UNSAFE_MESSAGE | 14 | +from torch.serialization import _default_to_weights_only, UNSAFE_MESSAGE, _open_zipfile_writer |
| 14 | 15 | ||
| 15 | import torch_npu | 16 | import torch_npu |
| 16 | from torch_npu.utils._error_code import ErrCode, pta_error | 17 | from torch_npu.utils._error_code import ErrCode, pta_error |
| 17 | from .utils import _should_print_warning | 18 | from .utils import _should_print_warning |
| 18 | 19 | ||
| 20 | +__all__ = ["load", "save", "save_async"] | ||
| 21 | + | ||
| 19 | ALWAYS_WARN_LEGACY_SERIALIZATION = False | 22 | ALWAYS_WARN_LEGACY_SERIALIZATION = False |
| 20 | RE_MAP_CPU = False | 23 | RE_MAP_CPU = False |
| 21 | - | 24 | +save_async_stream_map = {} |
| 22 | -__all__ = ["load", "save"] | ||
| 23 | 25 | ||
| 24 | 26 | ||
| 25 | def _get_always_warn_legacy_serialization(): | 27 | def _get_always_warn_legacy_serialization(): |
| @@ -429,6 +431,146 @@ def save( | |||
| 429 | return torch.serialization.save(obj, f, pickle_module, pickle_protocol, True, _disable_byteorder_record) | 431 | return torch.serialization.save(obj, f, pickle_module, pickle_protocol, True, _disable_byteorder_record) |
| 430 | 432 | ||
| 431 | 433 | ||
| 434 | +def save_async( | ||
| 435 | + obj: object, | ||
| 436 | + f, | ||
| 437 | + pickle_module: Any = pickle, | ||
| 438 | + pickle_protocol: int = DEFAULT_PROTOCOL, | ||
| 439 | + _use_new_zipfile_serialization: bool = True, | ||
| 440 | + _disable_byteorder_record: bool = False, | ||
| 441 | + model: torch.nn.Module = None | ||
| 442 | +) -> None: | ||
| 443 | + if _use_new_zipfile_serialization is False: | ||
| 444 | + raise RuntimeError("Error: torch_npu.save_async with \"_use_new_zipfile_serialization = False\"\ | ||
| 445 | + is not recommended for npu tensor, which may bring unexpected errors and hopefully \ | ||
| 446 | + set \"_use_new_zipfile_serialization = True\"", | ||
| 447 | + "if it is necessary to use this, please convert the npu tensor to cpu tensor for saving" + | ||
| 448 | + pta_error(ErrCode.PARAM)) | ||
| 449 | + | ||
| 450 | + _check_dill_version(pickle_module) | ||
| 451 | + save_args = (obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization, _disable_byteorder_record) | ||
| 452 | + | ||
| 453 | + device = torch.npu.current_device() | ||
| 454 | + save_thread = threading.Thread(target=_save_data_thread, args=(save_args, device, model)) | ||
| 455 | + save_thread.start() | ||
| 456 | + | ||
| 457 | + | ||
| 458 | +def _save_data_thread(save_args, | ||
| 459 | + device, | ||
| 460 | + model: torch.nn.Module = None): | ||
| 461 | + global save_async_stream_map | ||
| 462 | + torch.npu.set_device(device) | ||
| 463 | + | ||
| 464 | + def hook_fn(*args): | ||
| 465 | + torch.npu.current_stream().wait_stream(save_async_stream_map.get(device)) | ||
| 466 | + | ||
| 467 | + if device not in save_async_stream_map: | ||
| 468 | + save_async_stream = torch.npu.Stream() | ||
| 469 | + save_async_stream_map[device] = save_async_stream | ||
| 470 | + if isinstance(model, torch.nn.Module): | ||
| 471 | + model.register_full_backward_hook(hook_fn) | ||
| 472 | + else: | ||
| 473 | + save_async_stream = save_async_stream_map[device] | ||
| 474 | + | ||
| 475 | + obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization, _disable_byteorder_record = save_args | ||
| 476 | + with torch.npu.stream(save_async_stream): | ||
| 477 | + data_value, serialized_storages = _save(obj, pickle_module, pickle_protocol) | ||
| 478 | + storage_value = [] | ||
| 479 | + for key in sorted(serialized_storages.keys()): | ||
| 480 | + name = f'data/{key}' | ||
| 481 | + storage = serialized_storages.get(key) | ||
| 482 | + # given that we copy things around anyway, we might use storage.cpu() | ||
| 483 | + # this means to that to get tensors serialized, you need to implement | ||
| 484 | + # .cpu() on the underlying Storage | ||
| 485 | + if storage.device.type != 'cpu': | ||
| 486 | + storage = storage.cpu() | ||
| 487 | + # Now that it is on the CPU we can directly copy it into the zip file | ||
| 488 | + if storage.device.type != "cpu": | ||
| 489 | + storage_tensor = torch_npu._C._tensor_construct_from_storage(storage) | ||
| 490 | + num_bytes = storage_tensor.size().numel() * storage_tensor.element_size() | ||
| 491 | + else: | ||
| 492 | + num_bytes = storage.nbytes() | ||
| 493 | + storage_value.append((name, storage, num_bytes)) | ||
| 494 | + | ||
| 495 | + with _open_zipfile_writer(f) as opened_zipfile: | ||
| 496 | + opened_zipfile.write_record('data.pkl', data_value, len(data_value)) | ||
| 497 | + | ||
| 498 | + for name, storage, num_bytes in storage_value: | ||
| 499 | + opened_zipfile.write_record(name, storage.data_ptr(), num_bytes) | ||
| 500 | + | ||
| 501 | + | ||
| 502 | +def _save(obj, pickle_module, pickle_protocol): | ||
| 503 | + serialized_storages = {} | ||
| 504 | + id_map: Dict[int, str] = {} | ||
| 505 | + | ||
| 506 | + # Since loading storages that view the same data with different dtypes is | ||
| 507 | + # not supported, we need to keep track of the dtype associated with each | ||
| 508 | + # storage data_ptr and throw an error if the dtype is ever different. | ||
| 509 | + storage_dtypes: Dict[int, torch.dtype] = {} | ||
| 510 | + | ||
| 511 | + def persistent_id(obj): | ||
| 512 | + if isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj): | ||
| 513 | + | ||
| 514 | + if isinstance(obj, torch.storage.TypedStorage): | ||
| 515 | + storage = obj._untyped_storage | ||
| 516 | + storage_dtype = obj.dtype | ||
| 517 | + storage_type_str = obj._pickle_storage_type() | ||
| 518 | + storage_type = getattr(torch, storage_type_str) | ||
| 519 | + if storage.device.type != "cpu": | ||
| 520 | + storage_tensor = torch_npu._C._tensor_construct_from_storage(storage) | ||
| 521 | + storage_numel = storage_tensor.size().numel() * storage_tensor.element_size() // obj._element_size() | ||
| 522 | + else: | ||
| 523 | + storage_numel = obj._size() | ||
| 524 | + | ||
| 525 | + else: | ||
| 526 | + storage = obj | ||
| 527 | + storage_dtype = torch.uint8 | ||
| 528 | + storage_type = normalize_storage_type(type(obj)) | ||
| 529 | + if storage.device.type != "cpu": | ||
| 530 | + storage_tensor = torch_npu._C._tensor_construct_from_storage(storage) | ||
| 531 | + storage_numel = storage_tensor.size().numel() * storage_tensor.element_size() | ||
| 532 | + else: | ||
| 533 | + storage_numel = storage.nbytes() | ||
| 534 | + | ||
| 535 | + # If storage is allocated, ensure that any other saved storages | ||
| 536 | + # pointing to the same data all have the same dtype. If storage is | ||
| 537 | + # not allocated, don't perform this check | ||
| 538 | + if storage.data_ptr() != 0: | ||
| 539 | + if storage.data_ptr() in storage_dtypes: | ||
| 540 | + if storage_dtype != storage_dtypes[storage.data_ptr()]: | ||
| 541 | + raise RuntimeError( | ||
| 542 | + 'Cannot save multiple tensors or storages that ' | ||
| 543 | + 'view the same data as different types' + pta_error(ErrCode.VALUE)) | ||
| 544 | + else: | ||
| 545 | + storage_dtypes[storage.data_ptr()] = storage_dtype | ||
| 546 | + | ||
| 547 | + storage_key = id_map.setdefault(storage._cdata, str(len(id_map))) | ||
| 548 | + location = location_tag(storage) | ||
| 549 | + serialized_storages[storage_key] = storage | ||
| 550 | + | ||
| 551 | + return ('storage', | ||
| 552 | + storage_type, | ||
| 553 | + storage_key, | ||
| 554 | + location, | ||
| 555 | + storage_numel) | ||
| 556 | + | ||
| 557 | + return None | ||
| 558 | + | ||
| 559 | + # Write the pickle data for `obj` | ||
| 560 | + data_buf = io.BytesIO() | ||
| 561 | + pickler = pickle_module.Pickler(data_buf, protocol=pickle_protocol) | ||
| 562 | + pickler.persistent_id = persistent_id | ||
| 563 | + if isinstance(obj, torch.nn.Module): | ||
| 564 | + hook_handle = obj._backward_hooks.copy() | ||
| 565 | + obj._backward_hooks.clear() | ||
| 566 | + pickler.dump(obj) | ||
| 567 | + obj._backward_hooks.update(hook_handle) | ||
| 568 | + else: | ||
| 569 | + pickler.dump(obj) | ||
| 570 | + data_value = data_buf.getvalue() | ||
| 571 | + return data_value, serialized_storages | ||
| 572 | + | ||
| 573 | + | ||
| 432 | def _add_serialization_methods(): | 574 | def _add_serialization_methods(): |
| 433 | torch.save = save | 575 | torch.save = save |
| 434 | torch.load = load | 576 | torch.load = load |