refactor for torch_npu init module.-pytorch-AtomGit

Atest/test_torch_npu_init.py+459-0

@@ -0,0 +1,459 @@
	1	+# Owner(s): ["module: unknown"]
	2	+import json
	3	+import os
	4	+import statistics
	5	+import subprocess
	6	+import sys
	7	+import textwrap
	8	+
	9	+from torch_npu.testing.testcase import run_tests, TestCase
	10	+
	11	+
	12	+REQUIRED_C_EXTENSION_CHILDREN = [
	13	+ "_profiler",
	14	+ "_distributed_c10d",
	15	+ "_cd",
	16	+ "_logging",
	17	+ "_flops_count",
	18	+]
	19	+
	20	+EXPECTED_LOADED_MODULES = [
	21	+ "torch_npu.npu",
	22	+ "torch_npu.npu.amp",
	23	+ "torch_npu.npu.aclnn",
	24	+ "torch_npu.optim",
	25	+ "torch_npu.dynamo",
	26	+ "torch_npu._logging",
	27	+ "torch_npu._afd",
	28	+ "torch_npu.profiler",
	29	+ "torch_npu.distributed",
	30	+ "torch_npu.distributed.rpc",
	31	+ "torch_npu.op_plugin",
	32	+ "torch_npu.op_plugin.meta",
	33	+ "torch_npu.op_plugin.meta._meta_registrations",
	34	+ "torch_npu.utils._dynamo",
	35	+ "torch_npu.utils._inductor",
	36	+ "torch_npu.utils.custom_ops",
	37	+ "torch_npu.utils.patch_getenv",
	38	+]
	39	+
	40	+EXPECTED_NOT_LOADED_MODULES = [
	41	+ "torch_npu._C._afd",
	42	+]
	43	+
	44	+EXPECTED_TOP_LEVEL_ATTRS = [
	45	+ "npu",
	46	+ "optim",
	47	+ "dynamo",
	48	+ "_afd",
	49	+ "profiler",
	50	+ "op_plugin",
	51	+ "utils",
	52	+]
	53	+
	54	+LAZY_TOP_LEVEL_APIS = [
	55	+ "HiFloat8Tensor",
	56	+ "erase_stream",
	57	+ "matmul_checksum",
	58	+]
	59	+
	60	+AFD_OPS = [
	61	+ "attention_worker_scheduler_",
	62	+ "attention_worker_scheduler",
	63	+ "ffn_worker_scheduler_",
	64	+ "ffn_worker_scheduler",
	65	+]
	66	+
	67	+
	68	+class TestTorchNpuBootstrap(TestCase):
	69	+ def _run_python(self, code: str, *, optional: bool = False):
	70	+ proc = subprocess.run(
	71	+ [sys.executable, "-c", textwrap.dedent(code)],
	72	+ text=True,
	73	+ capture_output=True,
	74	+ env=os.environ.copy(),
	75	+ )
	76	+
	77	+ if proc.returncode == 0:
	78	+ return
	79	+
	80	+ message = (
	81	+ f"subprocess failed with return code {proc.returncode}\n"
	82	+ f"stdout:\n{proc.stdout}\n"
	83	+ f"stderr:\n{proc.stderr}"
	84	+ )
	85	+
	86	+ if optional:
	87	+ self.skipTest(message)
	88	+
	89	+ self.fail(message)
	90	+
	91	+ def test_01_import_order_compatibility(self):
	92	+ cases = [
	93	+ "import torch_npu",
	94	+ "import torch\nimport torch_npu",
	95	+ "import torch_npu\nimport torch",
	96	+ "import torch_npu\nimport torch_npu",
	97	+ ]
	98	+
	99	+ for code in cases:
	100	+ self._run_python(code)
	101	+
	102	+ def test_02_import_state_snapshot(self):
	103	+ self._run_python(
	104	+ f"""
	105	+ import sys
	106	+ import torch
	107	+ import torch_npu
	108	+ import torch_npu._C as C
	109	+
	110	+ required_c_children = {REQUIRED_C_EXTENSION_CHILDREN!r}
	111	+ expected_loaded = {EXPECTED_LOADED_MODULES!r}
	112	+ expected_not_loaded = {EXPECTED_NOT_LOADED_MODULES!r}
	113	+ expected_top_attrs = {EXPECTED_TOP_LEVEL_ATTRS!r}
	114	+
	115	+ assert hasattr(torch, "npu"), "torch.npu is not registered"
	116	+ assert torch.npu is torch_npu.npu
	117	+
	118	+ assert hasattr(torch.Tensor, "npu"), "torch.Tensor.npu is missing"
	119	+ assert hasattr(torch.nn.Module, "npu"), "torch.nn.Module.npu is missing"
	120	+
	121	+ for name in required_c_children:
	122	+ assert hasattr(C, name), f"torch_npu._C.{{name}} is missing"
	123	+
	124	+ # Old behavior: AFD is exposed as torch_npu._afd, not torch_npu._C._afd.
	125	+ assert hasattr(C, "_afd") is False
	126	+
	127	+ missing_modules = [
	128	+ name for name in expected_loaded if name not in sys.modules
	129	+ ]
	130	+ assert not missing_modules, (
	131	+ f"init-time modules changed, missing: {{missing_modules}}"
	132	+ )
	133	+
	134	+ unexpected_modules = [
	135	+ name for name in expected_not_loaded if name in sys.modules
	136	+ ]
	137	+ assert not unexpected_modules, (
	138	+ f"unexpected eager modules loaded: {{unexpected_modules}}"
	139	+ )
	140	+
	141	+ missing_attrs = [
	142	+ name for name in expected_top_attrs if not hasattr(torch_npu, name)
	143	+ ]
	144	+ assert not missing_attrs, (
	145	+ f"torch_npu top-level attrs changed, missing: {{missing_attrs}}"
	146	+ )
	147	+
	148	+ import torch_npu.npu
	149	+ import torch_npu.npu.aclnn
	150	+
	151	+ assert torch_npu.npu is not None
	152	+ assert torch_npu.npu.aclnn is not None
	153	+
	154	+ # _op_plugin_docs is imported for side effect, then removed from top-level.
	155	+ assert "torch_npu._op_plugin_docs" in sys.modules
	156	+ assert not hasattr(torch_npu, "_op_plugin_docs")
	157	+ """
	158	+ )
	159	+
	160	+ def test_03_public_exports_snapshot(self):
	161	+ self._run_python(
	162	+ f"""
	163	+ import torch
	164	+ import torch_npu
	165	+ import torch_npu._C as C
	166	+ from torch_npu.utils.exposed_api import public_npu_functions
	167	+
	168	+ lazy_names = {LAZY_TOP_LEVEL_APIS!r}
	169	+
	170	+ for name in lazy_names:
	171	+ assert name in torch_npu.__all__, f"{{name}} is missing from __all__"
	172	+ assert name in dir(torch_npu), f"{{name}} is missing from dir(torch_npu)"
	173	+ assert name not in torch_npu.__dict__, (
	174	+ f"{{name}} should not be cached before lazy access"
	175	+ )
	176	+
	177	+ value = getattr(torch_npu, name)
	178	+
	179	+ assert value is not None
	180	+ assert name in torch_npu.__dict__, (
	181	+ f"{{name}} was not cached after lazy access"
	182	+ )
	183	+
	184	+ available_public_ops = []
	185	+ missing_from_torch_npu = []
	186	+ missing_from_all = []
	187	+ missing_torch_alias = []
	188	+
	189	+ for name in public_npu_functions:
	190	+ if not hasattr(torch.ops.npu, name):
	191	+ continue
	192	+
	193	+ available_public_ops.append(name)
	194	+
	195	+ if not hasattr(torch_npu, name):
	196	+ missing_from_torch_npu.append(name)
	197	+
	198	+ if name not in torch_npu.__all__:
	199	+ missing_from_all.append(name)
	200	+
	201	+ if not hasattr(torch, name):
	202	+ missing_torch_alias.append(name)
	203	+
	204	+ assert available_public_ops, "no available public torch.ops.npu ops found"
	205	+ assert not missing_from_torch_npu, (
	206	+ f"some public ops are missing from torch_npu: "
	207	+ f"{{missing_from_torch_npu[:20]}}"
	208	+ )
	209	+ assert not missing_from_all, (
	210	+ f"some public ops are missing from torch_npu.__all__: "
	211	+ f"{{missing_from_all[:20]}}"
	212	+ )
	213	+ assert not missing_torch_alias, (
	214	+ f"some public ops are missing deprecated torch aliases: "
	215	+ f"{{missing_torch_alias[:20]}}"
	216	+ )
	217	+
	218	+ dtype_names = [
	219	+ name
	220	+ for name in dir(C._cd.DType)
	221	+ if not name.startswith("_") and name not in ["_dir", "name"]
	222	+ ]
	223	+
	224	+ missing_dtype = []
	225	+ mismatch_dtype = []
	226	+
	227	+ for name in dtype_names:
	228	+ if not hasattr(torch_npu, name):
	229	+ missing_dtype.append(name)
	230	+ continue
	231	+
	232	+ exported = getattr(torch_npu, name)
	233	+ source = getattr(C._cd.DType, name)
	234	+
	235	+ # Pybind objects may not preserve Python identity across getattr calls.
	236	+ if exported != source and repr(exported) != repr(source):
	237	+ mismatch_dtype.append((name, repr(exported), repr(source)))
	238	+
	239	+ assert not missing_dtype, (
	240	+ f"some DType symbols are missing from torch_npu: {{missing_dtype}}"
	241	+ )
	242	+ assert not mismatch_dtype, (
	243	+ f"some DType symbols do not match torch_npu._C._cd.DType: "
	244	+ f"{{mismatch_dtype[:10]}}"
	245	+ )
	246	+ """
	247	+ )
	248	+
	249	+ def test_04_framework_registration_snapshot(self):
	250	+ self._run_python(
	251	+ """
	252	+ import torch_npu
	253	+ import torch.distributed as dist
	254	+ import torch.distributed.rpc as rpc
	255	+ import torch.distributed.tensor # noqa: F401
	256	+ from torch._dynamo.device_interface import get_interface_for_device
	257	+ from torch._dynamo.backends.registry import _BACKENDS
	258	+ from torch._inductor.codegen.common import device_op_overrides_dict
	259	+
	260	+ iface = get_interface_for_device("npu")
	261	+ assert iface is not None
	262	+
	263	+ assert "npu" in _BACKENDS, "npu dynamo backend is not registered"
	264	+ assert "npugraph_ex" in _BACKENDS, (
	265	+ "npugraph_ex dynamo backend is not registered"
	266	+ )
	267	+
	268	+ assert "npu" in device_op_overrides_dict
	269	+ assert device_op_overrides_dict.get("npu") is not None
	270	+
	271	+ assert "hccl" in dist.Backend.backend_list
	272	+ assert "lccl" in dist.Backend.backend_list
	273	+
	274	+ names = [
	275	+ name for name in dir(rpc.BackendType)
	276	+ if "NPU" in name or "TENSORPIPE" in name
	277	+ ]
	278	+ assert hasattr(rpc, "BackendType")
	279	+ assert "NPU_TENSORPIPE" in names
	280	+ """,
	281	+ optional=True,
	282	+ )
	283	+
	284	+ def test_05_runtime_lazy_init_semantics(self):
	285	+ self._run_python(
	286	+ """
	287	+ import torch
	288	+ import torch_npu
	289	+
	290	+ assert torch_npu.npu.is_initialized() is False, (
	291	+ "import torch_npu unexpectedly triggered NPU lazy init"
	292	+ )
	293	+
	294	+ torch.npu.is_available()
	295	+ torch.npu.device_count()
	296	+
	297	+ assert torch_npu.npu.is_initialized() is False, (
	298	+ "availability query unexpectedly triggered NPU lazy init"
	299	+ )
	300	+ """
	301	+ )
	302	+
	303	+ self._run_python(
	304	+ """
	305	+ import sys
	306	+ import torch_npu
	307	+
	308	+ assert torch_npu.npu.is_initialized() is False
	309	+
	310	+ if torch_npu.npu.device_count() <= 0:
	311	+ sys.exit(0)
	312	+
	313	+ torch_npu.npu.get_device_properties(0)
	314	+
	315	+ assert torch_npu.npu.is_initialized() is True, (
	316	+ "runtime NPU API did not trigger lazy init"
	317	+ )
	318	+ """
	319	+ )
	320	+
	321	+ self._run_python(
	322	+ """
	323	+ import torch_npu
	324	+
	325	+ assert torch_npu.npu.is_initialized() is False
	326	+
	327	+ torch_npu.npu.init()
	328	+
	329	+ assert torch_npu.npu.is_initialized() is True
	330	+ """
	331	+ )
	332	+
	333	+ def test_06_component_behavior_snapshot(self):
	334	+ self._run_python(
	335	+ f"""
	336	+ import os
	337	+ import sys
	338	+ import torch_npu
	339	+ import torch_npu._C as C
	340	+ import torch_npu._afd
	341	+ import torch_npu.utils as utils
	342	+ import torch_npu.utils.asd_detector as asd_detector
	343	+ import torch_npu.utils.patch_getenv as patch_getenv
	344	+
	345	+ afd_ops = {AFD_OPS!r}
	346	+
	347	+ # patch_getenv behavior.
	348	+ assert os.getenv is patch_getenv._patched_getenv
	349	+ assert os.environ.get is patch_getenv._patched_environ_get
	350	+
	351	+ # ASD compatibility APIs.
	352	+ for module_name in [
	353	+ "torch_npu.utils._asd_detector",
	354	+ "torch_npu.utils.asd_detector",
	355	+ ]:
	356	+ assert module_name in sys.modules, (
	357	+ f"{{module_name}} is not loaded after import torch_npu"
	358	+ )
	359	+
	360	+ for api_name in ["set_asd_loss_scale", "register_asd_hook"]:
	361	+ assert hasattr(utils, api_name), (
	362	+ f"torch_npu.utils.{{api_name}} is missing"
	363	+ )
	364	+ assert hasattr(asd_detector, api_name), (
	365	+ f"torch_npu.utils.asd_detector.{{api_name}} is missing"
	366	+ )
	367	+
	368	+ utils_api = getattr(utils, api_name)
	369	+ detector_api = getattr(asd_detector, api_name)
	370	+
	371	+ assert callable(utils_api)
	372	+ assert callable(detector_api)
	373	+ assert utils_api is detector_api
	374	+
	375	+ # AFD compatibility behavior.
	376	+ assert hasattr(C, "_afd") is False
	377	+ assert "torch_npu._afd" in sys.modules
	378	+ assert "torch_npu._C._afd" not in sys.modules
	379	+
	380	+ try:
	381	+ import torch_npu._C._afd # noqa: F401
	382	+ raise AssertionError("import torch_npu._C._afd should fail")
	383	+ except ModuleNotFoundError:
	384	+ pass
	385	+
	386	+ for name in afd_ops:
	387	+ assert hasattr(torch_npu._afd, name), (
	388	+ f"torch_npu._afd.{{name}} is missing"
	389	+ )
	390	+ """
	391	+ )
	392	+
	393	+ def test_07_distributed_patch_behavior(self):
	394	+ self._run_python(
	395	+ """
	396	+ import sys
	397	+ import torch
	398	+ import torch_npu
	399	+ import torch.distributed as dist
	400	+ import torch.distributed.distributed_c10d as c10d
	401	+ import torch.distributed.launcher.api as launcher_api
	402	+ from torch.distributed.fsdp import sharded_grad_scaler
	403	+ from torch.distributed.fsdp._fully_shard import _fsdp_collectives
	404	+ from torch.distributed.fsdp._fully_shard._fsdp_param_group import (
	405	+ FSDPParamGroup,
	406	+ )
	407	+ from torch_npu.distributed.fsdp._add_fsdp_patch import (
	408	+ _patched_finalize_backward,
	409	+ _patched_get_param_all_gather_inputs,
	410	+ _patched_all_gather_copy_in,
	411	+ )
	412	+ from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler
	413	+
	414	+ assert torch._C._distributed_c10d._verify_params_across_processes is (
	415	+ torch_npu.distributed._verify_params_across_processes
	416	+ )
	417	+
	418	+ assert torch._C._distributed_c10d.ProcessGroup._get_sequence_number_for_group is (
	419	+ torch_npu.distributed.distributed_c10d._hccl_get_sequence_number_for_group
	420	+ )
	421	+
	422	+ assert c10d._add_ephemeral_timeout_for_all_pgs is (
	423	+ torch_npu.distributed.distributed_c10d._hccl_add_ephemeral_timeout_for_all_pgs
	424	+ )
	425	+
	426	+ assert dist.batch_isend_irecv is (
	427	+ torch_npu.distributed.distributed_c10d._batch_isend_irecv
	428	+ )
	429	+ assert c10d.batch_isend_irecv is (
	430	+ torch_npu.distributed.distributed_c10d._batch_isend_irecv
	431	+ )
	432	+
	433	+ assert dist.gather is torch_npu.distributed.distributed_c10d._gather
	434	+ assert c10d.gather is torch_npu.distributed.distributed_c10d._gather
	435	+
	436	+ assert dist.gather_object is torch_npu.distributed.distributed_c10d._gather_object
	437	+ assert c10d.gather_object is torch_npu.distributed.distributed_c10d._gather_object
	438	+
	439	+ assert dist.is_hccl_available is torch_npu.distributed.is_hccl_available
	440	+ assert dist.reinit_process_group is torch_npu.distributed.reinit_process_group
	441	+
	442	+ assert callable(c10d.rendezvous)
	443	+ assert callable(launcher_api._get_addr_and_port)
	444	+
	445	+ assert sharded_grad_scaler.ShardedGradScaler is _ShardedGradScaler
	446	+ assert FSDPParamGroup.finalize_backward is _patched_finalize_backward
	447	+ assert _fsdp_collectives._get_param_all_gather_inputs is (
	448	+ _patched_get_param_all_gather_inputs
	449	+ )
	450	+ assert torch.ops.fsdp.all_gather_copy_in is _patched_all_gather_copy_in
	451	+ assert torch.ops.fsdp.all_gather_copy_in.default is (
	452	+ _patched_all_gather_copy_in
	453	+ )
	454	+ """
	455	+ )
	456	+
	457	+
	458	+if __name__ == "__main__":
	459	+ run_tests()

Mtorch_npu/__init__.py+44-398

	@@ -1,13 +1,8 @@
1	-__all__ = ["erase_stream", "matmul_checksum", "HiFloat8Tensor"]	1	+__all__ = ["HiFloat8Tensor", "erase_stream", "matmul_checksum"]
2		2
3	import atexit	3	import atexit
4	import ctypes	4	import ctypes
5	import os	5	import os
6	-import sys
7	-import traceback
8	-import types
9	-import warnings
10	-from functools import wraps
11		6
12		7
13	# Disable autoloading before running 'import torch' to avoid circular dependencies	8	# Disable autoloading before running 'import torch' to avoid circular dependencies
	@@ -15,408 +10,59 @@ ORG_AUTOLOAD = os.getenv("TORCH_DEVICE_BACKEND_AUTOLOAD", "1")
15	os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"	10	os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
16		11
17	import torch	12	import torch
18	-import torch_npu
19	-from torch.distributed.fsdp import sharded_grad_scaler
20	-from torch.utils.checkpoint import DefaultDeviceType
21		13
22	-	14	+# Import-time env access logging patch. Keep early to capture initialization-time getenv.
23	-acc = torch._C._get_accelerator()	15	+import torch_npu.utils.patch_getenv
24	-if acc.type != "cpu":	16	+from torch_npu._init.core.module_loader import _load_core_modules
25	- import time	17	+from torch_npu._init.core.optional_features import _enable_optional_features
26	-	18	+from torch_npu._init.core.runtime_lifecycle import _initialize_runtime_lifecycle
27	- # torch_npu.utils._error_code.ErrCode.NOT_SUPPORT	19	+from torch_npu._init.patches.patch_manager import _apply_patches
28	- error_code = "ERR00007"	20	+from torch_npu._init.registry.registry_manager import _register_components
29	- error_code_msg = "feature not supported"
30	- submodule_name = "PTA"
31	- raise RuntimeError(
32	- f"Two accelerators cannot be used at the same time "
33	- f"in PyTorch: npu and {acc.type}. You can install "
34	- f"the cpu version of PyTorch to use your npu device, "
35	- f"or use the {acc.type} device with "
36	- f"'export TORCH_DEVICE_BACKEND_AUTOLOAD=0'.\n"
37	- f"[ERROR] {time.strftime('%Y-%m-%d-%H:%M:%S', time.localtime())} "
38	- f"(PID:{os.getpid()}, Device:-1, RankID:-1) "
39	- f"{error_code} {submodule_name} {error_code_msg}"
40	- )
41	-
42	-try:
43	- import torch_npu.npu
44	-except ImportError as e:
45	- if "libhccl.so" in str(e):
46	- if "ASCEND_OPP_PATH" in os.environ:
47	- # Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!!
48	- e.msg += (
49	- ". Please check that the compiler package is installed. "
50	- "Please run 'source set_env.sh' in the CANN installation path."
51	- )
52	- else:
53	- # Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!!
54	- e.msg += (
55	- ". Please check that the cann package is installed. "
56	- "Please run 'source set_env.sh' in the CANN installation path."
57	- )
58	- elif "libascendcl.so" in str(e):
59	- # Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!!
60	- e.msg += (
61	- ". Please check that the runtime package is installed. "
62	- "Please run 'source set_env.sh' in the CANN installation path."
63	- )
64	- raise
65	-
66	-import torch_npu._afd
67	-import torch_npu._C
68	-import torch_npu._logging
69	-import torch_npu.distributed.rpc
70	-import torch_npu.dynamo
71	-import torch_npu.npu.aclnn
72	-import torch_npu.npu.amp
73	-import torch_npu.op_plugin
74	-import torch_npu.optim
75	-import torch_npu.utils._afd_ops
76	-import torch_npu.utils.custom_ops
77	-from torch_npu import _op_plugin_docs, profiler
78	-from torch_npu._C._distributed_c10d import ParallelStore
79	-from torch_npu.asd.asd import _asd_patch
80	-from torch_npu.asd.checksum import _matmul_checksum as matmul_checksum
81	-from torch_npu.contrib.function import npu_functional
82	-from torch_npu.contrib.module import npu_modules
83	-from torch_npu.distributed.fsdp._add_fsdp_patch import _apply_fsdp_patch
84	-from torch_npu.distributed.rpc.backend_registry import _rpc_backend_registry
85	-from torch_npu.dynamo import _patch_npu_trace_rules
86	-from torch_npu.multiprocessing.reductions import _add_reductions_methods
87	-from torch_npu.npu._format import _apply_npu_format_patch
88	-from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler
89	-from torch_npu.npu.utils import _erase_stream as erase_stream
90	-from torch_npu.op_plugin.meta import _meta_registrations
91	-from torch_npu.profiler._add_mstx_patch import _apply_mstx_patch
92	-from torch_npu.utils import (
93	- _add_collect_env_methods,
94	- _add_intercept_methods,
95	- _add_serialization_methods,
96	- _add_storage_methods,
97	- _add_tensor_methods,
98	- _apply_dlpack_patch,
99	- _apply_module_patch,
100	- _apply_npu_show_warning,
101	- _apply_npugraph_tree_methods,
102	- _cann_package_check,
103	- _inductor_register_device_op_overrides,
104	- add_dynamo_methods,
105	- add_optim_method,
106	- add_perf_dump_patch,
107	- npu_patch_meta,
108	- patch_getenv,
109	-)
110	-from torch_npu.utils._dynamo_device import _dynamo_register_interface_for_device
111	-from torch_npu.utils._error_code import _except_handler, ErrCode, pta_error
112	-from torch_npu.utils.exposed_api import public_npu_functions
113	-from torch_npu.utils.hif8_tensor import _HiFloat8Tensor as HiFloat8Tensor
114	-from torch_npu.utils.utils import _is_interactive_command_line
115	from torch_npu.version import __version__ as __version__	21	from torch_npu.version import __version__ as __version__
116		22
117		23
118	-del _op_plugin_docs	24	+def _check_device_conflict():
		25	+ acc = torch._C._get_accelerator()
		26	+ if acc.type not in ("cpu", "npu"):
		27	+ import time
119		28
120	-_cann_package_check()	29	+ # torch_npu.utils._error_code.ErrCode.NOT_SUPPORT
121	-	30	+ error_code = "ERR00007"
122	-	31	+ error_code_msg = "feature not supported"
123	-def _wrap_torch_error_func(func):	32	+ submodule_name = "PTA"
124	- @wraps(func)
125	- def wrapper(args, *kwargs):
126	raise RuntimeError(	33	raise RuntimeError(
127	- f"torch.{func.__name__} is deprecated and will be removed in future version. "	34	+ f"Two accelerators cannot be used at the same time "
128	- f"Use torch_npu.{func.__name__} instead." + pta_error(ErrCode.NOT_SUPPORT)	35	+ f"in PyTorch: npu and {acc.type}. You can install "
		36	+ f"the cpu version of PyTorch to use your npu device, "
		37	+ f"or use the {acc.type} device with "
		38	+ f"'export TORCH_DEVICE_BACKEND_AUTOLOAD=0'.\n"
		39	+ f"[ERROR] {time.strftime('%Y-%m-%d-%H:%M:%S', time.localtime())} "
		40	+ f"(PID:{os.getpid()}, Device:-1, RankID:-1) "
		41	+ f"{error_code} {submodule_name} {error_code_msg}"
129	)	42	)
130		43
131	- return wrapper	44	+
		45	+def _initialize():
		46	+ # 1. pre-init checks
		47	+ _check_device_conflict()
		48	+
		49	+ # 2. core modules, registration side effects and public API export
		50	+ _load_core_modules()
		51	+
		52	+ # 3. backend and framework integration registration
		53	+ _register_components()
		54	+
		55	+ # 4. apply patches
		56	+ _apply_patches()
		57	+
		58	+ # 5. optional runtime features
		59	+ _enable_optional_features()
		60	+
		61	+ # 6. final extension barrier and shutdown hook
		62	+ _initialize_runtime_lifecycle()
132		63
133		64
134	-for name in dir(torch.ops.npu):	65	+_initialize()
135	- if name.startswith("__") or name in ["_dir", "name"]:
136	- continue
137	- globals()[name] = getattr(torch.ops.npu, name)
138	- if name in public_npu_functions:
139	- __all__.append(name)
140	- setattr(torch, name, _wrap_torch_error_func(getattr(torch.ops.npu, name)))
141	-
142	-for name in dir(torch_npu._C._cd.DType):
143	- if name.startswith("__") or name in ["_dir", "name"]:
144	- continue
145	- setattr(torch_npu, name, getattr(torch_npu._C._cd.DType, name))
146	-
147	-all_monkey_patches = [
148	- ["nn.functional", npu_functional],
149	- ["nn", npu_modules],
150	-]
151	-
152	-
153	-def _apply_patches(monkey_patches):
154	- def _getattr(module_list, root_module=torch):
155	- if len(module_list) <= 1:
156	- return root_module
157	-
158	- if hasattr(root_module, module_list[0]):
159	- return _getattr(module_list[1:], getattr(root_module, module_list[0]))
160	- else:
161	- empty_module_name = f"{root_module.__name__}.{module_list[0]}"
162	- sys.modules[empty_module_name] = types.ModuleType(empty_module_name)
163	- setattr(root_module, module_list[0], sys.modules.get(empty_module_name))
164	- return _getattr(module_list[1:], getattr(root_module, module_list[0]))
165	-
166	- for patch_pair in monkey_patches:
167	- dest, patch = patch_pair
168	- dest_module = _getattr(dest.split("."), root_module=torch)
169	- last_module_level = dest.split(".")[-1]
170	- if not isinstance(patch, types.ModuleType):
171	- setattr(dest_module, last_module_level, patch)
172	- continue
173	-
174	- if not hasattr(dest_module, last_module_level) or not hasattr(patch, "__all__"):
175	- setattr(dest_module, last_module_level, patch)
176	- sys.modules[f"{dest_module.__name__}.{last_module_level}"] = patch
177	- continue
178	-
179	- if not hasattr(patch, "__all__"):
180	- raise NotImplementedError(
181	- "Patch module must have __all__ definition."
182	- + pta_error(ErrCode.NOT_SUPPORT)
183	- )
184	- dest_module = getattr(dest_module, last_module_level)
185	- for attr in patch.__all__:
186	- setattr(dest_module, attr, getattr(patch, attr))
187	-
188	-
189	-def _apply_sharded_grad_scaler_patch():
190	- torch.distributed.fsdp.sharded_grad_scaler.ShardedGradScaler = _ShardedGradScaler
191	-
192	-
193	-def _apply_class_patches():
194	- _apply_npu_show_warning()
195	- _add_storage_methods()
196	- _apply_dlpack_patch()
197	- _apply_module_patch()
198	- _add_tensor_methods()
199	- _add_serialization_methods()
200	- _add_intercept_methods()
201	- _add_collect_env_methods()
202	- add_dynamo_methods()
203	- add_optim_method()
204	- _apply_sharded_grad_scaler_patch()
205	- add_perf_dump_patch()
206	- _apply_distributed_methods_patch()
207	- _apply_mstx_patch()
208	- _add_reductions_methods()
209	- _apply_npu_format_patch()
210	- _apply_fsdp_patch()
211	- _apply_npugraph_tree_methods()
212	- npu_patch_meta()
213	-
214	-
215	-def _apply_distributed_methods_patch():
216	- torch._C._distributed_c10d._verify_params_across_processes = (
217	- torch_npu.distributed._verify_params_across_processes
218	- )
219	- torch.distributed.batch_isend_irecv = (
220	- torch_npu.distributed.distributed_c10d._batch_isend_irecv
221	- )
222	- torch.distributed.distributed_c10d.batch_isend_irecv = (
223	- torch_npu.distributed.distributed_c10d._batch_isend_irecv
224	- )
225	- torch.distributed.gather = torch_npu.distributed.distributed_c10d._gather
226	- torch.distributed.distributed_c10d.gather = (
227	- torch_npu.distributed.distributed_c10d._gather
228	- )
229	- torch.distributed.gather_object = (
230	- torch_npu.distributed.distributed_c10d._gather_object
231	- )
232	- torch.distributed.distributed_c10d.gather_object = (
233	- torch_npu.distributed.distributed_c10d._gather_object
234	- )
235	- torch.distributed.is_hccl_available = torch_npu.distributed.is_hccl_available
236	- torch.distributed.reinit_process_group = torch_npu.distributed.reinit_process_group
237	- torch.distributed.distributed_c10d.rendezvous = (
238	- torch_npu.distributed.distributed_c10d._trigger_rendezvous_decorator(
239	- torch.distributed.distributed_c10d.rendezvous
240	- )
241	- )
242	- torch.distributed.launcher.api._get_addr_and_port = (
243	- torch_npu.distributed.distributed_c10d._trigger__get_addr_and_port_decorator(
244	- torch.distributed.launcher.api._get_addr_and_port
245	- )
246	- )
247	- torch._C._distributed_c10d.ProcessGroup._get_sequence_number_for_group = (
248	- torch_npu.distributed.distributed_c10d._hccl_get_sequence_number_for_group
249	- )
250	- torch.distributed.nn.functional._AllGatherBase.backward = (
251	- torch_npu.distributed.nn.functional._allgather_base_backward_hccl
252	- )
253	- torch.distributed.distributed_c10d._add_ephemeral_timeout_for_all_pgs = (
254	- torch_npu.distributed.distributed_c10d._hccl_add_ephemeral_timeout_for_all_pgs
255	- )
256	-
257	-
258	-torch.utils.rename_privateuse1_backend("npu")
259	-# rename device name to 'npu' and register funcs
260	-torch._register_device_module("npu", torch_npu.npu)
261	-unsupported_dtype = [
262	- torch.quint8,
263	- torch.quint4x2,
264	- torch.quint2x4,
265	- torch.qint32,
266	- torch.qint8,
267	-]
268	-torch.utils.generate_methods_for_privateuse1_backend(
269	- for_tensor=True,
270	- for_module=True,
271	- for_storage=True,
272	- unsupported_dtype=unsupported_dtype,
273	-)
274	-torch.nn.parameter.UninitializedTensorMixin._allowed_methods.append(torch.Tensor.npu)
275	-
276	-# register npu device interface for dynamo
277	-_dynamo_register_interface_for_device()
278	-
279	-# Apply monkey-patches.
280	-_apply_patches(all_monkey_patches)
281	-_apply_class_patches()
282	-_asd_patch()
283	-_except_handler.patch_excepthook()
284	-
285	-_warn_msg = {
286	- "DropoutWithByteMask": (
287	- "torch.nn.DropoutWithByteMask is deprecated and will be removed in future version. "
288	- "Use torch_npu.contrib.module.DropoutWithByteMask instead."
289	- ),
290	- "dropout_with_byte_mask": (
291	- "torch.nn.functional.dropout_with_byte_mask is deprecated and will be removed in future version. "
292	- "Use torch_npu.contrib.function.dropout_with_byte_mask instead."
293	- ),
294	-}
295	-
296	-
297	-def _wrap_torch_patch_warning_func(func):
298	- @wraps(func)
299	- def wrapper(args, *kwargs):
300	- warnings.warn(_warn_msg[func.__name__])
301	- return func(args, *kwargs)
302	-
303	- return wrapper
304	-
305	-
306	-torch.nn.DropoutWithByteMask = _wrap_torch_patch_warning_func(
307	- torch.nn.DropoutWithByteMask
308	-)
309	-torch.nn.functional.dropout_with_byte_mask = _wrap_torch_patch_warning_func(
310	- torch.nn.functional.dropout_with_byte_mask
311	-)
312	-# this must be placed at the end
313	-torch_npu._C._initExtension()
314	-
315	-
316	-def _new_process_group_hccl_helper(dist_backend_opts, pg_options):
317	- store = dist_backend_opts.store
318	- group_rank = dist_backend_opts.group_rank
319	- group_size = dist_backend_opts.group_size
320	- if pg_options is None or not isinstance(
321	- pg_options, torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options
322	- ):
323	- pg_options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options()
324	- pg_options.is_high_priority_stream = False
325	- pg_options._timeout = dist_backend_opts.timeout
326	- pg_options.global_ranks_in_group = dist_backend_opts.global_ranks_in_group
327	- pg_options.group_id = dist_backend_opts.group_id
328	- return torch_npu._C._distributed_c10d.ProcessGroupHCCL(
329	- store, group_rank, group_size, pg_options
330	- )
331	-
332	-
333	-def _new_process_group_lccl_helper(dist_backend_opts, pg_options):
334	- store = dist_backend_opts.store
335	- group_rank = dist_backend_opts.group_rank
336	- group_size = dist_backend_opts.group_size
337	- return torch_npu._C._distributed_c10d.ProcessGroupLCCL(
338	- store, group_rank, group_size
339	- )
340	-
341	-
342	-def _register_distributed_backend_for_npu():
343	- # init and register hccl backend
344	- # Note: Since torch 2.8, the hccl backend must be registered at first to keep a right default_device_backend_map
345	- torch.distributed.Backend.register_backend(
346	- "hccl",
347	- lambda dist_backend_opts, pg_options: _new_process_group_hccl_helper(
348	- dist_backend_opts, pg_options
349	- ),
350	- extended_api=True,
351	- devices=["npu"],
352	- )
353	-
354	- # init and register lccl backend
355	- torch.distributed.Backend.register_backend(
356	- "lccl",
357	- lambda dist_backend_opts, pg_options: _new_process_group_lccl_helper(
358	- dist_backend_opts, pg_options
359	- ),
360	- extended_api=True,
361	- devices=["npu"],
362	- )
363	-
364	-
365	-# init and register distributed backend
366	-_register_distributed_backend_for_npu()
367	-
368	-
369	-# set default device type for gradient checkpointing
370	-DefaultDeviceType.set_device_type("npu")
371	-del DefaultDeviceType
372	-
373	-
374	-# NPU exit, need to synchronize devices
375	-def _npu_shutdown():
376	- success = torch_npu._C._npu_shutdown_synchronize()
377	- torch_npu.distributed.distributed_c10d._destructor_process_group()
378	- torch_npu._C._npu_shutdown(success)
379	- _except_handler.handle_exception()
380	- torch_npu.asd.asd.matmul_check._cleanup()
381	- if torch_npu.npu.aclnn._use_static_aclnn_kernel:
382	- from torch_npu._inductor.npu_static_kernel import uninstall_static_kernel
383	-
384	- uninstall_static_kernel()
385	-
386	-
387	-# register npu shutdown hook on exit
388	-atexit.register(_npu_shutdown)
389	-
390	-# init and register rpc npu backend
391	-_rpc_backend_registry()
392	-
393	-# Enable NPU Sanitizer
394	-if "TORCH_NPU_SANITIZER" in os.environ:
395	- import torch_npu.npu._sanitizer as csan
396	-
397	- csan.enable_npu_sanitizer()
398	-
399	-# register npu device op overrides for inductor
400	-_inductor_register_device_op_overrides()
401	-
402	-# Support stream into Dynamo charts
403	-_patch_npu_trace_rules()
404	-
405	-if _is_interactive_command_line():
406	- os.environ["TASK_QUEUE_ENABLE"] = "0"
407	- warnings.warn(
408	- "On the interactive interface, the value of TASK_QUEUE_ENABLE is set to 0 by default. \
409	- Do not set it to 1 to prevent some unknown errors"
410	- )
411	-
412	-# Enable transfer_to_npu via environment variable
413	-_transfer_to_npu_env = os.getenv("TORCH_TRANSFER_TO_NPU", "0")
414	-if _transfer_to_npu_env == "1":
415	- from torch_npu.contrib import transfer_to_npu
416	-elif _transfer_to_npu_env != "0":
417	- raise ValueError(
418	- f"Invalid value for TORCH_TRANSFER_TO_NPU: {_transfer_to_npu_env}. Only '0' or '1' is supported."
419	- )
420		66
421		67
422	# This function is an entrypoint called by PyTorch	68	# This function is an entrypoint called by PyTorch

Mtorch_npu/_afd/__init__.py+1-0

	@@ -1,5 +1,6 @@
1	__all__ = ["create_schedule_context_holder"]	1	__all__ = ["create_schedule_context_holder"]
2		2
		3	+
3	from ._schedule_context import _create_schedule_context_holder, ScheduleContextHolder	4	from ._schedule_context import _create_schedule_context_holder, ScheduleContextHolder
4		5
5		6

Atorch_npu/_init/__init__.py+2-0

@@ -0,0 +1,2 @@
	1	+# Internal implementation package for torch_npu.
	2	+# Do not import runtime modules here to avoid import-time side effects.

Atorch_npu/_init/common/warning_utils.py+10-0

@@ -0,0 +1,10 @@
	1	+import os
	2	+
	3	+
	4	+def _should_print_warning():
	5	+ disabled_warning = os.environ.get("TORCH_NPU_DISABLED_WARNING", "0")
	6	+ if disabled_warning == "1":
	7	+ return False
	8	+
	9	+ rank = os.environ.get("RANK", None)
	10	+ return rank is None or rank == "0"

Atorch_npu/_init/core/__init__.py+2-0

@@ -0,0 +1,2 @@
	1	+# Core initialization helpers for torch_npu.
	2	+# Keep this package initializer side-effect free.

Atorch_npu/_init/core/_exports.py+122-0

@@ -0,0 +1,122 @@
	1	+from functools import wraps
	2	+from importlib import import_module
	3	+
	4	+import torch
	5	+
	6	+import torch_npu
	7	+
	8	+
	9	+_LAZY_PYTHON_SYMBOLS = {
	10	+ "HiFloat8Tensor": ("torch_npu.utils.hif8_tensor", "_HiFloat8Tensor"),
	11	+ "erase_stream": ("torch_npu.npu.utils", "_erase_stream"),
	12	+ "matmul_checksum": ("torch_npu.asd.checksum", "_matmul_checksum"),
	13	+}
	14	+
	15	+
	16	+def _append_unique(all_list, names):
	17	+ for name in names:
	18	+ all_list.append(name)
	19	+
	20	+
	21	+def _export_npu_ops(globals_dict, all_list):
	22	+ """
	23	+ Export NPU custom ops from torch.ops.npu.
	24	+ Rules:
	25	+ - torch.ops.npu.<op_name> -> torch_npu.<op_name>
	26	+ - torch.ops.npu.<op_name> -> torch.<op_name> deprecated wrapper
	27	+ """
	28	+
	29	+ def _wrap_torch_error_func(func):
	30	+ @wraps(func)
	31	+ def wrapper(args, *kwargs):
	32	+ # lazy import to avoid early utils import / circular dependency.
	33	+ from torch_npu.utils._error_code import ErrCode, pta_error
	34	+
	35	+ raise RuntimeError(
	36	+ f"torch.{func.__name__} is deprecated and will be removed in future version. "
	37	+ f"Use torch_npu.{func.__name__} instead."
	38	+ + pta_error(ErrCode.NOT_SUPPORT)
	39	+ )
	40	+
	41	+ return wrapper
	42	+
	43	+ from torch_npu.utils.exposed_api import public_npu_functions
	44	+
	45	+ for name in dir(torch.ops.npu):
	46	+ if name.startswith("__") or name in ["_dir", "name"]:
	47	+ continue
	48	+ globals_dict[name] = getattr(torch.ops.npu, name)
	49	+ if name in public_npu_functions and name not in all_list:
	50	+ _append_unique(all_list, [name])
	51	+ setattr(torch, name, _wrap_torch_error_func(getattr(torch.ops.npu, name)))
	52	+
	53	+
	54	+def _export_dtype_symbols():
	55	+ """
	56	+ Export DType symbols to from C extension.
	57	+ Rule:
	58	+ - torch_npu._C._cd.DType.<dtype_name> -> torch_npu.<dtype_name>
	59	+ """
	60	+ for name in dir(torch_npu._C._cd.DType):
	61	+ if name.startswith("__") or name in ["_dir", "name"]:
	62	+ continue
	63	+ setattr(torch_npu, name, getattr(torch_npu._C._cd.DType, name))
	64	+
	65	+
	66	+def _export_lazy_python_apis(globals_dict, all_list):
	67	+ """
	68	+ Lazily export Python-defined top-levels APIs.
	69	+ Excample:
	70	+ torch_npu.utils.hif8_tensor._HiFloat8Tensor -> torch_npu.HiFloat8Tensor
	71	+ """
	72	+ module_name = globals_dict.get("__name__", "torch_npu")
	73	+
	74	+ def _lazy_import_api(name: str):
	75	+ try:
	76	+ import_path, attr_name = _LAZY_PYTHON_SYMBOLS[name]
	77	+ except KeyError as exc:
	78	+ raise AttributeError(
	79	+ f"module {module_name!r} has no attribute {name!r}"
	80	+ ) from exc
	81	+
	82	+ module = import_module(import_path)
	83	+ value = getattr(module, attr_name)
	84	+ globals_dict[name] = value
	85	+ return value
	86	+
	87	+ def __getattr__(name: str):
	88	+ return _lazy_import_api(name)
	89	+
	90	+ def __dir__():
	91	+ return sorted(set(globals_dict) \| set(_LAZY_PYTHON_SYMBOLS))
	92	+
	93	+ globals_dict["__getattr__"] = __getattr__
	94	+ globals_dict["__dir__"] = __dir__
	95	+ _append_unique(all_list, _LAZY_PYTHON_SYMBOLS.keys())
	96	+
	97	+
	98	+def _export_public_apis():
	99	+ """
	100	+ Export torch_npu public APIs.
	101	+ 1. python APIs:
	102	+ - torch_npu.utils.hif8_tensor._HiFloat8Tensor -> torch_npu.HiFloat8Tensor
	103	+ - torch_npu.npu.utils._erase_stream -> torch_npu.erase_stream
	104	+ - torch_npu.asd.checksum._matmul_checksum -> torch_npu.matmul_checksum
	105	+
	106	+ 2. NPU custom ops:
	107	+ - torch.ops.npu.<op_name> -> torch_npu.<op_name>
	108	+ - torch.ops.npu.<op_name> -> torch.<op_name> deprecated wrapper
	109	+
	110	+ 3. DType symbols:
	111	+ - torch_npu._C._cd.DType.<dtype_name> -> torch_npu.<dtype_name>
	112	+ """
	113	+
	114	+ _export_dtype_symbols()
	115	+ import torch_npu as _torch_npu
	116	+
	117	+ globals_dict = _torch_npu.__dict__
	118	+ all_list = _torch_npu.__all__
	119	+
	120	+ _export_lazy_python_apis(globals_dict, all_list)
	121	+ _export_npu_ops(globals_dict, all_list)
	122	+ _export_dtype_symbols()

Atorch_npu/_init/core/module_loader.py+211-0

@@ -0,0 +1,211 @@
	1	+import importlib
	2	+import inspect
	3	+import os
	4	+import sys
	5	+
	6	+import torch_npu
	7	+from torch_npu._init.core._exports import _export_public_apis
	8	+
	9	+
	10	+_REQUIRED_C_EXTENSION_CHILDREN = [
	11	+ "_cd",
	12	+ "_logging",
	13	+ "_flops_count",
	14	+ "_profiler",
	15	+ "_distributed_c10d",
	16	+]
	17	+
	18	+
	19	+def _register_c_extension_submodules(module, memo=None):
	20	+ """
	21	+ Mirror torch/__init__.py behavior:
	22	+ expose nested extension modules, e.g. torch_npu._C._distributed_c10d,
	23	+ through sys.modules so that Python import machinery can resolve them.
	24	+ """
	25	+ if memo is None:
	26	+ memo = set()
	27	+ if module in memo:
	28	+ return
	29	+ memo.add(module)
	30	+
	31	+ module_name = module.__name__
	32	+ for name in dir(module):
	33	+ member = getattr(module, name)
	34	+ member_name = getattr(member, "__name__", "")
	35	+ if inspect.ismodule(member) and member_name.startswith(module_name):
	36	+ sys.modules.setdefault(member_name, member)
	37	+ _register_c_extension_submodules(member, memo)
	38	+
	39	+
	40	+def _create_child_once(_C, child_attr: str, init_method: str):
	41	+ if hasattr(_C, child_attr):
	42	+ return
	43	+ fn = getattr(_C, init_method, None)
	44	+ if callable(fn):
	45	+ fn()
	46	+
	47	+
	48	+def _initialize_c_extension_children(required_children):
	49	+ """
	50	+ Create and expose torch_npu._C child submodules.
	51	+ Every child submodule should be created here exactly once.
	52	+ Business Python modules must only consume them, not create them.
	53	+ """
	54	+ import torch_npu._C as _C # ensure torch_npu._C is imported
	55	+
	56	+ # Fixed order for child-module creation.
	57	+ _create_child_once(_C, "_profiler", "_profiler_init")
	58	+ _create_child_once(_C, "_distributed_c10d", "_c10d_npu_init")
	59	+ _create_child_once(_C, "_cd", "_cd_init")
	60	+ _create_child_once(_C, "_logging", "_logging_init")
	61	+ _create_child_once(_C, "_flops_count", "_flops_count_init")
	62	+
	63	+ # Optional RPC child, only if built.
	64	+ _create_child_once(_C, "_distributed_rpc", "_rpc_npu_init")
	65	+
	66	+ _register_c_extension_submodules(_C)
	67	+ missing = [name for name in required_children if not hasattr(_C, name)]
	68	+ if missing:
	69	+ raise RuntimeError(
	70	+ f"Required torch_npu._C child submodules are missing before Python import: {missing}"
	71	+ )
	72	+
	73	+
	74	+def _initialize_logging_if_needed():
	75	+ """
	76	+ Initialize logging runtime after _C._logging is ready.
	77	+ """
	78	+ if not hasattr(torch_npu._C, "_logging"):
	79	+ raise RuntimeError("torch_npu._C._logging is not initialized")
	80	+
	81	+ from torch_npu._logging._internal import (
	82	+ _add_logging_module,
	83	+ _logging_patch,
	84	+ _update_log_state_from_env,
	85	+ )
	86	+
	87	+ _logging_patch()
	88	+ _add_logging_module()
	89	+ _update_log_state_from_env()
	90	+
	91	+
	92	+def _initialize_profiler_if_needed():
	93	+ """
	94	+ Initialize profiler by enabling non-intrusive profiling hooks after _C._profiler is ready.
	95	+ """
	96	+ if not hasattr(torch_npu._C, "_profiler"):
	97	+ raise RuntimeError("torch_npu._C._profiler is not initialized")
	98	+
	99	+ from torch_npu.profiler._non_intrusive_profile import _NonIntrusiveProfile
	100	+
	101	+ _NonIntrusiveProfile.init()
	102	+
	103	+
	104	+def _initialize_rendezvous_if_needed():
	105	+ """
	106	+ Initialize rendezvous after distributed C-extension support is ready.
	107	+ """
	108	+ from torch_npu.distributed import is_available, rendezvous
	109	+ from torch_npu.utils._error_code import dist_error, ErrCode
	110	+
	111	+ if is_available() and not hasattr(torch_npu._C, "_distributed_c10d"):
	112	+ raise RuntimeError(
	113	+ "torch_npu._C._distributed_c10d is not initialized"
	114	+ + dist_error(ErrCode.INTERNAL)
	115	+ )
	116	+
	117	+ rendezvous._rendezvous_init()
	118	+
	119	+
	120	+def _check_npu_import():
	121	+ """
	122	+ Probe-import torch_npu.npu to convert environment issues (e.g. missing
	123	+ libhccl.so / libascendcl.so) into friendlier errors.
	124	+ Must be called only after _C / required child submodules / torch.npu are ready.
	125	+ """
	126	+ try:
	127	+ import torch_npu.npu # noqa: F401
	128	+ except ImportError as e:
	129	+ from torch_npu.utils._error_code import ErrCode, pta_error
	130	+
	131	+ if "libhccl.so" in str(e):
	132	+ if "ASCEND_OPP_PATH" in os.environ:
	133	+ # Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!!
	134	+ e.msg += (
	135	+ ". Please check that the compiler package is installed. "
	136	+ "Please run 'source set_env.sh' in the CANN installation path."
	137	+ + pta_error(ErrCode.NOT_FOUND)
	138	+ )
	139	+ else:
	140	+ # Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!!
	141	+ e.msg += (
	142	+ ". Please check that the cann package is installed. "
	143	+ "Please run 'source set_env.sh' in the CANN installation path."
	144	+ + pta_error(ErrCode.NOT_FOUND)
	145	+ )
	146	+ elif "libascendcl.so" in str(e):
	147	+ # Warning: key logs in the fault mode library!!! Don't make arbitrary modifications!!!
	148	+ e.msg += (
	149	+ ". Please check that the runtime package is installed. "
	150	+ "Please run 'source set_env.sh' in the CANN installation path."
	151	+ + pta_error(ErrCode.NOT_FOUND)
	152	+ )
	153	+ raise
	154	+
	155	+
	156	+def _load_core_modules():
	157	+ """
	158	+ Load torch_npu core modules.
	159	+
	160	+ Includes:
	161	+ 1. C extension child modules creation and exposure.
	162	+ 2. Python runtime core support initialization:
	163	+ logging / profiler / distributed runtime.
	164	+ 3. Check npu backend.
	165	+ 4. Python registration modules imported for side effects.
	166	+ 5. Public API export.
	167	+
	168	+ Dependency:
	169	+ - Runtime support must run after _C child modules are ready.
	170	+ """
	171	+ _initialize_c_extension_children(required_children=_REQUIRED_C_EXTENSION_CHILDREN)
	172	+
	173	+ # Do not hide these in another wrapper. Keep dependencies explicit.
	174	+ _initialize_logging_if_needed()
	175	+ _initialize_profiler_if_needed()
	176	+ _initialize_rendezvous_if_needed()
	177	+
	178	+ _check_npu_import()
	179	+
	180	+ _load_registration_modules()
	181	+ _export_public_apis()
	182	+
	183	+
	184	+def _load_registration_modules():
	185	+ """
	186	+ Import Python modules that rely on import-time side effects or old top-level submodule availability.
	187	+
	188	+ Dependency:
	189	+ - Must run after core module loading.
	190	+ - _C child modules must already be ready.
	191	+
	192	+ Includes:
	193	+ - ACLNN backend config module
	194	+ - old-compatible submodule imports: torch_npu.optim, torch_npu._afd
	195	+ - AFD op bindings: torch.ops.npu.<afd_op> -> torch_npu._afd.<afd_op>
	196	+ - custom ops import
	197	+ - op-plugin registration / meta registration / generated docs side effects
	198	+ """
	199	+ import torch_npu._afd # noqa: F401
	200	+ import torch_npu.npu.aclnn # noqa: F401
	201	+ import torch_npu.op_plugin
	202	+ import torch_npu.optim # noqa: F401
	203	+ from torch_npu.op_plugin.meta import _meta_registrations # noqa: F401
	204	+ from torch_npu.utils import custom_ops # noqa: F401
	205	+ from torch_npu.utils._afd_ops import initialize_afd_bindings
	206	+
	207	+ importlib.import_module("torch_npu._op_plugin_docs")
	208	+ if hasattr(torch_npu, "_op_plugin_docs"):
	209	+ delattr(torch_npu, "_op_plugin_docs")
	210	+
	211	+ initialize_afd_bindings()

Atorch_npu/_init/core/optional_features.py+43-0

@@ -0,0 +1,43 @@
	1	+import os
	2	+import warnings
	3	+
	4	+from torch_npu.utils.utils import _is_interactive_command_line
	5	+
	6	+
	7	+def _enable_sanitizer_if_needed():
	8	+ """
	9	+ Enable NPU Sanitizer.
	10	+ """
	11	+ if "TORCH_NPU_SANITIZER" in os.environ:
	12	+ import torch_npu.npu._sanitizer as csan
	13	+
	14	+ csan.enable_npu_sanitizer()
	15	+
	16	+
	17	+def _configure_interactive_mode():
	18	+ if _is_interactive_command_line():
	19	+ os.environ["TASK_QUEUE_ENABLE"] = "0"
	20	+ warnings.warn(
	21	+ "On the interactive interface, the value of TASK_QUEUE_ENABLE is set to 0 by default. "
	22	+ "Do not set it to 1 to prevent some unknown errors"
	23	+ )
	24	+
	25	+
	26	+def _enable_transfer_to_npu_if_needed():
	27	+ """
	28	+ Enable transfer_to_npu via environment variable
	29	+ """
	30	+ transfer_to_npu_env = os.getenv("TORCH_TRANSFER_TO_NPU", "0")
	31	+ if transfer_to_npu_env == "1":
	32	+ from torch_npu.contrib import transfer_to_npu # noqa: F401
	33	+ elif transfer_to_npu_env != "0":
	34	+ raise ValueError(
	35	+ f"Invalid value for TORCH_TRANSFER_TO_NPU: {transfer_to_npu_env}. "
	36	+ "Only '0' or '1' is supported."
	37	+ )
	38	+
	39	+
	40	+def _enable_optional_features():
	41	+ _enable_sanitizer_if_needed()
	42	+ _configure_interactive_mode()
	43	+ _enable_transfer_to_npu_if_needed()

Atorch_npu/_init/core/runtime_lifecycle.py+40-0

@@ -0,0 +1,40 @@
	1	+import atexit
	2	+
	3	+import torch_npu
	4	+
	5	+
	6	+def _npu_shutdown():
	7	+ """
	8	+ NPU exit, need to synchronize devices
	9	+ """
	10	+ from torch_npu.asd.asd import matmul_check
	11	+ from torch_npu.utils._error_code import _except_handler
	12	+
	13	+ success = torch_npu._C._npu_shutdown_synchronize()
	14	+ torch_npu.distributed.distributed_c10d._destructor_process_group()
	15	+ torch_npu._C._npu_shutdown(success)
	16	+ _except_handler.handle_exception()
	17	+
	18	+ matmul_check._cleanup()
	19	+
	20	+ if torch_npu.npu.aclnn._use_static_aclnn_kernel:
	21	+ from torch_npu._inductor.npu_static_kernel import uninstall_static_kernel
	22	+
	23	+ uninstall_static_kernel()
	24	+
	25	+
	26	+def _initialize_runtime_lifecycle():
	27	+ """
	28	+ Complete C extension initialization and register process-exit cleanup.
	29	+
	30	+ This entry is expected to be called once by torch_npu top-level import.
	31	+ """
	32	+ if not hasattr(torch_npu, "_C"):
	33	+ raise RuntimeError("torch_npu._C is not available before extension init")
	34	+
	35	+ if not hasattr(torch_npu._C, "_initExtension"):
	36	+ raise RuntimeError("torch_npu._C._initExtension is not available")
	37	+
	38	+ # final extension barrier, this must be placed at the end
	39	+ torch_npu._C._initExtension()
	40	+ atexit.register(_npu_shutdown)

Atorch_npu/_init/patches/__init__.py+4-0

@@ -0,0 +1,4 @@
	1	+# Internal patch management package for torch_npu.
	2	+# Keep this package initializer side-effect free.
	3	+#
	4	+# Concrete *_patches modules are auto-discovered by PatchManager.

Atorch_npu/_init/patches/api_patches.py+24-0

@@ -0,0 +1,24 @@
	1	+from torch_npu._init.patches.patch_manager import PatchManager
	2	+
	3	+
	4	+@PatchManager.register_patch("api")
	5	+def apply_torch_api_patches():
	6	+ from torch_npu.distributed.fsdp._add_fsdp_patch import _apply_fsdp_patch
	7	+ from torch_npu.multiprocessing.reductions import _add_reductions_methods
	8	+ from torch_npu.utils._module import _apply_module_patch
	9	+ from torch_npu.utils._optim import add_optim_method
	10	+ from torch_npu.utils.collect_env import _add_collect_env_methods
	11	+ from torch_npu.utils.dlpack import _apply_dlpack_patch
	12	+ from torch_npu.utils.serialization import _add_serialization_methods
	13	+ from torch_npu.utils.storage import _add_storage_methods
	14	+ from torch_npu.utils.tensor_methods import _add_tensor_methods
	15	+
	16	+ _add_storage_methods()
	17	+ _apply_dlpack_patch()
	18	+ _apply_module_patch()
	19	+ _add_tensor_methods()
	20	+ _add_serialization_methods()
	21	+ _add_collect_env_methods()
	22	+ add_optim_method()
	23	+ _add_reductions_methods()
	24	+ _apply_fsdp_patch()

Atorch_npu/_init/patches/asd_patches.py+8-0

@@ -0,0 +1,8 @@
	1	+from torch_npu._init.patches.patch_manager import PatchManager
	2	+
	3	+
	4	+@PatchManager.register_patch("asd")
	5	+def apply_asd_patch():
	6	+ from torch_npu.asd.asd import _asd_patch
	7	+
	8	+ _asd_patch()

Atorch_npu/_init/patches/distributed_patches.py+154-0

@@ -0,0 +1,154 @@
	1	+import torch
	2	+import torch.distributed.launcher.api
	3	+
	4	+import torch_npu
	5	+from torch_npu._init.patches.patch_manager import PatchManager
	6	+
	7	+
	8	+# 1. Replace PyTorch internal distributed implementations with NPU/HCCL implementations.
	9	+_INTERNAL_REPLACEMENTS = [
	10	+ (
	11	+ "_C._distributed_c10d._verify_params_across_processes",
	12	+ "distributed._verify_params_across_processes",
	13	+ ),
	14	+ (
	15	+ "_C._distributed_c10d.ProcessGroup._get_sequence_number_for_group",
	16	+ "distributed.distributed_c10d._hccl_get_sequence_number_for_group",
	17	+ ),
	18	+ (
	19	+ "distributed.distributed_c10d._add_ephemeral_timeout_for_all_pgs",
	20	+ "distributed.distributed_c10d._hccl_add_ephemeral_timeout_for_all_pgs",
	21	+ ),
	22	+]
	23	+
	24	+
	25	+# 2. Expose torch_npu distributed implementations through torch.distributed APIs.
	26	+_PUBLIC_API_ALIASES = [
	27	+ (
	28	+ "distributed.batch_isend_irecv",
	29	+ "distributed.distributed_c10d._batch_isend_irecv",
	30	+ ),
	31	+ (
	32	+ "distributed.distributed_c10d.batch_isend_irecv",
	33	+ "distributed.distributed_c10d._batch_isend_irecv",
	34	+ ),
	35	+ (
	36	+ "distributed.gather",
	37	+ "distributed.distributed_c10d._gather",
	38	+ ),
	39	+ (
	40	+ "distributed.distributed_c10d.gather",
	41	+ "distributed.distributed_c10d._gather",
	42	+ ),
	43	+ (
	44	+ "distributed.gather_object",
	45	+ "distributed.distributed_c10d._gather_object",
	46	+ ),
	47	+ (
	48	+ "distributed.distributed_c10d.gather_object",
	49	+ "distributed.distributed_c10d._gather_object",
	50	+ ),
	51	+ (
	52	+ "distributed.is_hccl_available",
	53	+ "distributed.is_hccl_available",
	54	+ ),
	55	+ (
	56	+ "distributed.reinit_process_group",
	57	+ "distributed.reinit_process_group",
	58	+ ),
	59	+]
	60	+
	61	+
	62	+def _resolve_attr(root, attr_path: str):
	63	+ obj = root
	64	+ for part in attr_path.split("."):
	65	+ obj = getattr(obj, part)
	66	+ return obj
	67	+
	68	+
	69	+def _assign_attr(target_root, target_path: str, source_root, source_path: str):
	70	+ parts = target_path.split(".")
	71	+ owner = (
	72	+ _resolve_attr(target_root, ".".join(parts[:-1]))
	73	+ if len(parts) > 1
	74	+ else target_root
	75	+ )
	76	+ setattr(owner, parts[-1], _resolve_attr(source_root, source_path))
	77	+
	78	+
	79	+def _apply_internal_replacements(torch, torch_npu):
	80	+ """
	81	+ Replace PyTorch internal distributed implementations with NPU/HCCL versions.
	82	+
	83	+ Example:
	84	+ torch._C._distributed_c10d._verify_params_across_processes
	85	+ -> torch_npu.distributed._verify_params_across_processes
	86	+ """
	87	+ for target_path, source_path in _INTERNAL_REPLACEMENTS:
	88	+ _assign_attr(torch, target_path, torch_npu, source_path)
	89	+
	90	+
	91	+def _apply_public_api_aliases(torch, torch_npu):
	92	+ """
	93	+ Patch torch.distributed public APIs with torch_npu implementations.
	94	+
	95	+ Example:
	96	+ torch.distributed.gather(...) -> torch_npu.distributed.distributed_c10d._gather(...)
	97	+ """
	98	+ for target_path, source_path in _PUBLIC_API_ALIASES:
	99	+ _assign_attr(torch, target_path, torch_npu, source_path)
	100	+
	101	+
	102	+def _apply_wrapped_functions(torch, torch_npu):
	103	+ """
	104	+ Wrap PyTorch distributed helpers with torch_npu NPU/HCCL logic.
	105	+
	106	+ Example:
	107	+ torch.distributed.distributed_c10d.rendezvous(...)
	108	+ -> torch_npu.distributed.distributed_c10d._trigger_rendezvous_decorator(...)
	109	+
	110	+ torch.distributed.launcher.api._get_addr_and_port(...)
	111	+ -> torch_npu.distributed.distributed_c10d._trigger__get_addr_and_port_decorator(...)
	112	+ """
	113	+ torch.distributed.distributed_c10d.rendezvous = (
	114	+ torch_npu.distributed.distributed_c10d._trigger_rendezvous_decorator(
	115	+ torch.distributed.distributed_c10d.rendezvous
	116	+ )
	117	+ )
	118	+
	119	+ torch.distributed.launcher.api._get_addr_and_port = (
	120	+ torch_npu.distributed.distributed_c10d._trigger__get_addr_and_port_decorator(
	121	+ torch.distributed.launcher.api._get_addr_and_port
	122	+ )
	123	+ )
	124	+
	125	+
	126	+def _apply_sharded_grad_scaler_patch(torch):
	127	+ """
	128	+ Replace PyTorch FSDP ShardedGradScaler with torch_npu implementation.
	129	+
	130	+ Example:
	131	+ torch.distributed.fsdp.sharded_grad_scaler.ShardedGradScaler
	132	+ -> torch_npu.npu.amp.sharded_grad_scaler._ShardedGradScaler
	133	+ """
	134	+ from torch.distributed.fsdp import sharded_grad_scaler
	135	+
	136	+ from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler
	137	+
	138	+ sharded_grad_scaler.ShardedGradScaler = _ShardedGradScaler
	139	+
	140	+
	141	+@PatchManager.register_patch("distributed")
	142	+def apply_distributed_methods_patch():
	143	+ """
	144	+ Patch PyTorch distributed APIs with torch_npu NPU/HCCL implementations.
	145	+
	146	+ Categories:
	147	+ 1. Internal replacements
	148	+ 2. Public API aliases
	149	+ 3. Wrapped functions
	150	+ """
	151	+ _apply_internal_replacements(torch, torch_npu)
	152	+ _apply_public_api_aliases(torch, torch_npu)
	153	+ _apply_sharded_grad_scaler_patch(torch)
	154	+ _apply_wrapped_functions(torch, torch_npu)

Atorch_npu/_init/patches/dynamo_patches.py+15-0

@@ -0,0 +1,15 @@
	1	+from torch_npu._init.patches.patch_manager import PatchManager
	2	+
	3	+
	4	+@PatchManager.register_patch("dynamo")
	5	+def apply_dynamo_methods_patch():
	6	+ from torch_npu.utils._dynamo import add_dynamo_methods
	7	+
	8	+ add_dynamo_methods()
	9	+
	10	+
	11	+@PatchManager.register_patch("dynamo")
	12	+def apply_npugraph_tree_patch():
	13	+ from torch_npu.utils._graph_tree import _apply_npugraph_tree_methods
	14	+
	15	+ _apply_npugraph_tree_methods()

Atorch_npu/_init/patches/monkey_patches.py+57-0

@@ -0,0 +1,57 @@
	1	+import sys
	2	+import types
	3	+
	4	+import torch
	5	+
	6	+from torch_npu._init.patches.patch_manager import PatchManager
	7	+from torch_npu.contrib.function import npu_functional
	8	+from torch_npu.contrib.module import npu_modules
	9	+from torch_npu.utils._error_code import ErrCode, pta_error
	10	+
	11	+
	12	+all_monkey_patches = [
	13	+ ["nn.functional", npu_functional],
	14	+ ["nn", npu_modules],
	15	+]
	16	+
	17	+
	18	+def _apply_patches(monkey_patches):
	19	+ def _getattr(module_list, root_module=torch):
	20	+ if len(module_list) <= 1:
	21	+ return root_module
	22	+
	23	+ if hasattr(root_module, module_list[0]):
	24	+ return _getattr(module_list[1:], getattr(root_module, module_list[0]))
	25	+
	26	+ empty_module_name = f"{root_module.__name__}.{module_list[0]}"
	27	+ sys.modules[empty_module_name] = types.ModuleType(empty_module_name)
	28	+ setattr(root_module, module_list[0], sys.modules.get(empty_module_name))
	29	+ return _getattr(module_list[1:], getattr(root_module, module_list[0]))
	30	+
	31	+ for dest, patch in monkey_patches:
	32	+ dest_module = _getattr(dest.split("."), root_module=torch)
	33	+ last_module_level = dest.split(".")[-1]
	34	+
	35	+ if not isinstance(patch, types.ModuleType):
	36	+ setattr(dest_module, last_module_level, patch)
	37	+ continue
	38	+
	39	+ if not hasattr(dest_module, last_module_level) or not hasattr(patch, "__all__"):
	40	+ setattr(dest_module, last_module_level, patch)
	41	+ sys.modules[f"{dest_module.__name__}.{last_module_level}"] = patch
	42	+ continue
	43	+
	44	+ if not hasattr(patch, "__all__"):
	45	+ raise NotImplementedError(
	46	+ "Patch module must have __all__ definition."
	47	+ + pta_error(ErrCode.NOT_SUPPORT)
	48	+ )
	49	+
	50	+ dest_module = getattr(dest_module, last_module_level)
	51	+ for attr in patch.__all__:
	52	+ setattr(dest_module, attr, getattr(patch, attr))
	53	+
	54	+
	55	+@PatchManager.register_patch("monkey")
	56	+def apply_monkey_patches():
	57	+ _apply_patches(all_monkey_patches)

Atorch_npu/_init/patches/npu_patches.py+22-0

@@ -0,0 +1,22 @@
	1	+from torch_npu._init.patches.patch_manager import PatchManager
	2	+
	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")
	12	+def apply_npu_format_patch():
	13	+ from torch_npu.npu._format import _apply_npu_format_patch
	14	+
	15	+ _apply_npu_format_patch()
	16	+
	17	+
	18	+@PatchManager.register_patch("npu")
	19	+def apply_npu_meta_patch():
	20	+ from torch_npu.utils._npu_meta_registration import npu_patch_meta
	21	+
	22	+ npu_patch_meta()

Atorch_npu/_init/patches/patch_manager.py+216-0

@@ -0,0 +1,216 @@
	1	+import pkgutil
	2	+from collections import defaultdict
	3	+from collections.abc import Callable
	4	+from importlib import import_module
	5	+
	6	+
	7	+PatchFn = Callable[[], None]
	8	+
	9	+
	10	+class PatchManager:
	11	+ """
	12	+ Central patch registry for torch_npu.
	13	+
	14	+ PatchManager discovers patch modules in two ways:
	15	+ 1. Auto-import modules under torch_npu._init.patches whose names end with
	16	+ '_patches';
	17	+ 2. Import external patch modules declared by DEFAULT_EXTRA_PATCH_MODULES
	18	+ or registered through register_patch_module().
	19	+
	20	+ Patch modules should register patch functions through:
	21	+
	22	+ @PatchManager.register_patch("group_name")
	23	+ def apply_xxx_patch():
	24	+ ...
	25	+ """
	26	+
	27	+ DEFAULT_PATCH_ORDER = [
	28	+ "monkey",
	29	+ "api",
	30	+ "distributed",
	31	+ "dynamo",
	32	+ "profiler",
	33	+ "npu",
	34	+ "warning",
	35	+ "asd",
	36	+ ]
	37	+
	38	+ DEFAULT_EXTRA_PATCH_MODULES = [
	39	+ # Component-owned patch modules outside torch_npu._init.patches.
	40	+ # Example:
	41	+ # "torch_npu.some_component.foo",
	42	+ ]
	43	+
	44	+ PATCH_MODULE_SUFFIX = "_patches"
	45	+
	46	+ _applied_patch_count = defaultdict(int)
	47	+ _builtin_patches_registered = False
	48	+ _custom_full_patch_order: list[str] \| None = None
	49	+ _patch_groups = defaultdict(list)
	50	+ _patch_modules: list[str] = []
	51	+
	52	+ @classmethod
	53	+ def _add_patch(cls, group: str, fn: PatchFn):
	54	+ if fn not in cls._patch_groups[group]:
	55	+ cls._patch_groups[group].append(fn)
	56	+
	57	+ @classmethod
	58	+ def register_patch(cls, group: str, fn: PatchFn \| None = None):
	59	+ """
	60	+ Register a patch function into a patch group.
	61	+
	62	+ Supports:
	63	+
	64	+ PatchManager.register_patch("graph", apply_graph_patch)
	65	+
	66	+ and:
	67	+
	68	+ @PatchManager.register_patch("graph")
	69	+ def apply_graph_patch():
	70	+ ...
	71	+ """
	72	+ if not isinstance(group, str) or not group:
	73	+ raise ValueError("patch group must be a non-empty string")
	74	+
	75	+ def decorator(real_fn: PatchFn):
	76	+ cls._add_patch(group, real_fn)
	77	+ return real_fn
	78	+
	79	+ if fn is not None:
	80	+ return decorator(fn)
	81	+
	82	+ return decorator
	83	+
	84	+ @classmethod
	85	+ def _resolve_patch_order(cls) -> list[str]:
	86	+ """
	87	+ Resolve final patch execution order.
	88	+
	89	+ - Built-in groups follow DEFAULT_PATCH_ORDER.
	90	+ - Groups not listed in DEFAULT_PATCH_ORDER are appended after default groups.
	91	+ """
	92	+ if cls._custom_full_patch_order is not None:
	93	+ base_order = list(cls._custom_full_patch_order)
	94	+ else:
	95	+ base_order = list(cls.DEFAULT_PATCH_ORDER)
	96	+
	97	+ extra_groups = [group for group in cls._patch_groups if group not in base_order]
	98	+
	99	+ return base_order + extra_groups
	100	+
	101	+ @classmethod
	102	+ def _register_builtin_patches(cls):
	103	+ """
	104	+ Discover and import patch modules.
	105	+
	106	+ This method imports:
	107	+ 1. built-in patch modules under torch_npu._init.patches whose names end
	108	+ with '_patches';
	109	+ 2. default external patch modules declared in DEFAULT_EXTRA_PATCH_MODULES;
	110	+ 3. external patch modules registered by register_patch_module().
	111	+
	112	+ Importing a patch module triggers @PatchManager.register_patch(...)
	113	+ decorators inside that module.
	114	+ """
	115	+ if cls._builtin_patches_registered:
	116	+ return
	117	+
	118	+ import torch_npu._init.patches as patches_pkg
	119	+
	120	+ # 1. Auto import torch_npu._init.patches/*_patches.py
	121	+ prefix = patches_pkg.__name__ + "."
	122	+ for module_info in sorted(
	123	+ pkgutil.iter_modules(patches_pkg.__path__), key=lambda x: x.name
	124	+ ):
	125	+ module_name = module_info.name
	126	+
	127	+ if module_name == "patch_manager":
	128	+ continue
	129	+ if not module_name.endswith(cls.PATCH_MODULE_SUFFIX):
	130	+ continue
	131	+
	132	+ import_module(prefix + module_name)
	133	+
	134	+ # 2. Register built-in external patch modules
	135	+ for module_name in cls.DEFAULT_EXTRA_PATCH_MODULES:
	136	+ cls.register_patch_module(module_name)
	137	+
	138	+ # 3. Import registered external patch modules
	139	+ for module_name in cls._patch_modules:
	140	+ import_module(module_name)
	141	+
	142	+ cls._builtin_patches_registered = True
	143	+
	144	+ @classmethod
	145	+ def apply_registered_patches(cls, group: str):
	146	+ """
	147	+ Apply newly registered patches in one group.
	148	+
	149	+ This supports delayed apply:
	150	+ if new patch functions are registered after this group was applied,
	151	+ calling this method again only applies newly added patch functions.
	152	+ """
	153	+ cls._register_builtin_patches()
	154	+
	155	+ patches = cls._patch_groups.get(group, [])
	156	+ start = cls._applied_patch_count[group]
	157	+
	158	+ for patch in patches[start:]:
	159	+ patch()
	160	+
	161	+ cls._applied_patch_count[group] = len(patches)
	162	+
	163	+ @staticmethod
	164	+ def _patch_excepthook():
	165	+ """
	166	+ Patch Python global exception hook.
	167	+
	168	+ Kept separate from normal patches because it is paired with shutdown
	169	+ exception handling.
	170	+ """
	171	+ from torch_npu.utils._error_code import _except_handler
	172	+
	173	+ _except_handler.patch_excepthook()
	174	+
	175	+ @classmethod
	176	+ def register_patch_module(cls, module_name: str):
	177	+ """
	178	+ Register an extra module that contains @PatchManager.register_patch(...)
	179	+ decorators. The module will be imported before patches are applied.
	180	+ """
	181	+ if not isinstance(module_name, str) or not module_name:
	182	+ raise ValueError("patch module name must be a non-empty string")
	183	+
	184	+ if module_name not in cls._patch_modules:
	185	+ cls._patch_modules.append(module_name)
	186	+
	187	+ @classmethod
	188	+ def set_patch_order(cls, order: list[str]):
	189	+ """
	190	+ Override base patch group order.
	191	+
	192	+ Must be called before PatchManager.run().
	193	+ Registered groups not listed in this order will still be appended after
	194	+ the base order by _resolve_patch_order().
	195	+ """
	196	+ cls._custom_full_patch_order = list(order)
	197	+
	198	+ @classmethod
	199	+ def clear_for_test(cls):
	200	+ """
	201	+ Test-only helper. Do not use in normal runtime.
	202	+ """
	203	+ cls._patch_groups.clear()
	204	+ cls._patch_modules.clear()
	205	+ cls._custom_full_patch_order = None
	206	+ cls._builtin_patches_registered = False
	207	+ cls._applied_patch_count.clear()
	208	+
	209	+
	210	+def _apply_patches():
	211	+ PatchManager._register_builtin_patches()
	212	+
	213	+ for group in PatchManager._resolve_patch_order():
	214	+ PatchManager.apply_registered_patches(group)
	215	+
	216	+ PatchManager._patch_excepthook()

Atorch_npu/_init/patches/profiler_patches.py+15-0

@@ -0,0 +1,15 @@
	1	+from torch_npu._init.patches.patch_manager import PatchManager
	2	+
	3	+
	4	+@PatchManager.register_patch("profiler")
	5	+def apply_mstx_patch():
	6	+ from torch_npu.profiler._add_mstx_patch import _apply_mstx_patch
	7	+
	8	+ _apply_mstx_patch()
	9	+
	10	+
	11	+@PatchManager.register_patch("profiler")
	12	+def apply_perf_dump_patch():
	13	+ from torch_npu.utils._step import add_perf_dump_patch
	14	+
	15	+ add_perf_dump_patch()

Atorch_npu/_init/patches/warning_patches.py+44-0

@@ -0,0 +1,44 @@
	1	+import warnings
	2	+from functools import wraps
	3	+
	4	+from torch_npu._init.patches.patch_manager import PatchManager
	5	+
	6	+
	7	+_WARN_MSG = {
	8	+ "DropoutWithByteMask": (
	9	+ "torch.nn.DropoutWithByteMask is deprecated and will be removed in future version. "
	10	+ "Use torch_npu.contrib.module.DropoutWithByteMask instead."
	11	+ ),
	12	+ "dropout_with_byte_mask": (
	13	+ "torch.nn.functional.dropout_with_byte_mask is deprecated and will be removed in future version. "
	14	+ "Use torch_npu.contrib.function.dropout_with_byte_mask instead."
	15	+ ),
	16	+}
	17	+
	18	+
	19	+def _wrap_torch_patch_warning_func(func):
	20	+ @wraps(func)
	21	+ def wrapper(args, *kwargs):
	22	+ warnings.warn(_WARN_MSG[func.__name__])
	23	+ return func(args, *kwargs)
	24	+
	25	+ return wrapper
	26	+
	27	+
	28	+@PatchManager.register_patch("warning")
	29	+def apply_npu_show_warning_patch():
	30	+ from torch_npu.utils.utils import _apply_npu_show_warning
	31	+
	32	+ _apply_npu_show_warning()
	33	+
	34	+
	35	+@PatchManager.register_patch("warning")
	36	+def apply_deprecated_api_warning_patch():
	37	+ import torch
	38	+
	39	+ torch.nn.DropoutWithByteMask = _wrap_torch_patch_warning_func(
	40	+ torch.nn.DropoutWithByteMask
	41	+ )
	42	+ torch.nn.functional.dropout_with_byte_mask = _wrap_torch_patch_warning_func(
	43	+ torch.nn.functional.dropout_with_byte_mask
	44	+ )

Atorch_npu/_init/registry/__init__.py+2-0

@@ -0,0 +1,2 @@
	1	+# Internal registry management package for torch_npu.
	2	+# Keep this package initializer side-effect free.

Atorch_npu/_init/registry/backend.py+25-0

@@ -0,0 +1,25 @@
	1	+import torch
	2	+
	3	+import torch_npu
	4	+
	5	+
	6	+def register_privateuse1_backend():
	7	+ torch.utils.rename_privateuse1_backend("npu")
	8	+ # rename device name to 'npu' and register funcs
	9	+ torch._register_device_module("npu", torch_npu.npu)
	10	+ unsupported_dtype = [
	11	+ torch.quint8,
	12	+ torch.quint4x2,
	13	+ torch.quint2x4,
	14	+ torch.qint32,
	15	+ torch.qint8,
	16	+ ]
	17	+ torch.utils.generate_methods_for_privateuse1_backend(
	18	+ for_tensor=True,
	19	+ for_module=True,
	20	+ for_storage=True,
	21	+ unsupported_dtype=unsupported_dtype,
	22	+ )
	23	+ torch.nn.parameter.UninitializedTensorMixin._allowed_methods.append(
	24	+ torch.Tensor.npu
	25	+ )

Atorch_npu/_init/registry/distributed.py+54-0

@@ -0,0 +1,54 @@
	1	+import torch
	2	+
	3	+
	4	+def _new_process_group_hccl_helper(dist_backend_opts, pg_options):
	5	+ import torch_npu
	6	+
	7	+ store = dist_backend_opts.store
	8	+ group_rank = dist_backend_opts.group_rank
	9	+ group_size = dist_backend_opts.group_size
	10	+ if pg_options is None or not isinstance(
	11	+ pg_options, torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options
	12	+ ):
	13	+ pg_options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options()
	14	+ pg_options.is_high_priority_stream = False
	15	+ pg_options._timeout = dist_backend_opts.timeout
	16	+ pg_options.global_ranks_in_group = dist_backend_opts.global_ranks_in_group
	17	+ pg_options.group_id = dist_backend_opts.group_id
	18	+ return torch_npu._C._distributed_c10d.ProcessGroupHCCL(
	19	+ store, group_rank, group_size, pg_options
	20	+ )
	21	+
	22	+
	23	+def _new_process_group_lccl_helper(dist_backend_opts, pg_options):
	24	+ import torch_npu
	25	+
	26	+ store = dist_backend_opts.store
	27	+ group_rank = dist_backend_opts.group_rank
	28	+ group_size = dist_backend_opts.group_size
	29	+ return torch_npu._C._distributed_c10d.ProcessGroupLCCL(
	30	+ store, group_rank, group_size
	31	+ )
	32	+
	33	+
	34	+def register_distributed_backend_for_npu():
	35	+ # init and register hccl backend
	36	+ # Note: Since torch 2.8, the hccl backend must be registered at first to keep a right default_device_backend_map
	37	+ torch.distributed.Backend.register_backend(
	38	+ "hccl",
	39	+ lambda dist_backend_opts, pg_options: _new_process_group_hccl_helper(
	40	+ dist_backend_opts, pg_options
	41	+ ),
	42	+ extended_api=True,
	43	+ devices=["npu"],
	44	+ )
	45	+
	46	+ # init and register lccl backend
	47	+ torch.distributed.Backend.register_backend(
	48	+ "lccl",
	49	+ lambda dist_backend_opts, pg_options: _new_process_group_lccl_helper(
	50	+ dist_backend_opts, pg_options
	51	+ ),
	52	+ extended_api=True,
	53	+ devices=["npu"],
	54	+ )

Atorch_npu/_init/registry/dynamo.py+32-0

@@ -0,0 +1,32 @@
	1	+from torch._dynamo.device_interface import register_interface_for_device
	2	+
	3	+from torch_npu.utils._dynamo_device import NpuInterface
	4	+
	5	+
	6	+def _dynamo_register_interface_for_device():
	7	+ register_interface_for_device("npu", NpuInterface)
	8	+ for i in range(32):
	9	+ register_interface_for_device(f"npu:{i}", NpuInterface)
	10	+
	11	+
	12	+def register_dynamo_backends():
	13	+ from torch_npu.dynamo import _register_backends
	14	+
	15	+ _register_backends()
	16	+
	17	+
	18	+def register_dynamo_device_interface():
	19	+ """
	20	+ Register NPU device interface for Dynamo
	21	+ """
	22	+ _dynamo_register_interface_for_device()
	23	+
	24	+
	25	+def register_dynamo_trace_rules():
	26	+ """
	27	+ # Support stream into Dynamo charts. Enable Dynamo to recognize NPU
	28	+ stream/device/memory/random APIs and related torch_npu._C bindings during graph capture.
	29	+ """
	30	+ from torch_npu.dynamo.trace_rule import _patch_npu_trace_rules
	31	+
	32	+ _patch_npu_trace_rules()

Atorch_npu/_init/registry/registry_manager.py+121-0

@@ -0,0 +1,121 @@
	1	+import torch
	2	+from torch.utils.checkpoint import DefaultDeviceType
	3	+
	4	+import torch_npu
	5	+
	6	+
	7	+def _register_npu_backend():
	8	+ """
	9	+ Register core NPU backend capability:
	10	+ - privateuse1 backend
	11	+ - torch.npu device module
	12	+ - Tensor / Module / Storage npu methods
	13	+ - CANN package / environment check
	14	+
	15	+ Note:
	16	+ This function must not initialize NPU runtime.
	17	+ NPU runtime initialization is ownde by torch_npu.npu._lazy_init().
	18	+ """
	19	+ from torch_npu._init.registry.backend import register_privateuse1_backend
	20	+ from torch_npu.utils.npu_intercept import _cann_package_check
	21	+
	22	+ register_privateuse1_backend()
	23	+ _cann_package_check()
	24	+
	25	+ if not hasattr(torch, "npu"):
	26	+ raise RuntimeError(
	27	+ "torch.npu is not registered after privateuse1 backend registration"
	28	+ )
	29	+
	30	+
	31	+def _register_distributed():
	32	+ """
	33	+ Register distributed backend for NPU.
	34	+
	35	+ Dependency:
	36	+ - _C._distributed_c10d must be ready.
	37	+ - distributed runtime should have been initialized by ModuleLoader.
	38	+ """
	39	+ if not hasattr(torch_npu._C, "_distributed_c10d"):
	40	+ raise RuntimeError(
	41	+ "torch_npu._C._distributed_c10d must be ready before distributed backend registration"
	42	+ )
	43	+
	44	+ from torch_npu._init.registry.distributed import (
	45	+ register_distributed_backend_for_npu,
	46	+ )
	47	+
	48	+ # init and register distributed backend
	49	+ register_distributed_backend_for_npu()
	50	+
	51	+
	52	+def _register_dynamo():
	53	+ """
	54	+ Register Dynamo integration:
	55	+ - Dynamo backend
	56	+ - Dynamo device interface
	57	+ - NPU trace rules for Dynamo
	58	+ """
	59	+ from torch_npu._init.registry.dynamo import (
	60	+ register_dynamo_backends,
	61	+ register_dynamo_device_interface,
	62	+ register_dynamo_trace_rules,
	63	+ )
	64	+
	65	+ register_dynamo_backends()
	66	+ register_dynamo_device_interface()
	67	+
	68	+ # Do not repeat this call for register_dynamo_trace_rules appends rules into
	69	+ # Dynamo's global rules maps.
	70	+ register_dynamo_trace_rules()
	71	+
	72	+
	73	+def _register_rpc():
	74	+ """
	75	+ Register and init RPC NPU backend.
	76	+ """
	77	+ from torch_npu.distributed.rpc.backend_registry import _rpc_backend_registry
	78	+
	79	+ _rpc_backend_registry()
	80	+
	81	+
	82	+def _register_inductor():
	83	+ """
	84	+ Register lightweight NPU device op overrides for Inductor.
	85	+ Do not import toch_npu._inductor here: toch_npu._inductor performs full NPU
	86	+ Inductor backend loading and heavy global patches lazily when torch.compile
	87	+ and Inductor path is actually used.
	88	+ """
	89	+ from torch_npu.utils._inductor import _inductor_register_device_op_overrides
	90	+
	91	+ _inductor_register_device_op_overrides()
	92	+
	93	+
	94	+def _register_default_gradient_device_type():
	95	+ """
	96	+ Set default device type for gradient checkpointing.
	97	+ """
	98	+ DefaultDeviceType.set_device_type("npu")
	99	+
	100	+
	101	+def _register_components():
	102	+ """
	103	+ Register torch_npu backend and integration capabilities.
	104	+
	105	+ Order matters:
	106	+ 1. NPU backend is the base capability.
	107	+ 2. Distributed and Dynamo depend on NPU backend / _C children.
	108	+ 3. RPC, dtensor and inductor are Python-side framework integrations.
	109	+ 4. DefaultDeviceType is set after NPU backend is registered.
	110	+ """
	111	+ if not hasattr(torch_npu, "_C"):
	112	+ raise RuntimeError(
	113	+ "torch_npu._C is not available before torch_npu registry init"
	114	+ )
	115	+
	116	+ _register_npu_backend()
	117	+ _register_distributed()
	118	+ _register_dynamo()
	119	+ _register_rpc()
	120	+ _register_inductor()
	121	+ _register_default_gradient_device_type()

Mtorch_npu/_logging/__init__.py+0-25

	@@ -1,25 +0,0 @@
1	-__all__ = []
2	-
3	-import os
4	-import re
5	-import logging
6	-import torch._logging._internal
7	-from torch_npu import _C
8	-from ._internal import _logging_patch, _add_logging_module
9	-
10	-
11	-_C._logging_init()
12	-_logging_patch()
13	-_add_logging_module()
14	-
15	-
16	-def _update_log_state_from_env():
17	- log_setting = os.environ.get("TORCH_NPU_LOGS", None)
18	- if log_setting is not None:
19	- torch._logging._internal.LOG_ENV_VAR = "TORCH_NPU_LOGS"
20	- torch._logging._internal._init_logs()
21	- _C._logging._LogContext.GetInstance().setLogs(torch._logging._internal.log_state.log_qname_to_level)
22	- elif os.environ.get("TORCH_LOGS", None) is not None:
23	- _C._logging._LogContext.GetInstance().setLogs(torch._logging._internal.log_state.log_qname_to_level)
24	-
25	-_update_log_state_from_env()

Mtorch_npu/_logging/_internal.py+10-1

	@@ -1,5 +1,4 @@
1	import os	1	import os
2	-import logging
3	import torch._logging._internal	2	import torch._logging._internal
4	from torch_npu import _C	3	from torch_npu import _C
5		4
	@@ -45,3 +44,13 @@ def _add_logging_module():
45	torch._logging._internal.register_log("aclgraph", "torch_npu.aclgraph")	44	torch._logging._internal.register_log("aclgraph", "torch_npu.aclgraph")
46	torch._logging._internal.register_log("npugraph", "torch_npu.npugraph")	45	torch._logging._internal.register_log("npugraph", "torch_npu.npugraph")
47	torch._logging._internal.register_log("cudagraphs", "torch_npu.npugraph")	46	torch._logging._internal.register_log("cudagraphs", "torch_npu.npugraph")
		47	+
		48	+
		49	+def _update_log_state_from_env():
		50	+ log_setting = os.environ.get("TORCH_NPU_LOGS", None)
		51	+ if log_setting is not None:
		52	+ torch._logging._internal.LOG_ENV_VAR = "TORCH_NPU_LOGS"
		53	+ torch._logging._internal._init_logs()
		54	+ _C._logging._LogContext.GetInstance().setLogs(torch._logging._internal.log_state.log_qname_to_level)
		55	+ elif os.environ.get("TORCH_LOGS", None) is not None:
		56	+ _C._logging._LogContext.GetInstance().setLogs(torch._logging._internal.log_state.log_qname_to_level)

Mtorch_npu/distributed/__init__.py+1-8

	@@ -5,7 +5,6 @@ __all__ = [
5	from torch.distributed import _make_nccl_premul_sum as _make_hccl_premul_sum	5	from torch.distributed import _make_nccl_premul_sum as _make_hccl_premul_sum
6		6
7	import torch_npu	7	import torch_npu
8	-from torch_npu.utils._error_code import ErrCode, dist_error
9		8
10		9
11	def is_available():	10	def is_available():
	@@ -20,10 +19,6 @@ def is_available():
20	return hasattr(torch_npu._C, "_c10d_npu_init")	19	return hasattr(torch_npu._C, "_c10d_npu_init")
21		20
22		21
23	-if is_available() and not torch_npu._C._c10d_npu_init():
24	- raise RuntimeError("Failed to initialize torch_npu.distributed" + dist_error(ErrCode.INTERNAL))
25	-
26	-
27	from torch_npu._C._distributed_c10d import (	22	from torch_npu._C._distributed_c10d import (
28	ParallelStore,	23	ParallelStore,
29	_verify_params_across_processes,	24	_verify_params_across_processes,
	@@ -31,7 +26,5 @@ from torch_npu._C._distributed_c10d import (
31	)	26	)
32		27
33		28
34	-from torch_npu.distributed import rendezvous, tensor, nn	29	+from torch_npu.distributed import tensor, nn
35	from .distributed_c10d import is_hccl_available, reinit_process_group, _reduce_scatter_tensor_uneven as reduce_scatter_tensor_uneven, _all_gather_into_tensor_uneven as all_gather_into_tensor_uneven	30	from .distributed_c10d import is_hccl_available, reinit_process_group, _reduce_scatter_tensor_uneven as reduce_scatter_tensor_uneven, _all_gather_into_tensor_uneven as all_gather_into_tensor_uneven
36	-
37	-rendezvous._rendezvous_init()

Mtorch_npu/distributed/rpc/backend_registry.py+9-6

	@@ -286,13 +286,16 @@ def _npu_tensorpipe_init_backend_handler(
286		286
287		287
288	def _rpc_backend_registry():	288	def _rpc_backend_registry():
289	- if hasattr(torch_npu._C, "_rpc_npu_init"):	289	+ if not hasattr(torch_npu._C, "_distributed_rpc"):
290	- torch_npu._C._rpc_npu_init()	290	+ raise RuntimeError(
291	- rpc.backend_registry.register_backend(	291	+ "torch_npu._C._distributed_rpc must be initialized before RPC backend registration"
292	- "NPU_TENSORPIPE",
293	- _npu_tensorpipe_construct_rpc_backend_options_handler,
294	- _npu_tensorpipe_init_backend_handler,
295	)	292	)
296		293
		294	+ rpc.backend_registry.register_backend(
		295	+ "NPU_TENSORPIPE",
		296	+ _npu_tensorpipe_construct_rpc_backend_options_handler,
		297	+ _npu_tensorpipe_init_backend_handler,
		298	+ )
		299	+
297	import torch.distributed.rpc as _rpc_module	300	import torch.distributed.rpc as _rpc_module
298	_rpc_module.BackendType = rpc.backend_registry.BackendType	301	_rpc_module.BackendType = rpc.backend_registry.BackendType

Mtorch_npu/dynamo/__init__.py+7-8

	@@ -6,9 +6,8 @@ import warnings
6	from torch._dynamo import register_backend as _register_backend	6	from torch._dynamo import register_backend as _register_backend
7	from torch._dynamo.backends.registry import _BACKENDS	7	from torch._dynamo.backends.registry import _BACKENDS
8		8
		9	+from torch_npu._init.common.warning_utils import _should_print_warning
9	from torch_npu.utils._error_code import ErrCode, pta_error	10	from torch_npu.utils._error_code import ErrCode, pta_error
10	-from torch_npu.utils.utils import _should_print_warning
11	-from .trace_rule import _patch_npu_trace_rules
12		11
13	_global_npu_backend = {}	12	_global_npu_backend = {}
14	__all__ = []	13	__all__ = []
	@@ -161,15 +160,15 @@ def _get_npugraph_ex_backend():
161	return _exec	160	return _exec
162		161
163		162
164	-_global_backend = _get_default_backend(name="npu")
165	-_npugraph_ex_backend = _get_npugraph_ex_backend()
166	-
167	-
168	def _register_npu_backend(backend, name="npu"):	163	def _register_npu_backend(backend, name="npu"):
169	if name in _BACKENDS.keys():	164	if name in _BACKENDS.keys():
170	del _BACKENDS[name]	165	del _BACKENDS[name]
171	_register_backend(backend, name)	166	_register_backend(backend, name)
172		167
173		168
174	-_register_npu_backend(_global_backend)	169	+def _register_backends():
175	-_register_npu_backend(_npugraph_ex_backend, NPUGRAPH_EX_BACKEND)	170	+ global_backend = _get_default_backend(name="npu")
		171	+ npugraph_ex_backend = _get_npugraph_ex_backend()
		172	+
		173	+ _register_npu_backend(global_backend)
		174	+ _register_npu_backend(npugraph_ex_backend, NPUGRAPH_EX_BACKEND)

Mtorch_npu/npu/__init__.py+1-4

	@@ -151,7 +151,7 @@ import re
151	import torch	151	import torch
152	from torch.storage import _LegacyStorage, _warn_typed_storage_removal	152	from torch.storage import _LegacyStorage, _warn_typed_storage_removal
153	from torch._utils import classproperty	153	from torch._utils import classproperty
154	-from torch_npu.utils import _should_print_warning	154	+from torch_npu._init.common.warning_utils import _should_print_warning
155		155
156	import torch_npu	156	import torch_npu
157	from torch_npu.utils._error_code import ErrCode, pta_error, prof_error	157	from torch_npu.utils._error_code import ErrCode, pta_error, prof_error
	@@ -190,9 +190,6 @@ from ._npugraph_handlers import (
190	register_npu_graph_handler,	190	register_npu_graph_handler,
191	)	191	)
192		192
193	-# init profiler
194	-if not torch_npu._C._profiler_init():
195	- raise RuntimeError("proflier initialization failed" + prof_error(ErrCode.UNAVAIL))
196		193
197	config = npu_config._npuConfig()	194	config = npu_config._npuConfig()
198		195

Mtorch_npu/npu/utils.py+1-1

	@@ -9,7 +9,7 @@ from torch._utils import _get_device_index as _torch_get_device_index
9		9
10	import torch_npu	10	import torch_npu
11	import torch_npu._C	11	import torch_npu._C
12	-from torch_npu.utils._error_code import ErrCode, pta_error, _except_handler	12	+from torch_npu.utils._error_code import ErrCode, pta_error
13	from torch_npu.npu._backends import get_soc_version	13	from torch_npu.npu._backends import get_soc_version
14		14
15		15

Mtorch_npu/profiler/__init__.py+0-3

	@@ -14,6 +14,3 @@ from ._non_intrusive_profile import _NonIntrusiveProfile
14	__all__ = ["profile", "ProfilerActivity", "supported_activities", "tensorboard_trace_handler", "schedule",	14	__all__ = ["profile", "ProfilerActivity", "supported_activities", "tensorboard_trace_handler", "schedule",
15	"ProfilerAction", "_ExperimentalConfig", "supported_profiler_level", "supported_ai_core_metrics",	15	"ProfilerAction", "_ExperimentalConfig", "supported_profiler_level", "supported_ai_core_metrics",
16	"supported_export_type", "ProfilerLevel", "AiCMetrics", "ExportType", "HostSystem"]	16	"supported_export_type", "ProfilerLevel", "AiCMetrics", "ExportType", "HostSystem"]
17	-
18	-
19	-_NonIntrusiveProfile.init()

Mtorch_npu/profiler/analysis/prof_common_func/_constant.py+1-2

	@@ -3,10 +3,9 @@ import os
3	import time	3	import time
4	from typing import Union	4	from typing import Union
5		5
6	-from torch_npu.utils import _should_print_warning	6	+from torch_npu._init.common.warning_utils import _should_print_warning
7	from torch_npu.utils._error_code import ErrCode, prof_error	7	from torch_npu.utils._error_code import ErrCode, prof_error
8		8
9	-
10	__all__ = []	9	__all__ = []
11		10
12		11

Mtorch_npu/utils/__init__.py+1-27

	@@ -9,43 +9,17 @@ __all__ = [
9	"get_cann_version",	9	"get_cann_version",
10	]	10	]
11		11
12	-from torch_npu import _C
13	from torch_npu.npu.utils import get_cann_version	12	from torch_npu.npu.utils import get_cann_version
14	-from torch_npu.utils._error_code import ErrCode, pta_error
15		13
16	-from ._dynamo import add_dynamo_methods
17	-from ._graph_tree import _apply_npugraph_tree_methods
18	-from ._inductor import _inductor_register_device_op_overrides
19	-from ._module import _apply_module_patch
20	-from ._npu_meta_registration import npu_patch_meta
21	-from ._optim import add_optim_method
22	-from ._step import add_perf_dump_patch
23	from .affinity import (	14	from .affinity import (
24	_reset_thread_affinity as reset_thread_affinity,	15	_reset_thread_affinity as reset_thread_affinity,
25	_set_thread_affinity as set_thread_affinity,	16	_set_thread_affinity as set_thread_affinity,
26	)	17	)
27	from .asd_detector import register_asd_hook, set_asd_loss_scale	18	from .asd_detector import register_asd_hook, set_asd_loss_scale
28	-from .collect_env import _add_collect_env_methods
29	from .combine_tensors import (	19	from .combine_tensors import (
30	get_part_combined_tensor,	20	get_part_combined_tensor,
31	is_combined_tensor_valid,	21	is_combined_tensor_valid,
32	npu_combine_tensors,	22	npu_combine_tensors,
33	)	23	)
34	-from .dlpack import _apply_dlpack_patch
35	from .flops_count import _FlopsCounter as FlopsCounter	24	from .flops_count import _FlopsCounter as FlopsCounter
36	-from .npu_intercept import _add_intercept_methods, _cann_package_check	25	+from .serialization import save_async
37	-from .serialization import _add_serialization_methods, save_async
38	-from .storage import _add_storage_methods
39	-from .tensor_methods import _add_tensor_methods
40	-from .utils import (
41	- _apply_npu_show_warning,
42	- _print_error_log,
43	- _print_info_log,
44	- _print_warn_log,
45	- _should_print_warning,
46	-)
47	-
48	-
49	-# init flopcount
50	-if not _C._flops_count_init():
51	- raise RuntimeError("flopcount initialization failed" + pta_error(ErrCode.UNAVAIL))

Mtorch_npu/utils/_afd_ops.py+12-5

	@@ -1,8 +1,15 @@
1	import torch	1	import torch
2	-import torch_npu	2	+
		3	+__all__ = ["initialize_afd_bindings"]
		4	+
		5	+_OP_NAMES = [
		6	+ "attention_worker_scheduler_",
		7	+ "attention_worker_scheduler",
		8	+ "ffn_worker_scheduler_",
		9	+ "ffn_worker_scheduler"]
3		10
4		11
5	-torch_npu._afd.attention_worker_scheduler_ = torch.ops.npu.attention_worker_scheduler_	12	+def initialize_afd_bindings():
6	-torch_npu._afd.attention_worker_scheduler = torch.ops.npu.attention_worker_scheduler	13	+ import torch_npu._afd as npu_afd
7	-torch_npu._afd.ffn_worker_scheduler_ = torch.ops.npu.ffn_worker_scheduler_	14	+ for name in _OP_NAMES:
8	-torch_npu._afd.ffn_worker_scheduler = torch.ops.npu.ffn_worker_scheduler	15	+ setattr(npu_afd, name, getattr(torch.ops.npu, name))

Mtorch_npu/utils/_dynamo_device.py+0-6

	@@ -82,9 +82,3 @@ class NpuInterface(DeviceInterface):
82	@staticmethod	82	@staticmethod
83	def is_bf16_supported(including_emulation: bool = False):	83	def is_bf16_supported(including_emulation: bool = False):
84	return True	84	return True
85	-
86	-
87	-def _dynamo_register_interface_for_device():
88	- register_interface_for_device("npu", NpuInterface)
89	- for i in range(32):
90	- register_interface_for_device(f"npu:{i}", NpuInterface)

Mtorch_npu/utils/hif8_tensor.py+0-3

	@@ -16,9 +16,6 @@ import torch_npu
16	from torch_npu.utils._error_code import ErrCode, pta_error	16	from torch_npu.utils._error_code import ErrCode, pta_error
17		17
18		18
19	-# init transformer engine
20	-torch_npu._C._cd_init()
21	-
22	tex = torch_npu._C._cd	19	tex = torch_npu._C._cd
23	aten = torch.ops.aten	20	aten = torch.ops.aten
24		21

Mtorch_npu/utils/serialization.py+1-2

	@@ -30,8 +30,7 @@ from torch.serialization import ( # noqa: F401
30		30
31	import torch_npu	31	import torch_npu
32	from torch_npu.utils._error_code import ErrCode, pta_error	32	from torch_npu.utils._error_code import ErrCode, pta_error
33	-	33	+from torch_npu._init.common.warning_utils import _should_print_warning
34	-from .utils import _should_print_warning
35		34
36		35
37	__all__ = ["load", "save_async"]	36	__all__ = ["load", "save_async"]

Mtorch_npu/utils/utils.py+1-10

	@@ -3,6 +3,7 @@ import sys
3	import time	3	import time
4	import warnings	4	import warnings
5	from warnings import _showwarnmsg_impl	5	from warnings import _showwarnmsg_impl
		6	+from torch_npu._init.common.warning_utils import _should_print_warning
6		7
7	__all__ = []	8	__all__ = []
8		9
	@@ -31,16 +32,6 @@ def _print_error_log(error_msg: str):
31	_print_log(_LogLevel.ERROR, error_msg)	32	_print_log(_LogLevel.ERROR, error_msg)
32		33
33		34
34	-def _should_print_warning():
35	- disabled_warning = os.environ.get("TORCH_NPU_DISABLED_WARNING", "0")
36	- if disabled_warning == "1":
37	- return False
38	- rank = os.environ.get("RANK", None)
39	- if rank is None or rank == "0":
40	- return True
41	- return False
42	-
43	-
44	def _apply_npu_show_warning():	35	def _apply_npu_show_warning():
45	def npu_show_warning(message, category, filename, lineno, file=None, line=None):	36	def npu_show_warning(message, category, filename, lineno, file=None, line=None):
46	npu_path = os.path.dirname(os.path.dirname(__file__))	37	npu_path = os.path.dirname(os.path.dirname(__file__))