advanced/privateuseone

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

使用 PrivateUse1 实现新的后端集成

创建于：2023 年 10 月 03 日 | 最后更新：2024 年 05 月 07 日 | 最后验证：2024 年 11 月 05 日

在本教程中，我们将介绍将位于 pytorch/pytorch 仓库之外的新后端通过 PrivateUse1 集成到 PyTorch 所需的必要步骤。请注意，本教程假设您已具备 PyTorch 的基本知识。您是 PyTorch 的高级用户。

注意

本教程仅涉及与 PrivateUse1 机制相关的、用于集成新设备的部分，其他部分将不予涵盖。同时，本教程中并非所有涉及的模块都必需，您可以根据实际需求选择对您有帮助的模块。

什么是 PrivateUse1？

在 Pytorch 2.0 之前，PyTorch 提供了三个预留的调度键（及其对应的 Autograd 键）用于原型化 out-of-tree 后端扩展，这三个调度键如下：

• PrivateUse1/AutogradPrivateUse1

• PrivateUse2/AutogradPrivateUse2

• PrivateUse3/AutogradPrivateUse3

原型验证通过后，您可以为新后端申请私有键，例如 CUDA、XLA、MPS 等。

然而，随着 PyTorch 的飞速发展，越来越多的硬件厂商尝试将他们的后端集成到 PyTorch 中，这可能会导致以下问题：

• 每个新的后端集成都会涉及大量文件修改

• 目前调度键（DispatchKeySet）的数量存在硬限制（64 位限制）

注意

通过 PrivateUse1 键将新后端集成到 PyTorch 中也存在问题，因为不可能同时集成多个后端。幸运的是，这些 out-of-tree 后端很少同时使用。

鉴于以上原因，社区开始推荐新后端通过 PrivateUse1 集成到 PyTorch 中。

然而，先前的 PrivateUse1 机制在集成新后端方面能力并不完善，因为它在某些模块（如 Storage、AMP、Distributed 等）中缺乏相关的支持。

随着 Pytorch 2.1.0 的到来，对 PrivateUse1 在新后端集成方面进行了一系列优化和增强，现在可以快速高效地支持新设备的集成。

如何通过 PrivateUse1 集成新后端

在本节中，我们将讨论通过 PrivateUse1 将新后端集成到 Pytorch 中的详细信息，主要包括以下几个部分：

1. 为新后端注册算子。

2. 为新后端注册生成器。

3. 为新后端注册设备保护。

4. 为新后端元数据注册序列化和反序列化函数。

5. 其他模块。

为新后端注册算子

新后端可能有一些高性能的算子实现，可以通过 在 C++ 中注册一个调度的算子 中描述的 TORCH_LIBRARY_IMPL API 注册到调度器。这涉及几种情况：

1. 将新后端支持的所有前向算子注册到调度器，并同时注册回退，这样当新后端不支持某些算子时，这些算子可以回退到 CPU 执行，以确保功能的可用性。

at::Tensor wrapper_Custom_Tensor_add(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  // Implementation of add kernel in new backend
  ...
}

TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
  ...
  m.impl("add.Tensor", TORCH_FN(wrapper_Custom_Tensor_add));
  ...
}

void custom_cpu_fallback(const c10::OperatorHandle& op, torch::jit::Stack* stack) {
  // Add some hints about new devices that do not support and need to fall back to cpu
  at::native::cpu_fallback(op, stack);
}

TORCH_LIBRARY_IMPL(_, PrivateUse1, m) {
  m.fallback(torch::CppFunction::makeFromBoxedFunction<&custom_cpu_fallback>());
}

2. 通过 AutogradPrivateUse1 将 torch::autograd::Function 的算子注册到调度器，如果新后端需要覆盖 PyTorch Autograd layer，调度器和自动求导系统将自动调用这些算子的前向和后向实现。

class CumtomSeluFunction : public torch::autograd::Function<CumtomSeluFunction> {
  // Implementation of selu kernel in new backend
}

at::Tensor wrapper_AutogradCumstom__selu(const at::Tensor & self) {
  return CumtomSeluFunction::apply(self);
}

TORCH_LIBRARY_IMPL(aten, AutogradPrivateUse1, m) {
  ...
  m.impl("selu", TORCH_FN(wrapper_AutogradCustom__selu));
  ...
}

3. 通过 AutocastPrivateUse1 将希望支持 自动混合精度 (AMP) 和回退机制的算子注册到调度器，自动转换系统将在需要时自动调用这些算子。

TORCH_LIBRARY_IMPL(aten, AutocastPrivateUse1, m) {
  ...
  KERNEL_PRIVATEUSEONE(<operator>, <policy>)
  ...
}

TORCH_LIBRARY_IMPL(_, AutocastPrivateUse1, m) {
  m.fallback(torch::CppFunction::makeFallthrough());
}

需要补充的是，如果您希望在新后端中支持 AMP，需要通过 torch._register_device_module("backend_name", BackendModule) 注册一个新的 BackendModule，并且 BackendModule 需要具有以下 API：

• get_amp_supported_dtype() -> List[torch.dtype]

  获取 AMP 中新后端支持的 dtype，这可能会支持一个额外的 dtype。

• is_autocast_enabled() -> bool

  检查 AMP 在新后端是否已启用。

• get_autocast_dtype() -> torch.dtype

  获取 AMP 中新后端支持的 dtype，它由 set_autocast_dtype 设置或为默认 dtype，默认 dtype 为 torch.float16。

• set_autocast_enabled(bool) -> None

  启用或禁用新后端的 AMP。

• set_autocast_dtype(dtype) -> None

  设置新后端中 AMP 支持的 dtype，并且该 dtype 必须包含在从 get_amp_supported_dtype 获取的 dtypes 中。

为新后端注册生成器

需要支持对应新设备的生成器。目前，PrivateUse1 可以动态注册自定义生成器，主要分为以下几个步骤：

1. 继承 GeneratorImpl 类来实现对应新后端的生成器类，并实现各种通用方法。

2. 定义一个新的后端 builder，它有一个参数：device index。

3. 调用 REGISTER_GENERATOR_PRIVATEUSE1 宏以完成动态注册。

struct CustomGeneratorImpl : public c10::GeneratorImpl {
  // Implementation of generator in new backend
}

at::Generator make_custom_generator(c10::DeviceIndex device_index) {
  return at::make_generator<CustomGeneratorImpl>(device_index);
}

REGISTER_GENERATOR_PRIVATEUSE1(make_cumstom_generator)

为新后端注册设备保护

PyTorch 通过 DeviceGuard 提供与设备、流和事件切换相关的功能。此功能也适用于 PrivateUse1 键。

1. 继承 DeviceGuardImplInterface 类来实现对应新后端的各种通用方法。

2. 调用 C10_REGISTER_GUARD_IMPL 宏以完成动态注册。

struct CustomGuardImpl final : public c10::impl::DeviceGuardImplInterface {
  // Implementation of guard in new backend
}

C10_REGISTER_GUARD_IMPL(PrivateUse1, CustomGuardImpl);

为新后端元数据注册序列化和反序列化函数

PyTorch 目前能够动态注册序列化/反序列化函数，以支持在 TensorImpl.ExtraMeta 类中对名为 backend_meta_ 的新后端附加元数据的序列化和反序列化。您可以参考以下步骤：

1. 继承 BackendMeta 类来实现对应新后端的 CustomBackendMetadata，并且新后端的各个字段都可以在该类中进行自定义。

2. 实现新后端的序列化和反序列化函数，函数签名如下：void(const at::Tensor&, std::unordered_map<std::string, bool>&)。

3. 调用 TensorBackendMetaRegistry 宏以完成动态注册。

struct CustomBackendMetadata : public c10::BackendMeta {
  // Implementation of backend metadata in new backend
}

void for_serialization(const at::Tensor& t, std::unordered_map<std::string, bool>& m) {
  // Implementation of serialization
}

void for_deserialization(const at::Tensor& t, std::unordered_map<std::string, bool>& m) {
  // Implementation of deserialization
}

TensorBackendMetaRegistry(c10::DeviceType::PrivateUse1, &for_serialization, &for_deserialization);

其他模块

除了上述部分，还有一些其他模块可以通过 PrivateUse1 进行扩展，例如 distributed collective communication、benchmark timer 等，这些将在未来添加。一个关于 PrivateUse1 集成的示例是 Ascend NPU。

如何通过 Privateuse1 提升用户体验

通过 PrivateUse1 集成新设备的首要目标是满足基本功能需求，接下来的工作是提升可用性，这主要涉及以下几个方面：

1. 将新后端模块注册到 Pytorch。

2. 将 PrivateUse1 重命名为新后端的自定义名称。

3. 为新后端生成与新后端名称相关的方��和属性。

将新后端模块注册到 Pytorch

PyTorch 中一些与 CUDA 相关的接口可以通过 torch.cuda.xxx 的形式调用。因此，为了符合用户习惯，通过 PrivateUse1 机制实现的新后端也应该提供类似的接口。

例如，使用 Ascend NPU：

torch._register_device_module('npu', torch_npu.npu)

完成上述操作后，用户可以通过 torch.npu.xxx 调用 Ascend NPU 的一些独有 API。

将 PrivateUse1 重命名为新后端的自定义名称

PrivateUse1 键是集成到 PyTorch 中的新后端的内部机制。对于用户而言，与 PrivateUse1 相比，与新后端密切相关的自定义名称会更友好。

以 Ascend NPU 为例，首次使用会更友好。

torch.rand((2,2),device='npu:0')
torch.rand((2,2),device='privateuse1:0')

现在，PyTorch 为自命名的 PrivateUse1 后端提供了一个新的 C++/Python API，使用起来非常简单。

PYTHON

torch.rename_privateuse1_backend("npu")

C++

c10::register_privateuse1_backend("npu")

为新后端生成方法和属性

在将 PrivateUse1 重命名为自定义名称后，将为新后端自动生成与 Tensor、nn、Storage 模块中新后端名称相关的属性和方法。

以下是 Ascend NPU 的一个示例：

torch.rename_privateuse1_backend("npu")
unsupported_dtype = [torch.quint8]
torch.utils.generate_methods_for_privateuse1_backend(for_tensor=True, for_module=True, for_storage=True, unsupported_dtype=unsupported_dtype)

然后，您就可以使用以下方法和属性了：

torch.Tensor.npu()
torch.Tensor.is_npu
torch.Storage.npu()
torch.Storage.is_npu
...

未来工作

对 PrivateUse1 机制的改进仍在进行中，因此新模块的 PrivateUse1 集成方法将陆续添加。以下是我们正在积极进行的一些项目：

• 添加 distributed collective communication 的集成方法。

• 添加 benchmark timer 的集成方法。

结论

本教程向您介绍了通过 PrivateUse1 将新后端集成到 PyTorch 的过程，包括但不限于算子注册、生成器注册、设备保护注册等。同时，还介绍了一些提高用户体验的方法。