引言
随着对各种硬件加速器的需求不断增长，对强大且适应性强的深度学习框架的需求也日益迫切。在处理这一集成过程中，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 中，调度器是框架后端的一个关键组件，负责管理操作如何路由到适当的设备特定实现。调度键是这个系统的一个组成部分，作为标识符，代表各种执行上下文——如设备（CPU、CUDA、XPU）、布局（稠密、稀疏）和自动微分功能。这些键确保操作被导向正确的实现。
PrivateUse1 是一个可定制的设备调度键，类似于 CUDA/CPU/XPU 等），为树外设备保留。它为开发者提供了一种扩展 PyTorch 功能的方法，而无需修改核心框架，允许集成新的设备、硬件加速器或其他专用计算环境。
为什么我们需要 PrivateUse1？
内部，派发键以位掩码的形式表示，每个位表示某个键是否激活。这种位掩码表示方式在快速查找和组合键方面效率很高，但本质上限制了不同键的数量（通常为 64 个或更少）。
PyTorch 中 BackendComponent 派发键的当前实现遇到了一个关键瓶颈，这限制了新后端的添加，从而限制了 PyTorch 生态系统的扩展。
针对这一挑战，对 PrivateUse1 机制进行了一系列优化，以提高其容量。
-
PrivateUse1 集成机制
最初，PrivateUse1、PrivateUse2 和 PrivateUse3 被预留为后备选项，旨在仅在现有关键资源变得稀缺时才激活。
PrivateUse1 现在正在开发中，以匹配 CUDA 和 CPU 等成熟键的鲁棒性和多功能性。实现这一点需要跨关键 PyTorch 模块的深度集成。这种集成不仅仅是简单的切换——它涉及对 AMP（自动混合精度）、Autograd、分布式训练、检查点、数据加载器、优化和量化等核心组件的显著更新。
PrivateUse1 的激活是一个巨大的协作努力，最终产生了超过 100 个旨在将其从占位符转变为完全操作分发键的拉取请求。
-
PrivateUse1 的 UT/CI 质量保证
虽然单元测试对于确保在 PrivateUse1 机制的开发过程中质量至关重要，但仅靠它们本身并不能防止新的拉取请求无意中影响现有功能或树外设备的兼容性。
为了减轻这种风险，社区已经将
pytorch_openreg模块添加到测试套件中。该模块利用 CPU 后端来模拟与加速器的交互，为严格的测试创建一个受控环境。实施后，这将使我们在相关代码更新时自动执行设备通用的测试用例，使我们能够快速检测并解决影响 PrivateUse1 集成机制的可能问题。 -
综合文档
通过提供全面且易于理解的文档，我们旨在降低开发者的入门门槛，并鼓励更广泛地采用 PrivateUse1 机制在 PyTorch 生态系统中。该文档包括：
- 使用 PrivateUse1 逐步集成新后端的指南
- 对 PrivateUse1 功能及优点的清晰解释
- 代码示例和高效实现的最佳实践
这些增强旨在提高 PrivateUse1 机制的健壮性和可靠性，促进新后端的更好集成，并扩展 PyTorch 的功能。
上游与下游的兼容性
设备通用的单元测试
PyTorch 中的大多数单元测试都关注 CPU 和 CUDA 设备，这限制了其他硬件用户参与。为了解决这个问题，我们计划修改 PyTorch 的单元测试框架，以更好地支持非 CUDA 设备。该计划包括移除现有的设备限制，实现动态数据类型加载，并将装饰器泛化以适应更广泛的设备。此外，我们旨在强制使用通用设备代码，并扩展分布式测试以支持非 NCCL 后端。
通过这些改进，我们希望显著提高非 CUDA 设备的测试覆盖率和通过率，将它们整合到 PyTorch 的持续集成过程中。初始更改已经实施，为新的硬件支持铺平了道路，并为其他设备创建了一个参考模板。
通过自动化测试确保设备集成稳健
为维护 PyTorch 的高质量保证标准，已建立独立的构建仓库和每日持续集成（CI）工作流程，重点关注冒烟测试和集成测试。
pytorch-integration-tests 仓库自动化测试 PyTorch 的设备特定功能，确保它们在各种硬件平台（如 NPU 和其他专用设备）上正确且高效地运行。在仓库中，我们试图构建一个完全自动化的系统，持续验证 PyTorch 与不同硬件后端的兼容性。
- 自动化集成测试：使用 GitHub Actions 在多种设备上运行自动化测试。这种自动化确保代码库中的每个更改都经过彻底测试，针对多个硬件平台，以便在开发早期阶段捕捉潜在问题。
- 可复用工作流程：本存储库中的工作流程是模块化和可复用的，这简化了测试过程。开发者可以轻松地将这些工作流程适应新的设备或测试场景，使系统既灵活又可扩展，随着 PyTorch 的发展。
- 树外设备意识：该存储库显示了所有树外设备的存在和行为，使社区保持知情。这种方法最大限度地减少了意外破坏下游功能的风险，并提供了对更改的快速反馈。
加强多设备集成的工作对于其在不断发展的深度学习领域的适应性至关重要。这些举措不仅有利于现有用户，还降低了新硬件供应商和开发者的入门门槛，促进了人工智能和机器学习领域的创新。随着 PyTorch 的不断发展，其对灵活性、健壮性和包容性的承诺使其成为能够满足深度学习社区多样化需求的领先框架。
