文件最后提交记录最后更新时间
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_alltorch_npu 顶层公开 API 统一导出到 globals()__all__ 中,保证 public API 行为与旧版兼容。 * lazy Python API:对 HiFloat8Tensorerase_streammatmul_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_npuimport torch; import torch_npuimport 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!3463820 天前