透過 PrivateUse1 促進新的後端整合

建立於：2023 年 10 月 03 日 | 最後更新：2024 年 05 月 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. 其他模組。

為新的後端註冊核心

新的後端可能具有運算符的一些高效能實作，可以透過 在 C++ 中註冊分發的運算符中描述的 TORCH_LIBRARY_IMPL API 註冊到分發器。這涉及以下幾種情況：

1. 將新後端支援的所有前向運算子註冊到分發器 (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. 透過 AutogradPrivateUse1 將來自 torch::autograd::Function 的核心 (kernel) 註冊到分發器。如果新的後端需要覆寫 PyTorch Autograd layer，分發器和 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 中新後端支援的 dtypes，它可能支援一個或多個 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)

為新的後端註冊裝置保護器 (device guard)

PyTorch 透過 DeviceGuard 提供與裝置、流 (stream) 和事件切換相關的功能。此功能也適用於 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 目前能夠動態註冊序列化/反序列化函數，以支援在類別 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 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 的過程，包括但不限於運算元註冊、產生器註冊、裝置守護註冊等等。同時，還介紹了一些改善使用者體驗的方法。