| refactor for torch_npu init module v2.
Co-authored-by: bellatan<tanmei2@huawei.com>
# message auto-generated for no-merge-commit merge:
!34638 merge v2.7.1_torch_npu_init_refactor into v2.7.1
refactor for torch_npu init module v2.
Created-by: bellatan
Commit-by: bellatan
Merged-by: ascend-robot
Description: <!--
PR描述模板更新日期：20260203
-->
# 【合入来源】
> <font color="red">**如有社区issue，请关联issue链接**</font>\
> <font color="red">**请勿携带内部流程信息（需求链接、问题单、内部issue等）**</font>
- [ ] 需求
- [ ] 问题单
- [ ] issue/工单
- [x] 重构优化
- [ ] 资料更新
# 【修改方案】
本 PR 对 torch_npu 初始化链路进行重构，将原先集中在 torch_npu/__init__.py 中的初始化逻辑拆分到 _init 目录下的多个职责模块中，形成“**顶层编排 + 子模块分阶段执行 + 内部能力统一收口**”的结构。重构后，torch_npu/__init__.py 不再承载大量具体业务初始化细节，只负责固定初始化时序。各类具体能力分别由 _check_device_conflict、_load_core_modules、_register_components、_apply_patches、_enable_optional_features、_initialize_runtime_lifecycle 等内部函数承接。
## 一、核心修改
### 1. 重构 torch_npu/__init__.py 顶层初始化入口
重构后的初始化流程如下：
```python
def _initialize():
_check_device_conflict()
_load_core_modules()
_register_components()
_apply_patches()
_enable_optional_features()
_initialize_runtime_lifecycle()
```
顶层入口主要负责：
1. 维护 __all__；
2. 在 import torch 前关闭 TORCH_DEVICE_BACKEND_AUTOLOAD，避免 PyTorch 后端自动加载导致循环依赖；
3. 提前导入 torch_npu.utils.patch_getenv，用于捕获初始化阶段的环境变量访问；
4. 按固定顺序调用各初始化阶段入口；
5. 保留 _autoload() 作为 PyTorch 后端自动加载入口，用于恢复 TORCH_DEVICE_BACKEND_AUTOLOAD。
---
### 2. 新增 _init 目录：
目录结构如下：
```text
torch_npu/_init/
__init__.py
common/
warning_utils.py
core/
_exports.py
module_loader.py
optional_features.py
runtime_lifecycle.py
patches/
__init__.py
api_patches.py
asd_patches.py
distributed_patches.py
dynamo_patches.py
monkey_patches.py
npu_patches.py
patch_manager.py
profiler_patches.py
warning_patches.py
registry/
__init__.py
backend.py
distributed.py
dynamo.py
registry_manager.py
```
---
### 3. 通过 _check_device_conflict() 处理前置设备冲突检查
**_check_device_conflict()**：负责最早期的设备冲突检查，避免 NPU 与其他 accelerator 同时启用。该接口属于初始化内部逻辑，不作为 public API 暴露。
---
### 4. 通过 _load_core_modules() 统一管理核心模块加载、注册副作用和顶层 API 导出
_load_core_modules() 将原先散落在 torch_npu/__init__.py 中的模块导入、底层 _C 子模块准备、基础 runtime 支撑模块初始化、导入即注册副作用以及顶层 API 导出统一收口。该阶段主要负责：
1. 加载 torch_npu 初始化所需的核心模块；
2. 统一准备 _C child submodules；
3. 初始化 logging、profiler、distributed 等基础组件；
4. 在 _C 准备完成后进行 torch_npu.npu 导入检查；
5. 加载需要通过 import 触发注册副作用的 Python 模块；
6. 导出 torch_npu 顶层 public API。
具体包括：
* _C 子模块初始化：统一创建并注册 _profiler、_distributed_c10d、_cd、_logging、_flops_count 等 _C child submodules，保证业务 Python 模块只消费这些子模块，不再各自创建。
* torch_npu.npu 导入检查：在 _C 子模块完成准备后再检查 torch_npu.npu 导入状态，既保留对底层依赖缺失的友好报错，又避免 _C 未就绪时提前 import torch_npu.npu 导致循环导入。
* 导入副作用模块加载：统一加载需要通过 import 触发注册副作用的模块，例如 aclnn、optim、afd、custom ops、op_plugin、meta registrations 等，避免注册类副作用散落在初始化流程中。
* 顶层 API 导出：通过 export_all 将 torch_npu 顶层公开 API 统一导出到 globals() 和 __all__ 中，保证 public API 行为与旧版兼容。
* lazy Python API：对 HiFloat8Tensor、erase_stream、matmul_checksum 等接口采用 lazy export，保证接口可见但不在 import 阶段立即加载对应模块，减少循环导入风险。
* NPU custom ops：将 torch.ops.npu 下的公开算子导出到 torch_npu 顶层，并保留 torch.<op> deprecated wrapper。
* dtype symbols：将 _C._cd.DType 中的 dtype 符号导出到 torch_npu 顶层。
---
### 5. 通过 _register_components() 统一管理框架集成注册
_register_components() 负责 backend 和 framework integration 注册，将原先散落在顶层入口中的 NPU backend、distributed、Dynamo、RPC、Inductor 等注册逻辑统一收口。通过该阶段统一收口后，框架集成注册逻辑不再散落在顶层 __init__.py 中，后续新增集成能力时可直接在 registry 目录下维护。该阶段主要负责：
* NPU backend 注册：将 PyTorch PrivateUse1 backend 映射为 NPU，并注册 torch.npu 设备模块和相关方法。
* distributed backend 注册：注册 HCCL、LCCL backend，保证 NPU distributed 能力可用。
* Dynamo 注册：注册 Dynamo backend、NPU device interface 和 trace rules，保证 NPU 能接入 Dynamo 编译链路。
* RPC 注册：注册 NPU RPC backend，保证 RPC 场景下 NPU backend 可用。
* Inductor lightweight override 注册：只注册轻量级 NPU device op override，避免 import 阶段提前加载 heavy module。
* 默认 gradient device type 配置：保持 checkpoint 等场景下默认设备类型与旧行为兼容。
---
### 6. 通过 _apply_patches() 统一管理 patch 注册与执行
引入集中式 patch 管理机制，统一收口原先散落在初始化入口中的 patch 逻辑。顶层入口 _apply_patches() 负责触发 patch 发现、注册和执行，具体由 PatchManager 承接。_apply_patches() 主要完成以下工作：
1. **patch 分组注册**：各组件 patch 按 group 注册，例如 monkey、api、distributed、dynamo、profiler、npu、warning、asd 等。
2. **内置 patch 自动发现**：PatchManager 会自动扫描 _init/patches 下符合命名规则的 patch 模块。模块被导入后，内部 patch 会完成注册。
3. **固定 patch 执行顺序**：patch group 按默认顺序执行，避免 import 顺序变化导致 patch 行为漂移。
4. **支持自定义 patch 顺序**：PatchManager 支持调整 patch group 执行顺序，便于测试或特殊场景扩展。
5. **异常钩子统一处理**：全局异常钩子由 PatchManager.run() 统一处理，便于初始化失败和运行时异常场景的集中管理。
---
### 7. 通过 _enable_optional_features() 统一管理可选运行时能力
将 sanitizer、交互式模式配置、transfer_to_npu 等可选能力统一收口到 _enable_optional_features()，避免可选逻辑散落在顶层初始化入口。该阶段主要包括：
```python
_enable_sanitizer_if_needed()
_configure_interactive_mode()
_enable_transfer_to_npu_if_needed()
```
具体说明：
* sanitizer：仅在用户显式配置 TORCH_NPU_SANITIZER 时启用；
* interactive mode：在交互式命令行环境中自动设置相关运行配置，并给出 warning 提示；
* transfer_to_npu：通过 TORCH_TRANSFER_TO_NPU 控制是否启用，对非法配置进行显式报错。
---
### 8. 通过 _initialize_runtime_lifecycle() 统一管理 runtime 生命周期
_initialize_runtime_lifecycle() 专门负责最终 C++ extension 初始化屏障和进程退出阶段的 shutdown hook 注册。该阶段主要包括：
* extension finalize：调用 torch_npu._C._initExtension() 完成最终 C++ extension 绑定。该阶段放在核心模块加载、框架注册、API 导出和 patch 执行之后，保证 Python 侧初始化准备完成后再进入最终 extension barrier。
* shutdown hook 注册：负责注册进程退出阶段的 NPU 资源清理逻辑，包括设备同步、distributed 资源析构、异常处理和其他 runtime 清理流程。
---
## 三、重构目的和收益
本次重构的目标是把 torch_npu 初始化从“单文件集中式副作用堆叠”调整为“阶段化、组件化、可维护”的初始化框架。主要收益包括：
1. **顶层入口更清晰**
torch_npu/__init__.py 只保留初始化编排，不再堆叠大量具体 import、注册、patch 和 shutdown 逻辑。
2. **初始化顺序更稳定**
_C 子模块和基础 runtime 支撑能力统一由 _load_core_modules 准备，降低循环导入和 _C 未就绪时提前访问的风险。
3. **组件职责更清楚**
模块加载、框架注册、API 导出、patch、可选功能、runtime 生命周期分别由不同接口承接。
4. **patch 更易维护**
各组件 patch 可以在自己的文件中维护，由 PatchManager 自动发现和统一执行，减少顶层冲突。
5. **支持后续扩展**
新增初始化能力时，只需放到对应处理的接口 或 patch group 中，不需要继续膨胀 __init__.py。
6. **便于问题定位**
初始化链路被拆成明确阶段，出现问题时可以快速判断是模块加载、注册、导出、patch、optional feature 还是 runtime lifecycle 阶段异常。
---
## 四、兼容性说明
本次重构保持以下兼容性：
1. import torch_npu 行为保持兼容；
2. 顶层公开 API 保持兼容；
3. __version__ 仍从 torch_npu.version 导出。
---
## 五、PatchManager 机制说明
本 PR 引入 PatchManager，用于统一管理 torch_npu 初始化阶段的 patch 注册与执行。原先 patch 逻辑集中在 torch_npu/__init__.py 中，和初始化流程、模块导入、框架注册逻辑混在一起，导致顶层文件过重，也不利于各组件独立维护。本次重构后，patch 逻辑从顶层入口中解耦，由 _apply_patches() 作为顶层入口触发执行，具体注册、发现、排序、幂等保护由 PatchManager 管理。
PatchManager 主要支持以下能力：
1. patch 按 group 分组注册；
2. 内置 patch 模块自动发现；
3. patch 按固定顺序执行；
4. patch 执行具备幂等保护；
5. 支持组件自行维护 patch module；
6. 支持按 group 执行，为后续按需使能 patch 打基础；
7. 支持自定义 patch 顺序，便于测试和问题定位。
整体机制如下：
```text
组件 patch 文件自注册
↓
PatchManager 自动发现/加载
↓
按 group 统一管理
↓
按固定顺序执行
↓
幂等保护，避免重复 patch
```
---
### 场景一：新增 torch_npu 内置 patch
如果新增的是 torch_npu 内置 patch，例如 distributed patch、profiler patch、NPU API patch、warning patch、ASD patch 等，可以直接放到：torch_npu/_init/patches/ 目录下，并按 group 注册。
示例：
```python
from torch_npu._init.patches.patch_manager import PatchManager
@PatchManager.register_patch("profiler")
def apply_profiler_patch():
...
```
使用方式：
```text
1. 在 _init/patches 下新增或修改对应 *_patches.py 文件；
2. 在文件中通过 @PatchManager.register_patch(group) 注册 patch；
3. import torch_npu 时，由 _apply_patches() 统一触发；
4. PatchManager 自动发现并按 group 顺序执行。
```
---
### 场景二：组件自行维护 patch module
如果某个组件自己的目录下新加了patch 文件，通过 patch module 注册机制接入。
示例：
```python
PatchManager.register_patch_module("torch_npu.some_component.some_patches")
```
组件自己的 patch 文件中仍然使用 group 注册：
```python
from torch_npu._init.patches.patch_manager import PatchManager
@PatchManager.register_patch("some_component")
def apply_some_component_patch():
...
```
使用方式：
```text
1. 组件在自己的目录中维护 patch 文件；
2. 通过 register_patch_module 注册该 patch module；
3. module 被导入后，内部 patch 自动注册到 PatchManager；
4. 后续仍由 PatchManager 统一排序和执行。
```
适用场景：
```text
组件有独立维护边界；
patch 逻辑不适合放到中心化 patches 目录；
后续组件可能独立演进、迁移或删除。
```
---
### 场景三：按 group 执行 patch，用于测试或后续按需使能
PatchManager 支持按 group 执行 patch。当前默认初始化路径仍执行全部注册 patch，后续也可按需使能。
示例：
```python
PatchManager.apply_registered_patches("distributed")
```
使用方式：
```text
1. 指定需要执行的 patch group；
2. PatchManager 只执行该 group 下已注册的 patch；
3. 已执行过的 patch 不会重复执行；
4. 可用于单独验证某一类 patch 的行为。
```
适用场景：
```text
只验证 distributed patch；
只执行 profiler patch；
排查某一类 patch 对初始化流程的影响；
后续通过环境变量控制某个 patch group 是否启用。
```
如果需要调整 patch group 顺序，也可以使用：
```python
PatchManager.set_patch_order([
"monkey",
"api",
"distributed",
])
```
适用场景：
```text
测试 patch 顺序；
排查 patch 依赖问题；
特殊构建或实验场景调整 patch 执行顺序。
```
# 【资料变更】
> 不涉及
# 【接口变更】
> 不涉及
# 【功能验证】
新增 TestTorchNpuBootstrap 初始化专项测试，覆盖以下场景：
1. test_01_import_order_compatibility
验证 import torch_npu、import torch; import torch_npu、import torch_npu; import torch、重复 import torch_npu 等不同导入顺序保持兼容。
2. test_02_import_state_snapshot
验证 import torch_npu 后的初始化状态，包括 torch.npu 注册、Tensor/Module.npu 方法生成、_C child submodules 准备、旧版初始化副作用模块加载、非预期模块不 eager import、顶层关键属性可访问等。
3. test_03_public_exports_snapshot
验证顶层 public API 导出行为，包括 lazy Python APIs、torch.ops.npu public ops、deprecated torch.<op> alias、dtype symbols 等导出保持兼容。
4. test_04_framework_registration_snapshot
验证框架集成注册行为，包括 Dynamo NPU device interface、Dynamo backend、Inductor lightweight device op override、distributed backend、RPC backend 等注册保持生效。
5. test_05_runtime_lazy_init_semantics
验证 import 阶段不触发 NPU runtime lazy init，查询类 API 不触发完整 runtime 初始化，真实 runtime API 和显式 torch_npu.npu.init() 能正常触发 lazy init。
6. test_06_component_behavior_snapshot
验证关键组件行为保持兼容，包括 patch_getenv 生效、ASD detector 兼容 API、AFD 通过 torch_npu._afd 暴露、torch_npu._C._afd 不暴露、AFD ops 可访问等。
7. test_07_distributed_patch_behavior
验证 distributed patch 行为保持兼容，包括 distributed 内部函数替换、public API alias、rendezvous/launcher patch、FSDP 相关 patch 等。
新增测试用例本地验证通过。
# 【CheckList】
> PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]
- [x] 代码注释完备，正确记录错误日志
- [x] 代码实现进行了返回值、空指针等校验
- [x] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
- [x] PR持续集成流水线（CI）执行通过，代码检查无异常
See merge request: Ascend/pytorch!34638 | 19 天前 |
