通过 PrivateUse1 促进新的后端集成

创建时间： Oct 03， 2023 |上次更新时间：2024 年 5 月 7 日 |上次验证： Nov 05， 2024

在本教程中，我们将演练一些必要的步骤来集成新的后端 住在 repo 之外 由 .请注意，本教程假定 您已经对 PyTorch 有基本的了解。 您是 PyTorch 的高级用户。pytorch/pytorchPrivateUse1

注意

本教程仅涉及与促进新设备集成的 PrivateUse1 机制相关的部分， 和其他部分将不包括在内。同时，并非本教程中涉及的所有模块都是必需的， 并且您可以根据自己的实际需求选择对您有帮助的模块。

什么是 PrivateUse1？

在 Pytorch 2.0 之前，PyTorch 提供了三个保留的调度密钥（及其相应的 Autograd 密钥） 对于树外后端扩展的原型设计，三个 Dispatch Key 如下：

• PrivateUse1/AutogradPrivateUse1

• PrivateUse2/AutogradPrivateUse2

• PrivateUse3/AutogradPrivateUse3

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

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

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

• 目前对 Dispatch Key 的数量有硬性限制（64 位限制）DispatchKeySet

注意

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

鉴于以上原因，社区开始推荐新的后端进行集成 通过 .PrivateUse1

但是，以前的机制并不完全能够与新的后端集成，因为它 在某些模块中缺少一些相关支持，例如 Storage、AMP、Distributed 等。PrivateUse1

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

如何通过 PrivateUse1 集成新后端

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

1. 为新后端注册内核。

2. 为新后端注册 generator。

3. 为新后端注册 device guard。

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

5. 其他模块。

为新后端注册内核

新后端可能具有一些高性能的 operator 实现，这些实现可以注册到调度程序 由 API 在 C++ 中注册 Dispatched Operator 中所述。这包括 几种情况：TORCH_LIBRARY_IMPL

1. 将新后端支持的所有 forward 运算符注册到 Dispatcher，并注册 fallback 同时，当新的后端不支持某些算子时，这些算子可以回退 到 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. 通过 将内核从 注册到 Dispatcher （如果需要） new backend to override ，Dispatcher 和 autograd 系统将自动调用 forward 和 这些运算符的向后实现。torch::autograd::FunctionAutogradPrivateUse1PyTorch 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） 的内核，并且 回退机制传递给 Dispatcher，则自动转换系统将在需要时自动调用这些内核。AutocastPrivateUse1

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

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

需要补充的是，如果要在新的后端支持 AMP，则需要注册一个新的 BY ，并且需要具有以下 API：BackendModuletorch._register_device_module("backend_name", BackendModule)BackendModule

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

  在 AMP 中的新后端上获取受支持的 dtypes，这可能支持另外一个 .dtype

• is_autocast_enabled() -> bool

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

• get_autocast_dtype() -> torch.dtype

  在 AMP 中的新后端上获取受支持，该后端由 或 default ，默认值为 。dtypeset_autocast_dtypedtypedtypetorch.float16

• set_autocast_enabled(bool) -> None

  在新后端启用或禁用 AMP。

• set_autocast_dtype(dtype) -> None

  在 AMP 中的新后端上设置 supported，并将 be 包含在 GOT 中 从。dtypedtypedtypesget_amp_supported_dtype

为新后端注册生成器

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

1. 继承类实现新后端对应的 generator 类， 并实施各种通用方法。GeneratorImpl

2. 使用单个参数定义新的后端：。builderdevice index

3. 调用 macro 完成动态注册。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)

为新后端注册 device guard

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

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

2. 调用 macro 完成动态注册。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. 继承 class 实现对应的新后端和 可以在 class 中自定义 new backend 的各种字段。BackendMetaCustomBackendMetadata

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);

其他模块

除了上述部分外，还有一些其他模块可以通过 进行扩展 例如 、 、 等，这些函数将在将来添加。 关于集成的一个例子是 Ascend NPU。PrivateUse1distributed collective communicationbenchmark timerPrivateUse1

如何通过 Privateuse1 改善用户体验

集成新设备的主要目标是满足基本功能需求， 而接下来要做的就是提升易用性，主要涉及以下几个方面。PrivateUse1

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

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

3. 生成与新后端相关的方法和属性。

将新的后端模块注册到 Pytorch

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

例如，使用 ：Ascend NPU

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

完成以上操作后，用户可以通过Ascend NPUtorch.npu.xxx

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

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

以 为例，第一次使用将更加用户友好。Ascend NPU

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

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

蟒

torch.rename_privateuse1_backend("npu")

C++

c10::register_privateuse1_backend("npu")

生成与新后端相关的方法和属性

重命名为自定义名称后，自动生成与新后端名称相关的属性和方法 在新后端的模块中。PrivateUse1Tensor, 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
...

未来的工作

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

• 新增 的集成方式。distributed collective communication

• 新增 的集成方式。benchmark timer

结论

本教程向您介绍了通过以下方式将新后端集成到 PyTorch 中的过程，包括但不限于 操作员注册、生成器注册、Device Guard 注册等。同时，介绍了一些方法 以改善用户体验。PrivateUse1