自动加载树外扩展¶
创建日期:2024 年 10 月 10 日 | 最后更新:2024 年 10 月 10 日 | 最后验证:2024 年 10 月 10 日
作者: Yuanhao Ji
扩展自动加载机制使 PyTorch 能够自动加载树外后端扩展,而无需显式导入语句。此功能对用户有利,因为它增强了他们的体验,使他们无需显式加载或导入特定于设备的扩展即可遵循熟悉的 PyTorch 设备编程模型。此外,它还促进了在树外设备上轻松采用现有 PyTorch 应用程序,且无需进行代码更改。有关更多详细信息,请参阅 [RFC] 自动加载设备扩展。
如何在 PyTorch 中使用树外扩展自动加载
查看使用 Intel Gaudi HPU、华为昇腾 NPU 的示例
PyTorch v2.5 或更高版本
注意
此功能默认启用,可以通过使用 export TORCH_DEVICE_BACKEND_AUTOLOAD=0
禁用。如果您遇到类似“Failed to load the backend extension”(加载后端扩展失败)的错误,此错误与 PyTorch 无关,您应该禁用此功能并向树外扩展的维护者寻求帮助。
如何将此机制应用于树外扩展?¶
例如,假设您有一个名为 foo
的后端和一个名为 torch_foo
的对应包。确保您的包与 PyTorch 2.5 或更高版本兼容,并在其 __init__.py
文件中包含以下代码片段:
def _autoload():
print("Check things are working with `torch.foo.is_available()`.")
然后,您唯一需要做的是在您的 Python 包中定义一个入口点
setup(
name="torch_foo",
version="1.0",
entry_points={
"torch.backends": [
"torch_foo = torch_foo:_autoload",
],
}
)
现在,您只需添加 import torch
语句即可导入 torch_foo
模块,而无需添加 import torch_foo
>>> import torch
Check things are working with `torch.foo.is_available()`.
>>> torch.foo.is_available()
True
在某些情况下,您可能会遇到循环导入的问题。下面的示例展示了如何解决这些问题。
示例¶
在此示例中,我们将使用 Intel Gaudi HPU 和华为昇腾 NPU 来演示如何利用自动加载功能将树外扩展与 PyTorch 集成。
habana_frameworks.torch 是一个 Python 包,它通过使用 PyTorch HPU
设备键,使用户能够在 Intel Gaudi 上运行 PyTorch 程序。
habana_frameworks.torch
是 habana_frameworks
的一个子模块,我们在 habana_frameworks/setup.py
中添加了一个指向 __autoload()
的入口点
setup(
name="habana_frameworks",
version="2.5",
+ entry_points={
+ 'torch.backends': [
+ "device_backend = habana_frameworks:__autoload",
+ ],
+ }
)
在 habana_frameworks/init.py
中,我们使用一个全局变量来跟踪我们的模块是否已加载
import os
is_loaded = False # A member variable of habana_frameworks module to track if our module has been imported
def __autoload():
# This is an entrypoint for pytorch autoload mechanism
# If the following condition is true, that means our backend has already been loaded, either explicitly
# or by the autoload mechanism and importing it again should be skipped to avoid circular imports
global is_loaded
if is_loaded:
return
import habana_frameworks.torch
在 habana_frameworks/torch/init.py
中,我们通过更新全局变量的状态来防止循环导入
import os
# This is to prevent torch autoload mechanism from causing circular imports
import habana_frameworks
habana_frameworks.is_loaded = True
torch_npu 使用户能够在华为昇腾 NPU 上运行 PyTorch 程序,它利用了 PrivateUse1
设备键,并将设备名称 npu
暴露给终端用户。
我们在 torch_npu/setup.py 中定义了一个入口点
setup(
name="torch_npu",
version="2.5",
+ entry_points={
+ 'torch.backends': [
+ 'torch_npu = torch_npu:_autoload',
+ ],
+ }
)
与 habana_frameworks
不同,torch_npu
使用环境变量 TORCH_DEVICE_BACKEND_AUTOLOAD
来控制自动加载过程。例如,我们将其设置为 0
以禁用自动加载,从而防止循环导入
# Disable autoloading before running 'import torch'
os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'
import torch
工作原理¶

自动加载是基于 Python 的 Entrypoints 机制实现的。我们在 torch/__init__.py
中发现并加载所有由树外扩展定义的特定入口点。
如上所示,安装 torch_foo
后,当加载您定义的入口点时,您的 Python 模块即可被导入,然后您可以在调用它时执行一些必要的工作。
请参阅此拉取请求中的实现:[RFC] 添加对设备扩展自动加载的支持。
结论¶
在本教程中,我们了解了 PyTorch 中的树外扩展自动加载机制,该机制自动加载后端扩展,无需添加额外的导入语句。我们还学习了如何通过定义入口点将此机制应用于树外扩展,以及如何防止循环导入。此外,我们还回顾了如何使用自动加载机制与 Intel Gaudi HPU 和华为昇腾 NPU 的示例。