跳转到主要内容
博客

PyTorch 多设备集成中的挑战与努力：兼容性、可移植性和集成效率

引言

随着对多种硬件加速器需求的增长，对健壮且适应性强的深度学习框架的需求变得越来越重要。在进行这种集成时，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 生态系统的扩展。

bit diagram

为了应对这一挑战，PrivateUse1 机制进行了一系列优化，以增强其容量。

  • **PrivateUse1 集成机制** 最初保留作为备用选项，**PrivateUse1** 以及 **PrivateUse2** 和 **PrivateUse3** 被设计为仅在现有键资源稀缺时激活。**PrivateUse1** 现在正在开发中，以匹配 CUDA 和 CPU 等现有键的鲁棒性和多功能性。实现这一目标需要深度集成到关键 PyTorch 模块中。这种集成不仅仅是一个简单的切换，它涉及对核心组件的重大更新，例如 **AMP（自动混合精度）**、**Autograd**、**分布式训练**、**检查点**、**DataLoader**、**优化**和**量化**等。
flow diagram

**PrivateUse1** 的激活是一项巨大的协作努力，最终完成了 100 多个拉取请求，旨在使其从一个占位符发展成为一个完全可操作的调度键。

  • **PrivateUse1 UT/CI 质量保证** 尽管单元测试对于确保 **PrivateUse1** 机制开发过程中的质量至关重要，但它们本身不足以防止新的拉取请求无意中影响现有功能或树外设备的兼容性。为了降低这种风险，社区已将 `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 的不断发展，其对灵活性、健壮性和包容性的承诺使其成为一个能够满足深度学习社区多样化需求的领先框架。