自动加载树外扩展¶

创建于：2024年10月10日 | 最后更新：2024年10月10日 | 最后验证：2024年10月10日

扩展自动加载机制使PyTorch能够自动加载树外后端扩展，而无需显式的导入语句。此功能对用户有益，因为它增强了他们的体验，并使他们能够遵循熟悉的PyTorch设备编程模型，而无需显式加载或导入特定设备的扩展。此外，它还有助于轻松采用现有的PyTorch应用程序，无需对树外设备进行任何代码更改。有关更多详细信息，请参阅[RFC] Autoload Device Extension。

What you will learn

如何在PyTorch中使用树外扩展自动加载
使用英特尔Gaudi HPU、华为Ascend NPU的审查示例

Prerequisites

PyTorch v2.5 或更高版本

注意

此功能默认启用，可以通过使用 export TORCH_DEVICE_BACKEND_AUTOLOAD=0来禁用它。如果您遇到类似“无法加载后端扩展”的错误，此错误与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和华为Ascend 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] Add support for device extension autoloading。

结论¶

在本教程中，我们学习了PyTorch中的树外扩展自动加载机制，该机制自动加载后端扩展，无需添加额外的导入语句。我们还学习了如何通过定义入口点将此机制应用于树外扩展，以及如何防止循环导入。我们还回顾了一个示例，展示了如何将自动加载机制与Intel Gaudi HPU和华为Ascend NPU一起使用。