引言
随着对各种硬件加速器的需求增长，对强大且适应性强的深度学习框架的需求变得越来越关键。在进行这种集成的过程中，PyTorch 生态系统中出现了一些挑战，可能影响到各种硬件供应商。本博客旨在强调这些问题，并提出解决方案，以增强 PyTorch 在不同硬件平台上的适应性、可移植性和弹性。
通过加速器自动加载提高用户代码的可移植性
目前，用户在不同加速器上运行代码时面临额外的工作。其中一项任务是手动导入非官方设备 (out-of-tree devices) 的模块。这不仅要求用户理解不同加速器之间的不同使用模式，还需要让他们的代码感知到这些差异。如果你有原本在 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 能够屏蔽用户免受设备差异的影响是一项竞争优势。加速器自动加载 (Accelerator Autoloading) 允许用户继续使用熟悉的 PyTorch 设备编程模型，而无需显式加载或导入设备特定的扩展。
它是如何工作的？
利用 Python 的插件架构，通过 PyTorch 包中的入口点 (entry points) 实现设备扩展的自动加载。
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）、布局（dense、sparse）和自动微分 (autograd) 功能）的标识符。这些键确保操作被导向正确的实现。
PrivateUse1 是一个可定制的设备调度键，类似于 CUDA/CPU/XPU 等），专用于非官方设备。它为开发者提供了一种在不修改核心框架的情况下扩展 PyTorch 功能的方式，从而可以集成新设备、硬件加速器或其他专用计算环境。
为什么我们需要 PrivateUse1？
在内部，调度键表示为位掩码 (bit masks)，每个位代表某个键是否处于活动状态。这种位掩码表示形式对于快速查找和组合键是高效的，但它本质上限制了不同键的数量（通常为 64 个或更少）。
PyTorch 中 BackendComponent 调度键的当前实现遇到了一个关键瓶颈，这限制了新后端的添加，从而限制了 PyTorch 生态系统的扩展。
为了应对这一挑战，对 PrivateUse1 机制应用了一系列优化，以增强其容量。
-
PrivateUse1 集成机制
最初被保留为备用选项，PrivateUse1 以及 PrivateUse2 和 PrivateUse3 被设计为仅在现有键资源变得稀缺时才激活。
目前正在开发 PrivateUse1，使其具有与 CUDA 和 CPU 等已建立的键相同的稳健性和多功能性。实现这一目标需要在关键的 PyTorch 模块中进行深度集成。这种集成不仅仅是简单的切换，它涉及到对核心组件的重大更新，例如 AMP (Automatic Mixed Precision，自动混合精度)、Autograd (自动微分)、Distributed Training (分布式训练)、Checkpointing (检查点)、DataLoader (数据加载器)、Optimization (优化器) 和 Quantization (量化) 等。
激活 PrivateUse1 是一项巨大的协作努力，最终完成了 100 多个拉取请求，旨在将其从一个占位符变成一个完全可操作的调度键。
-
PrivateUse1 UT/CI 质量保证
虽然单元测试 (unit tests) 对于确保 PrivateUse1 机制开发过程中的质量至关重要，但仅靠它们不足以防止新的拉取请求无意中影响现有功能或非官方设备的兼容性。
为了降低这种风险，社区已将
pytorch_openreg模块添加到测试套件中。该模块利用 CPU 后端模拟与加速器的交互，创建受控环境进行严格测试。实现后，这将使设备通用测试用例在相关代码更新时自动执行，从而使我们能够快速检测和解决影响 PrivateUse1 集成机制的任何潜在问题。 -
全面文档
通过提供全面且易于理解的文档，我们旨在降低开发者的入门门槛，并鼓励在 PyTorch 生态系统中更广泛地采用 PrivateUse1 机制。此文档包括：
- 使用 PrivateUse1 集成新后端的逐步指南
- PrivateUse1 功能和优势的清晰解释
- 高效实现的示例代码和最佳实践
这些改进旨在提高 PrivateUse1 机制的稳健性和可靠性，促进新后端的更好集成并扩展 PyTorch 的能力。
上游与下游的兼容性
设备通用单元测试
PyTorch 中的大多数单元测试都侧重于 CPU 和 CUDA 设备，这限制了使用其他硬件的用户的参与。为了解决这个问题，我们计划修改 PyTorch 的单元测试框架，以便更好地支持非 CUDA 设备。此计划包括移除现有的设备限制、实现动态数据类型加载以及泛化装饰器 (decorators) 以适应更广泛的设备。此外，我们旨在强制使用通用设备代码，并扩展分布式测试以支持非 NCCL 后端。
通过这些改进，我们希望显著增加非 CUDA 设备的测试覆盖率和通过率，将它们集成到 PyTorch 的持续集成 (continuous integration) 流程中。初步更改已实施，为新的硬件支持铺平了道路，并为其他设备创建了参考模板。
通过自动化测试确保稳健的设备集成
为了维护 PyTorch 的高质量保证标准，我们建立了一个独立的构建仓库和日常持续集成 (CI) 工作流程，重点关注冒烟测试 (smoke testing) 和集成测试 (integration testing)。
pytorch-integration-tests 仓库自动化测试 PyTorch 的设备特定功能，确保它们在各种硬件平台（NPU 和其他专用设备）上正确高效地运行。在此仓库中，我们正努力建立一个完全自动化的系统，持续验证 PyTorch 与不同硬件后端的兼容性。
- 自动化集成测试 (Automated Integration Tests)：使用 GitHub Actions 在不同设备上运行自动化测试。这种自动化确保代码库中的每次更改都在多个硬件平台进行全面测试，从而在开发过程早期发现潜在问题。
- 可重用工作流 (Reusable Workflows)：此仓库中的工作流是模块化和可重用的，这简化了测试过程。开发者可以轻松地将这些工作流调整到新设备或测试场景，使系统随着 PyTorch 的发展而具有灵活性和可扩展性。
- 非官方设备感知 (Awareness of Out-of-Tree Devices)：仓库展示了所有非官方设备的存在和行为，使社区保持了解。这种方法最大程度地降低了意外破坏下游功能 (downstream functionalities) 的风险，并提供了对变更的快速反馈。
增强多设备集成的努力对于 PyTorch 在不断发展的深度学习领域的适应性至关重要。这些举措不仅惠及现有用户，还降低了新硬件供应商和开发者的入门门槛，促进了人工智能和机器学习的创新。随着 PyTorch 的持续发展，其对灵活性、稳健性和包容性的承诺使其成为能够满足深度学习社区多样化需求的领先框架。
