advanced/privateuseone

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

通过PrivateUse1促进新的后端集成

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

在本教程中，我们将介绍通过PrivateUse1集成pytorch/pytorch仓库外部的新后端所需的一些步骤。请注意，本教程假定您已经具备对PyTorch的基本了解。您是PyTorch的高级用户。

注意

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

什么是PrivateUse1？

在Pytorch 2.0之前，PyTorch提供了三个保留的分派键（及其对应的Autograd键），用于原型化树外后端扩展，这三个分派键如下

• PrivateUse1/AutogradPrivateUse1

• PrivateUse2/AutogradPrivateUse2

• PrivateUse3/AutogradPrivateUse3

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

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

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

• 分派键 (DispatchKeySet 64位限制) 的数量存在硬性限制

注意

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

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

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

随着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. 通过AutogradPrivateUse1从torch::autograd::Function注册内核，如果新的后端需要覆盖PyTorch Autograd 层，则分派器和autograd系统将自动调用这些算子的前向和反向实现。

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)和回退机制的内核，autocast系统将在需要时自动调用这些内核。

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，该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 通信、benchmark 计时器等，这些将在未来添加。一个关于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 提供了一个新的 C++/Python API，用于自命名的 PrivateUse1 后端，使用起来非常简单。

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 集成方法将陆续添加。以下是我们正在积极进行的一些项目

• 添加 分布式集合通信 的集成方法。

• 添加 基准计时器 的集成方法。

结论

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