通过PrivateUse1促进新后端集成

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

在这个教程中，我们将逐步介绍一些必要的步骤，以集成一个位于 pytorch/pytorch 仓库之外的后端 由 PrivateUse1 提供。请注意，本教程假设您已经具备 PyTorch 的基本知识。 您是 PyTorch 的高级用户。

注意

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

什么是PrivateUse1?

在 PyTorch 2.0 之前，PyTorch 提供了三个保留的 dispatch keys（以及对应的 Autograd keys）用于原型化外部后端扩展，这三个 dispatch keys 如下：

• PrivateUse1/AutogradPrivateUse1

• PrivateUse2/AutogradPrivateUse2

• PrivateUse3/AutogradPrivateUse3

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

然而，随着PyTorch的快速发展，越来越多的硬件制造商正在尝试将其后端集成到PyTorch中，这可能会导致以下问题：

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

• 目前对 Dispatch Keys 的数量存在硬性限制 (DispatchKeySet 64位限制)

注意

通过 PrivateUse1 Key 将新后端集成到 PyTorch 中也存在一个问题，因为无法同时集成许多后端。幸运的是，这些外部后端很少会同时被使用。

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

然而，之前的 PrivateUse1 机制无法完全与新后端集成，因为它在某些模块中缺乏相关支持，例如存储、AMP、分布式等。

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

如何通过PrivateUse1集成新后端

在本节中，我们将讨论通过 PrivateUse1 将新后端集成到Pytorch的细节， 这主要由以下部分组成：

1. 为新的后端注册内核。

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

3. 为新后端注册设备守护进程。

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

5. 其他模块。

为新后端注册内核

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

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. 从 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. 注册希望支持 自动混合精度（AMP） 和 通过 AutocastPrivateUse1 的分发器回退机制的内核，当需要时，自动转换系统将自动调用这些内核。

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

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

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

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

  获取AMP中新后端支持的dtypes，可能会支持一个额外的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

  设置支持的 dtype 在AMP的新后端中，并且 dtype 应包含在 dtypes 中，该 get_amp_supported_dtype 获得

新后端的注册生成器

需要支持对应新设备的生成器。目前，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 提供与设备、流和事件切换相关的功能 via DeviceGuard. 此功能也适用于 PrivateUse1 Key.

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目前能够动态注册序列化/反序列化函数以支持新后端附加元数据backend_meta_在类TensorImpl.ExtraMeta中的序列化和反序列化。您可以参考以下步骤：

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 Key 是新后端集成到 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 的过程，包括但不限于运算符注册、生成器注册、设备守护注册等。同时，还介绍了一些方法以提高用户体验。