已合并
fix_stash_obj_in_tls_bug #35969
Ambi创建于 5月18日
fix_stash_obj_in_tls_bug #35969
已合并
从已删除 :v2.7.1合入到Ascend/pytorchv2.7.1
共 3 个文件变更+10-10
| @@ -29,7 +29,7 @@ class TestRunWithRngState(TestUtils): | |||
| 29 | 29 | ||
| 30 | return res1, res2 | 30 | return res1, res2 |
| 31 | 31 | ||
| 32 | - @parametrize("shape", [(10,)]) | 32 | + @parametrize("shape", [(20,)]) |
| 33 | 33 | ||
| 34 | def test_rng_state_with_compile(self, shape, dtype): | 34 | def test_rng_state_with_compile(self, shape, dtype): |
| 35 | device = "npu" | 35 | device = "npu" |
| @@ -221,7 +221,7 @@ class TestTreeManagerIntegration(TestCase): | |||
| 221 | def test_reset_npugraph_trees(self): | 221 | def test_reset_npugraph_trees(self): |
| 222 | get_container(0) # Initialize a container | 222 | get_container(0) # Initialize a container |
| 223 | reset_npugraph_trees() | 223 | reset_npugraph_trees() |
| 224 | - container_dict = getattr(local, "tree_manager_containers", {}) | 224 | + container_dict = getattr(local, "npu_tree_manager_containers", {}) |
| 225 | self.assertEqual(len(container_dict), 0) | 225 | self.assertEqual(len(container_dict), 0) |
| 226 | 226 | ||
| 227 | 227 | ||
| @@ -255,8 +255,8 @@ class TreeManagerContainer: | |||
| 255 | local = threading.local() | 255 | local = threading.local() |
| 256 | 256 | ||
| 257 | # one tree manager per device | 257 | # one tree manager per device |
| 258 | -local.tree_manager_containers = {} | 258 | +local.npu_tree_manager_containers = {} |
| 259 | -local.tree_manager_locks = defaultdict(threading.Lock) | 259 | +local.npu_tree_manager_locks = defaultdict(threading.Lock) |
| 260 | 260 | ||
| 261 | 261 | ||
| 262 | # only incremented by user call of mark_step_begin | 262 | # only incremented by user call of mark_step_begin |
| @@ -266,8 +266,8 @@ class MarkStepBox: | |||
| 266 | 266 | ||
| 267 | # We need to register this as an object that will be copied over as TLS when new | 267 | # We need to register this as an object that will be copied over as TLS when new |
| 268 | # threads are created in autograd | 268 | # threads are created in autograd |
| 269 | -torch._C._stash_obj_in_tls("tree_manager_containers", local.tree_manager_containers) | 269 | +torch._C._stash_obj_in_tls("npu_tree_manager_containers", local.npu_tree_manager_containers) |
| 270 | -torch._C._stash_obj_in_tls("tree_manager_locks", local.tree_manager_locks) | 270 | +torch._C._stash_obj_in_tls("npu_tree_manager_locks", local.npu_tree_manager_locks) |
| 271 | 271 | ||
| 272 | 272 | ||
| 273 | def mark_step_begin() -> None: | 273 | def mark_step_begin() -> None: |
| @@ -280,8 +280,8 @@ def mark_step_begin() -> None: | |||
| 280 | def reset_npugraph_trees() -> None: | 280 | def reset_npugraph_trees() -> None: |
| 281 | "Clear all npugraph trees" | 281 | "Clear all npugraph trees" |
| 282 | # see shutdown below for why this is necessary | 282 | # see shutdown below for why this is necessary |
| 283 | - container_dict = get_obj(local, "tree_manager_containers") | 283 | + container_dict = get_obj(local, "npu_tree_manager_containers") |
| 284 | - locks_dict = get_obj(local, "tree_manager_locks") | 284 | + locks_dict = get_obj(local, "npu_tree_manager_locks") |
| 285 | for device, lock in locks_dict.items(): | 285 | for device, lock in locks_dict.items(): |
| 286 | with lock: | 286 | with lock: |
| 287 | container = container_dict.get(device) | 287 | container = container_dict.get(device) |
| @@ -306,8 +306,8 @@ def get_obj(thread_local: Any, attr_name: str) -> Any: | |||
| 306 | 306 | ||
| 307 | 307 | ||
| 308 | def get_container(device_index: int) -> TreeManagerContainer: | 308 | def get_container(device_index: int) -> TreeManagerContainer: |
| 309 | - container_dict = get_obj(local, "tree_manager_containers") | 309 | + container_dict = get_obj(local, "npu_tree_manager_containers") |
| 310 | - lock = get_obj(local, "tree_manager_locks")[device_index] | 310 | + lock = get_obj(local, "npu_tree_manager_locks")[device_index] |
| 311 | 311 | ||
| 312 | with lock: | 312 | with lock: |
| 313 | if device_index not in container_dict: | 313 | if device_index not in container_dict: |