引言
随着对多样化硬件加速器需求的增长，对强大且适应性强的深度学习框架的需求变得越来越重要。在进行这种集成的过程中，PyTorch 生态系统中浮现出一些挑战，可能会影响到各种硬件供应商。本博客旨在强调这些问题，并提出解决方案，以增强 PyTorch 在不同硬件平台上的适应性、可移植性和弹性。
通过加速器自动加载提高用户的代码可移植性
目前，当用户在不同的加速器上运行代码时，需要进行额外的工作。其中一项任务是手动导入树外设备的模块。这要求用户不仅要理解不同加速器之间的不同使用模式，还要使他们的代码意识到这些差异。如果您的项目最初在 GPU/CPU 上运行，并想迁移到其他加速器，这可能会导致大量的工作和潜在的挫败感。
额外的导入示例
# Case 1: Use HPU
import torch
import torchvision.models as models
import habana_frameworks.torch # <-- extra import
model = models.resnet50().eval().to("hpu")
input = torch.rand(128, 3, 224, 224).to("hpu")
output = model(input)
# Case 2: Use torch_npu
import torch
import torch_npu # <-- extra import
print(torch.ones(1, 2, device='npu'))
作为一个高水平的机器学习框架，PyTorch 屏蔽用户免受设备差异影响的能力是一项竞争优势。 加速器自动加载 允许用户继续使用熟悉的 PyTorch 设备编程模型，而无需显式加载或导入特定于设备的扩展。
它是如何工作的？
利用 Python 的插件架构，通过 PyTorch 包中的入口点实现设备扩展的自动加载。
Python 入口点为 Python 包提供了一种标准化的方式，用于公开和发现应用程序中的组件或插件。通过在加速器包 setup.py 中的定义，PyTorch 可以在调用 import torch 时自动初始化加速器模块，从而为用户提供不同后端设备之间一致的体验。
从设备角度来看，只需要在 setup.py 中声明以下设置（以 torch_npu 为例）
// setup.py
entry_points={
'torch.backends': ['torch_npu = torch_npu:_autoload', ],
}
当调用 import torch 时，加速器模块将自动加载。这为用户提供了跨树外设备的一致编程体验，消除了了解 CUDA、HPU 和 NPU 之间差异的需求。
# Case 1: Use HPU
import torch
import torchvision.models as models
model = models.resnet50().eval().to("hpu")
input = torch.rand(128, 3, 224, 224).to("hpu")
output = model(input)
# Case 2: Use torch_npu
import torch
print(torch.ones(1, 2, device='npu'))
设备集成优化
什么是 PrivateUse1？
在 PyTorch 中，dispatcher（调度器）是框架后端的一个关键组件，它管理操作如何路由到适当的设备特定实现。Dispatch keys（调度键）是该系统不可或缺的一部分，充当标识符，表示各种执行上下文，例如设备（CPU、CUDA、XPU）、布局（密集、稀疏）和自动微分功能。这些键确保操作被定向到正确的实现。
PrivateUse1 是一个可定制的设备调度键，类似于 CUDA/CPU/XPU 等，专为树外设备保留。它为开发者提供了一种扩展 PyTorch 功能的方式，而无需修改核心框架，从而可以集成新的设备、硬件加速器或其他专门的计算环境。
为什么我们需要 PrivateUse1？
在内部，调度键表示为位掩码，每个位代表某个键是否处于活动状态。这种位掩码表示对于键的快速查找和组合非常有效，但它固有地限制了不同键的数量（通常为 64 个或更少）。
PyTorch 中 BackendComponent 调度键的当前实现遇到了一个关键瓶颈，这限制了新后端的添加，并因此限制了 PyTorch 生态系统的扩展。
为了应对这一挑战，已对 PrivateUse1 机制应用了一系列优化，以增强其容量。
-
PrivateUse1 集成机制
最初保留作为回退选项，PrivateUse1 以及 PrivateUse2 和 PrivateUse3，旨在仅在现有键资源变得稀缺时才激活。
PrivateUse1 目前正在开发中，以匹配 CUDA 和 CPU 等已建立键的稳健性和多功能性。实现这一点需要在关键的 PyTorch 模块之间进行深度集成。这种集成不仅仅是一个简单的切换，它还涉及到对核心组件的重大更新，例如 AMP (自动混合精度)、Autograd、分布式训练、Checkpointing、DataLoader、Optimization 和 Quantization 等。
PrivateUse1 的激活是一项大规模的协作努力，最终完成了 100 多个 pull request，旨在使其从占位符变成完全可操作的调度键。
-
PrivateUse1 UT/CI 质量保证
虽然单元测试对于确保 PrivateUse1 机制开发过程中的质量至关重要，但仅靠它们不足以防止新的 pull request 无意中影响树外设备的现有功能或兼容性。
为了降低这种风险，社区已将
pytorch_openreg模块添加到测试套件中。该模块利用 CPU 后端来模拟与加速器的交互，从而为严格的测试创建受控环境。实施后，这将使设备通用测试用例能够在每次更新相关代码时自动执行，从而使我们能够快速检测和解决影响 PrivateUse1 集成机制的任何潜在问题。 -
全面的文档
通过提供全面且易于理解的文档，我们的目标是降低开发者的入门门槛，并鼓励在 PyTorch 生态系统中更广泛地采用 PrivateUse1 机制。该文档包括
- 使用 PrivateUse1 集成新后端的逐步指南
- PrivateUse1 功能和优势的清晰解释
- 高效实现的代码示例和最佳实践
这些增强功能旨在提高 PrivateUse1 机制的稳健性和可靠性，从而促进新后端的更好集成并扩展 PyTorch 的功能。
上游和下游之间的兼容性
设备通用单元测试
PyTorch 中的大多数单元测试都侧重于 CPU 和 CUDA 设备，这限制了来自其他硬件用户的参与。为了解决这个问题，我们计划修改 PyTorch 的单元测试框架，以便更好地支持非 CUDA 设备。该计划包括删除现有的设备限制、实现动态数据类型加载以及通用化装饰器以适应更广泛的设备。此外，我们的目标是强制使用通用设备代码，并扩展分布式测试以支持非 NCCL 后端。
通过这些改进，我们希望显着提高非 CUDA 设备的测试覆盖率和通过率，并将它们集成到 PyTorch 的持续集成过程中。最初的更改已经实施，为新的硬件支持铺平了道路，并为其他设备创建了参考模板。
通过自动化测试确保强大的设备集成
为了维护 PyTorch 高标准的质量保证，我们建立了一个独立的构建仓库和每日持续集成 (CI) 工作流程，专注于冒烟测试和集成测试。
pytorch-integration-tests 仓库自动化了 PyTorch 设备特定功能的测试，确保它们在各种硬件平台（NPU 和其他专用设备）上正确高效地运行。在仓库中，我们正在尝试构建一个完全自动化的系统，持续验证 PyTorch 与不同硬件后端的兼容性。
- 自动化集成测试：使用 GitHub Actions 在不同设备上运行自动化测试。这种自动化确保代码库中的每次更改都针对多个硬件平台进行全面测试，从而在开发过程的早期发现潜在问题。
- 可重用的工作流程：此仓库中的工作流程是模块化和可重用的，这简化了测试过程。开发者可以轻松地将这些工作流程调整到新的设备或测试场景，使系统在 PyTorch 不断发展的同时保持灵活性和可扩展性。
- 树外设备的感知：该仓库显示所有树外设备的存在和行为，使社区了解情况。这种方法最大限度地降低了意外破坏下游功能的风险，并为更改提供快速反馈。
增强多设备集成的努力对于其在不断发展的深度学习领域中的适应性至关重要。这些举措不仅使当前用户受益，还降低了新硬件供应商和开发者的入门门槛，从而促进了 AI 和机器学习领域的创新。随着 PyTorch 的不断发展，其对灵活性、稳健性和包容性的承诺使其成为能够满足深度学习社区多样化需求的领先框架。
