refactor: remove unsupported api interception patches-pytorch-AtomGit

Mtest/jit/test_script_module.py+6-27

	@@ -490,9 +490,8 @@ class TestScriptModuleExtraRepr(TestCase):
490	class TestScriptModuleShareMemory(TestCase):	490	class TestScriptModuleShareMemory(TestCase):
491	"""share_memory behavior differs by device:	491	"""share_memory behavior differs by device:
492	CPU: works, makes storage shared.	492	CPU: works, makes storage shared.
493	- GPU/CUDA: no-op (per torch.Tensor.share_memory_ docstring).	493	+ NPU/CUDA: no-op (per torch.Tensor.share_memory_ docstring).
494	- NPU: torch-npu intercepts with RuntimeError.	494	+ Tests document actual NPU behavior and CPU baseline"""
495	- Tests document actual NPU behavior and CPU baseline."""
496		495
497	def test_share_memory_cpu_returns_self(self):	496	def test_share_memory_cpu_returns_self(self):
498	sm = _make_cpu_linear()	497	sm = _make_cpu_linear()
	@@ -510,34 +509,14 @@ class TestScriptModuleShareMemory(TestCase):
510	sm.share_memory()	509	sm.share_memory()
511	self.assertTrue(sm.linear.weight.untyped_storage().is_shared())	510	self.assertTrue(sm.linear.weight.untyped_storage().is_shared())
512		511
513	- def test_share_memory_on_npu_raises(self):	512	+ def test_share_memory_npu_noop(self):
514	sm = _make_linear()	513	sm = _make_linear()
515	- with self.assertRaisesRegex(	514	+ result = sm.share_memory()
516	- RuntimeError, r"share_memory.*not supported in npu"):	515	+ self.assertIs(result, sm)
517	- sm.share_memory()	516	+ self.assertEqual(sm.linear.weight.device.type, device_type)
518		517
519		518
520	class TestScriptModuleMetadata(TestCase):	519	class TestScriptModuleMetadata(TestCase):
521	- """register_module/register_parameter on NPU:
522	- torch-npu intercepts with RuntimeError.
523	- On CPU they also raise RuntimeError (PyTorch limitation:
524	- "Cannot re-assign modules" / "Can't add a new parameter
525	- after ScriptModule construction")."""
526	-
527	- def test_register_module_raises_on_npu(self):
528	- sm = _make_linear()
529	- sub = nn.Linear(2, 2).to(device_type)
530	- with self.assertRaisesRegex(
531	- RuntimeError, r"register_module.*not supported in npu"):
532	- sm.register_module("new_sub", sub)
533	-
534	- def test_register_parameter_raises_on_npu(self):
535	- sm = _make_linear()
536	- param = nn.Parameter(torch.randn(2, 2)).to(device_type)
537	- with self.assertRaisesRegex(
538	- RuntimeError, r"register_parameter.*not supported in npu"):
539	- sm.register_parameter("new_param", param)
540	-
541	def test_set_submodule_raises(self):	520	def test_set_submodule_raises(self):
542	sm = _make_linear()	521	sm = _make_linear()
543	new_sub = nn.Linear(2, 2).to(device_type)	522	new_sub = nn.Linear(2, 2).to(device_type)

Mtest/npu/test_unsupport_api.py+4-52

	@@ -1,10 +1,8 @@
1	import torch	1	import torch
2	import torch.nn as nn	2	import torch.nn as nn
3		3
4	-import torch_npu
5		4
6	from torch_npu.testing.testcase import TestCase, run_tests	5	from torch_npu.testing.testcase import TestCase, run_tests
7	-from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
8		6
9		7
10	class SimpleModel(nn.Module):	8	class SimpleModel(nn.Module):
	@@ -16,15 +14,6 @@ class SimpleModel(nn.Module):
16	return self.fc(x)	14	return self.fc(x)
17		15
18		16
19	-class ScriptModel(nn.Module):
20	- def __init__(self):
21	- super(ScriptModel, self).__init__()
22	- self.linear = torch.nn.Linear(4, 4)
23	-
24	- def forward(self, x, h):
25	- return torch.tanh(self.linear(x) + h)
26	-
27	-
28	class TestPtaUnsupportApi(TestCase):	17	class TestPtaUnsupportApi(TestCase):
29		18
30	def test_crow_indices(self):	19	def test_crow_indices(self):
	@@ -67,47 +56,10 @@ class TestPtaUnsupportApi(TestCase):
67	coalesce_tensor = sparse_tensor.coalesce().npu()	56	coalesce_tensor = sparse_tensor.coalesce().npu()
68	coalesce_tensor.ccol_indices()	57	coalesce_tensor.ccol_indices()
69		58
70	- def test_Module_share_memory_runtimeerror(self):	59	+ def test_Module_share_memory_npu(self):
71	- with self.assertRaisesRegex(RuntimeError, r"(.*) is not supported in npu."):	60	+ model = SimpleModel().npu()
72	- model = SimpleModel().npu()	61	+ model.share_memory()
73	- model.share_memory()	62	+ self.assertEqual(model.fc.weight.device.type, "npu")
74	-
75	- def test_ScriptModule_register_parameter_runtimeerror(self):
76	- with self.assertRaisesRegex(RuntimeError, r"(.*) is not supported in npu."):
77	- model = ScriptModel().npu()
78	- x, h = torch.rand(3, 4).npu(), torch.rand(3, 4).npu()
79	- traced_cell = torch.jit.trace(model, (x, h))
80	- traced_cell.register_parameter("test_parameter", torch.nn.Parameter(torch.ones(1, 1)))
81	-
82	- def test_ScriptModule_add_module_runtimeerror(self):
83	- with self.assertRaisesRegex(RuntimeError, r"(.*) is not supported in npu."):
84	- model = ScriptModel().npu()
85	- x, h = torch.rand(3, 4).npu(), torch.rand(3, 4).npu()
86	- traced_cell = torch.jit.trace(model, (x, h))
87	- extra_linear = nn.Linear(5, 2)
88	- traced_cell.add_module("extra_linear", extra_linear)
C C chujinjin8月6日已过期这里社区有相关用例吗？ Z zhounan 8月11日评论： 8月11日评论：
89	-
90	- def test_ScriptModule_register_buffer_runtimeerror(self):
91	- with self.assertRaisesRegex(RuntimeError, r"(.*) is not supported in npu."):
92	- model = ScriptModel().npu()
93	- x, h = torch.rand(3, 4).npu(), torch.rand(3, 4).npu()
94	- traced_cell = torch.jit.trace(model, (x, h))
95	- traced_cell.register_buffer("test_buff", torch.zeros(3))
96	-
97	- def test_ScriptModule_register_module_runtimeerror(self):
98	- with self.assertRaisesRegex(RuntimeError, r"(.*) is not supported in npu."):
99	- model = ScriptModel().npu()
100	- x, h = torch.rand(3, 4).npu(), torch.rand(3, 4).npu()
101	- traced_cell = torch.jit.trace(model, (x, h))
102	- extra_linear = nn.Linear(5, 2)
103	- traced_cell.register_module("extra_linear", extra_linear)
104	-
105	- def test_ScriptModule_bfloat16_runtimeerror(self):
106	- with self.assertRaisesRegex(RuntimeError, r"(.*) is not supported in npu."):
107	- model = ScriptModel().npu()
108	- x, h = torch.rand(3, 4).npu(), torch.rand(3, 4).npu()
109	- traced_cell = torch.jit.trace(model, (x, h))
110	- traced_cell.bfloat16()
111		63
112	def test_Tensor_is_shared(self):	64	def test_Tensor_is_shared(self):
113	input_tensor = torch.tensor([1, 2, 3])	65	input_tensor = torch.tensor([1, 2, 3])

Mtorch_npu/_init/patches/npu_patches.py+0-7

	@@ -1,13 +1,6 @@
1	from torch_npu._init.patches.patch_manager import PatchManager	1	from torch_npu._init.patches.patch_manager import PatchManager
2		2
3		3
4	-@PatchManager.register_patch("npu")
5	-def apply_npu_intercept_patch():
6	- from torch_npu.utils.npu_intercept import _add_intercept_methods
7	-
8	- _add_intercept_methods()
9	-
10	-
11	@PatchManager.register_patch("npu")	4	@PatchManager.register_patch("npu")
12	def apply_npu_format_patch():	5	def apply_npu_format_patch():
13	from torch_npu.npu._format import _apply_npu_format_patch	6	from torch_npu.npu._format import _apply_npu_format_patch

Mtorch_npu/_init/registry/registry_manager.py+62-2

	@@ -1,7 +1,68 @@
1	-import torch	1	+import os
		2	+import torch
2	from torch.utils.checkpoint import DefaultDeviceType	3	from torch.utils.checkpoint import DefaultDeviceType
3		4
4	import torch_npu	5	import torch_npu
		6	+from torch_npu.utils._error_code import ErrCode, pta_error
		7	+from torch_npu.utils.collect_env import get_cann_version
		8	+
		9	+
		10	+cann_pytorch_version_map = {
		11	+ "6.3.RC2": ["1.8.1.post2", "1.11.0.post1", "2.0.0.rc1"],
		12	+ "6.3.RC1": ["1.8.1.post1", "1.11.0"],
		13	+ "6.1.RC1": ["1.8.1.post1", "1.11.0"],
		14	+ "6.0.1": ["1.8.1", "1.11.0.rc2"],
		15	+ "6.0.RC1": ["1.8.1", "1.11.0.rc1"]
		16	+}
		17	+
		18	+
		19	+def _cann_package_check():
		20	+ if "ASCEND_HOME_PATH" in os.environ:
		21	+ ascend_home_path = os.environ["ASCEND_HOME_PATH"]
		22	+ if not os.path.exists(ascend_home_path):
		23	+ raise Exception(f"ASCEND_HOME_PATH : {ascend_home_path} does not exist. "
		24	+ "Please run 'source set_env.sh' in the CANN installation path." +
		25	+ pta_error(ErrCode.NOT_FOUND))
		26	+
		27	+ # check whether environment variables are correctly configured
		28	+ if "ASCEND_OPP_PATH" not in os.environ:
		29	+ raise Exception("ASCEND_OPP_PATH environment variable is not set. "
		30	+ "Please check whether the opp package has been installed. If exist, please run "
		31	+ "'source set_env.sh' in the CANN installation path." +
		32	+ pta_error(ErrCode.NOT_FOUND))
		33	+
		34	+ ascend_opp_path = os.environ["ASCEND_OPP_PATH"]
		35	+ if not os.path.exists(ascend_opp_path):
		36	+ raise Exception(f"ASCEND_OPP_PATH : {ascend_opp_path} does not exist. "
		37	+ "Please check whether the opp package has been installed. If exist, please run "
		38	+ "'source set_env.sh' in the CANN installation path." +
		39	+ pta_error(ErrCode.NOT_FOUND))
		40	+
		41	+ ascend_runtime_path = os.path.join(ascend_home_path, "runtime")
		42	+ if not os.path.exists(ascend_runtime_path):
		43	+ raise Exception(f"ASCEND_RUNTIME_PATH : {ascend_runtime_path} does not exist. "
		44	+ "Please check whether the runtime package has been installed. If exist, please run "
		45	+ "'source set_env.sh' in the CANN installation path." +
		46	+ pta_error(ErrCode.NOT_FOUND))
		47	+
		48	+ ascend_compiler_path = os.path.join(ascend_home_path, "compiler")
		49	+ if not os.path.exists(ascend_compiler_path):
		50	+ raise Exception(f"ASCEND_COMPILER_PATH : {ascend_compiler_path} does not exist. "
		51	+ "Please check whether the compiler package has been installed. If exist, please run "
		52	+ "'source set_env.sh' in the CANN installation path." +
		53	+ pta_error(ErrCode.NOT_FOUND))
		54	+
		55	+ # get the cann version
		56	+ cann_version = get_cann_version()
		57	+
		58	+ # check whether the CANN package version matches the pytorch version
		59	+ if cann_version in cann_pytorch_version_map and \
		60	+ torch_npu.__version__ not in cann_pytorch_version_map[cann_version]:
		61	+ print(f"Warning: CANN package version {cann_version} and PyTorch version {torch_npu.__version__} "
		62	+ "do not match. Please check the README of the Ascend PyTorch repo.")
		63	+ else:
		64	+ print("Warning: ASCEND_HOME_PATH environment variable is not set.")
		65	+
5		66
6	def _register_npu_backend():	67	def _register_npu_backend():
7	"""	68	"""
	@@ -16,7 +77,6 @@ def _register_npu_backend():
16	NPU runtime initialization is ownde by torch_npu.npu._lazy_init().	77	NPU runtime initialization is ownde by torch_npu.npu._lazy_init().
17	"""	78	"""
18	from torch_npu._init.registry.backend import register_privateuse1_backend	79	from torch_npu._init.registry.backend import register_privateuse1_backend
19	- from torch_npu.utils.npu_intercept import _cann_package_check
20		80
21	register_privateuse1_backend()	81	register_privateuse1_backend()
22	_cann_package_check()	82	_cann_package_check()

Dtorch_npu/utils/npu_intercept.py+0-103

	@@ -1,103 +0,0 @@
1	-import os
2	-
3	-from functools import wraps
4	-
5	-import torch
6	-import torch_npu
7	-from torch_npu.utils._error_code import ErrCode, pta_error
8	-from .unsupport_api import unsupported_Tensor_api, unsupported_nn_api
9	-from .collect_env import get_cann_version
10	-
11	-
12	-cann_pytorch_version_map = {
13	- "6.3.RC2": ["1.8.1.post2", "1.11.0.post1", "2.0.0.rc1"],
14	- "6.3.RC1": ["1.8.1.post1", "1.11.0"],
15	- "6.1.RC1": ["1.8.1.post1", "1.11.0"],
16	- "6.0.1": ["1.8.1", "1.11.0.rc2"],
17	- "6.0.RC1": ["1.8.1", "1.11.0.rc1"]
18	-}
19	-
20	-
21	-__all__ = []
22	-
23	-
24	-def _cann_package_check():
25	- if "ASCEND_HOME_PATH" in os.environ:
26	- ascend_home_path = os.environ["ASCEND_HOME_PATH"]
27	- if not os.path.exists(ascend_home_path):
28	- raise Exception(f"ASCEND_HOME_PATH : {ascend_home_path} does not exist. "
29	- "Please run 'source set_env.sh' in the CANN installation path." +
30	- pta_error(ErrCode.NOT_FOUND))
31	-
32	- # check whether environment variables are correctly configured
33	- if "ASCEND_OPP_PATH" not in os.environ:
34	- raise Exception("ASCEND_OPP_PATH environment variable is not set. "
35	- "Please check whether the opp package has been installed. If exist, please run "
36	- "'source set_env.sh' in the CANN installation path." +
37	- pta_error(ErrCode.NOT_FOUND))
38	-
39	- ascend_opp_path = os.environ["ASCEND_OPP_PATH"]
40	- if not os.path.exists(ascend_opp_path):
41	- raise Exception(f"ASCEND_OPP_PATH : {ascend_opp_path} does not exist. "
42	- "Please check whether the opp package has been installed. If exist, please run "
43	- "'source set_env.sh' in the CANN installation path." +
44	- pta_error(ErrCode.NOT_FOUND))
45	-
46	- ascend_runtime_path = os.path.join(ascend_home_path, "runtime")
47	- if not os.path.exists(ascend_runtime_path):
48	- raise Exception(f"ASCEND_RUNTIME_PATH : {ascend_runtime_path} does not exist. "
49	- "Please check whether the runtime package has been installed. If exist, please run "
50	- "'source set_env.sh' in the CANN installation path." +
51	- pta_error(ErrCode.NOT_FOUND))
52	-
53	- ascend_compiler_path = os.path.join(ascend_home_path, "compiler")
54	- if not os.path.exists(ascend_compiler_path):
55	- raise Exception(f"ASCEND_COMPILER_PATH : {ascend_compiler_path} does not exist. "
56	- "Please check whether the compiler package has been installed. If exist, please run "
57	- "'source set_env.sh' in the CANN installation path." +
58	- pta_error(ErrCode.NOT_FOUND))
59	-
60	- # get the cann version
61	- cann_version = get_cann_version()
62	-
63	- # check whether the CANN package version matches the pytorch version
64	- if cann_version in cann_pytorch_version_map and \
65	- torch_npu.__version__ not in cann_pytorch_version_map[cann_version]:
66	- print(f"Warning: CANN package version {cann_version} and PyTorch version {torch_npu.__version__} "
67	- "do not match. Please check the README of the Ascend PyTorch repo.")
68	- else:
69	- print("Warning: ASCEND_HOME_PATH environment variable is not set.")
70	-
71	-
72	-def _create_wrap_func(check_func):
73	- def decorator(func):
74	- @wraps(func)
75	- def wrapper(args, *kwargs):
76	- if check_func(args, *kwargs):
77	- raise RuntimeError(f"{str(func)} is not supported in npu." + pta_error(ErrCode.NOT_SUPPORT))
78	-
79	- return func(args, *kwargs)
80	- return wrapper
81	- return decorator
82	-
83	-
84	-# Specific check functions
85	-def _is_tensor_npu_supported(args, *kwargs):
86	- return torch.is_tensor(args[0]) and args[0].is_npu
87	-
88	-
89	-def _is_module_parameters_supported(args, *kwargs):
90	- module_args = [m for m in args if isinstance(m, torch.nn.Module) and hasattr(m, "_modules")]
91	- module_parameters = [p for _, p in module_args[0].named_parameters()]
92	- return any(p.device is not None and p.device.type == "npu" for p in module_parameters)
93	-
94	-
95	-def _apply_wrap_func_to_modules(wrap_func, unsupported_modules):
96	- for attr_name, parent_module in unsupported_modules.items():
97	- setattr(parent_module, attr_name, wrap_func(getattr(parent_module, attr_name)))
98	-
99	-
100	-# Apply wrap functions to specific modules
101	-def _add_intercept_methods():
102	- _apply_wrap_func_to_modules(_create_wrap_func(_is_tensor_npu_supported), unsupported_Tensor_api)
103	- _apply_wrap_func_to_modules(_create_wrap_func(_is_module_parameters_supported), unsupported_nn_api)

Dtorch_npu/utils/unsupport_api.py+0-18

	@@ -1,18 +0,0 @@
1	-import torch
2	-
3	-"""
4	-key: attr_name(str)
5	-value: parent_module(object)
6	-"""
7	-
8	-unsupported_Tensor_api = {
9	-}
10	-
11	-unsupported_nn_api = {
12	- "share_memory": torch.nn.Module,
13	- "add_module": torch.jit.ScriptModule,
14	- "bfloat16": torch.jit.ScriptModule,
15	- "register_buffer": torch.jit.ScriptModule,
16	- "register_parameter": torch.jit.ScriptModule,
17	- "register_module": torch.jit.ScriptModule
18	-}