已合并
fix: pass map_location to _legacy_load in torch_npu.utils.load for torch.compile cache #43271
wuyouqi1创建于 7月29日
fix: pass map_location to _legacy_load in torch_npu.utils.load for torch.compile cache #43271
已合并
共 3 个文件变更+68-1
| @@ -6,8 +6,30 @@ from testutils import TestUtils | |||
| 6 | import torch_npu | 6 | import torch_npu |
| 7 | import torch_npu._inductor | 7 | import torch_npu._inductor |
| 8 | import unittest | 8 | import unittest |
| 9 | +import os | ||
| 10 | +import tempfile | ||
| 9 | 11 | ||
| 10 | class TestCodeCache(TestUtils): | 12 | class TestCodeCache(TestUtils): |
| 13 | + | ||
| 14 | + def setUp(self): | ||
| 15 | + super().setUp() | ||
| 16 | + self._saved_cache_dir = os.environ.get("TORCHINDUCTOR_CACHE_DIR") | ||
| 17 | + self._saved_fx_cache = torch._inductor.config.fx_graph_cache | ||
| 18 | + self._tmpdir = tempfile.TemporaryDirectory() | ||
| 19 | + cache_dir = os.path.join(self._tmpdir.name, ".inductor_cache") | ||
| 20 | + os.makedirs(cache_dir, exist_ok=True) | ||
| 21 | + os.environ["TORCHINDUCTOR_CACHE_DIR"] = cache_dir | ||
| 22 | + torch._inductor.config.fx_graph_cache = True | ||
| 23 | + | ||
| 24 | + def tearDown(self): | ||
| 25 | + if self._saved_cache_dir is None: | ||
| 26 | + os.environ.pop("TORCHINDUCTOR_CACHE_DIR", None) | ||
| 27 | + else: | ||
| 28 | + os.environ["TORCHINDUCTOR_CACHE_DIR"] = self._saved_cache_dir | ||
| 29 | + torch._inductor.config.fx_graph_cache = self._saved_fx_cache | ||
| 30 | + self._tmpdir.cleanup() | ||
| 31 | + super().tearDown() | ||
| 32 | + | ||
| 11 | 33 | ||
| 12 | def test_codecache(self): | 34 | def test_codecache(self): |
| 13 | device_properties = torch_npu.npu.get_device_properties( | 35 | device_properties = torch_npu.npu.get_device_properties( |
| @@ -23,6 +45,19 @@ class TestCodeCache(TestUtils): | |||
| 23 | self.assertEqual(system2["device"]["name"], device_properties.name) | 45 | self.assertEqual(system2["device"]["name"], device_properties.name) |
| 24 | self.assertEqual(system2["version"]["cann"], torch.version.cann) | 46 | self.assertEqual(system2["version"]["cann"], torch.version.cann) |
| 25 | 47 | ||
| 48 | + def test_fx_graph_cache_constant_device(self): | ||
| 49 | + """fx_graph_cache: tensor constants should remain on NPU after cache load.""" | ||
| 50 | + | ||
| 51 | + def fn(x): | ||
| 52 | + c = torch.tensor(list(range(15)), dtype=torch.int64) | ||
| 53 | + return x.index_select(0, c.to(x.device)) | ||
| 54 | + | ||
| 55 | + x = torch.randn(16, device="npu") | ||
| 56 | + expected = torch.compile(fn, backend="inductor")(x) | ||
| 57 | + torch._dynamo.reset() | ||
| 58 | + actual = torch.compile(fn, backend="inductor")(x) | ||
| 59 | + self.assertEqual(expected, actual) | ||
| 60 | + | ||
| 26 | 61 | ||
| 27 | if __name__ == "__main__": | 62 | if __name__ == "__main__": |
| 28 | run_tests() | 63 | run_tests() |
| @@ -6,6 +6,10 @@ import argparse | |||
| 6 | import torch | 6 | import torch |
| 7 | import torch.nn as nn | 7 | import torch.nn as nn |
| 8 | import torch.nn.functional as F | 8 | import torch.nn.functional as F |
| 9 | +from torch.testing._internal.common_utils import ( | ||
| 10 | + parametrize, | ||
| 11 | + instantiate_parametrized_tests, | ||
| 12 | +) | ||
| 9 | 13 | ||
| 10 | import torch_npu | 14 | import torch_npu |
| 11 | from torch_npu.testing.testcase import TestCase, run_tests | 15 | from torch_npu.testing.testcase import TestCase, run_tests |
| @@ -95,6 +99,31 @@ class TestSerialization(TestCase): | |||
| 95 | self.assertExpectedInline(f'{x_loaded.device.type}:{x_loaded.device.index}', 'npu:0') | 99 | self.assertExpectedInline(f'{x_loaded.device.type}:{x_loaded.device.index}', 'npu:0') |
| 96 | self.assertRtolEqual(x, x_loaded.cpu()) | 100 | self.assertRtolEqual(x, x_loaded.cpu()) |
| 97 | 101 | ||
| 102 | + | ||
H | |||
| 103 | + def test_legacy_load_maplocation(self, map_kind): | ||
| 104 | + """legacy + weights_only=False: cover None / Callable / Dict map_location.""" | ||
| 105 | + if map_kind == "none": | ||
| 106 | + x = torch.randn(5).npu() | ||
| 107 | + map_location = None | ||
| 108 | + elif map_kind == "callable": | ||
| 109 | + x = torch.randn(5) | ||
| 110 | + map_location = lambda storage, loc: storage.to(device="npu:0") | ||
| 111 | + else: | ||
| 112 | + x = torch.randn(5) | ||
| 113 | + map_location = {"cpu": "npu:0"} | ||
| 114 | + | ||
| 115 | + with tempfile.TemporaryDirectory() as tmpdir: | ||
| 116 | + path = os.path.join(tmpdir, "data.pt") | ||
| 117 | + torch.serialization.save(x, path, _use_new_zipfile_serialization=False) | ||
| 118 | + y = torch.load( | ||
| 119 | + path, | ||
| 120 | + map_location=map_location, | ||
| 121 | + weights_only=False, | ||
| 122 | + mmap=False, | ||
| 123 | + ) | ||
| 124 | + self.assertEqual(str(y.device), "npu:0") | ||
| 125 | + self.assertRtolEqual(x.cpu() if x.is_npu else x, y.cpu()) | ||
| 126 | + | ||
| 98 | def test_save_npu_format(self): | 127 | def test_save_npu_format(self): |
| 99 | with tempfile.TemporaryDirectory() as tmpdir: | 128 | with tempfile.TemporaryDirectory() as tmpdir: |
| 100 | path = os.path.join(tmpdir, 'data.pt') | 129 | path = os.path.join(tmpdir, 'data.pt') |
| @@ -275,5 +304,8 @@ class TestSerialization(TestCase): | |||
| 275 | save_load_check(a, b) | 304 | save_load_check(a, b) |
| 276 | 305 | ||
| 277 | 306 | ||
| 307 | +instantiate_parametrized_tests(TestSerialization) | ||
| 308 | + | ||
| 309 | + | ||
| 278 | if __name__ == "__main__": | 310 | if __name__ == "__main__": |
| 279 | run_tests() | 311 | run_tests() |
| @@ -336,7 +336,7 @@ def load( | |||
| 336 | return _remap_result(cpu_result, map_location) | 336 | return _remap_result(cpu_result, map_location) |
| 337 | else: | 337 | else: |
| 338 | return _legacy_load( | 338 | return _legacy_load( |
| 339 | - opened_file, "cpu", pickle_module, **pickle_load_args | 339 | + opened_file, map_location, pickle_module, **pickle_load_args |
| 340 | ) | 340 | ) |
| 341 | 341 | ||
| 342 | 342 | ||
建议再补一个最小的 pickle.dumps/pickle.loads 用例，覆盖Inductor 缓存实际经过的 UntypedStorage._load_from_bytes -> torch.load(BytesIO,map_location=None) 路径，避免仅文件路径加载通过而缓存反序列化回归