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 的過程，包括但不限於運算元註冊、生成器註冊、裝置保護註冊等。同時，還介紹了一些提高使用者體驗的方法。