博客

OpenReg：一个自包含的 PyTorch 加速器模拟器

引言

PyTorch 社区正积极致力于构建一个日益增长的专用加速器生态系统，为开发者提供一种标准化且灵活的方式来添加自定义硬件支持。随着其应用范围的扩大，社区正聚焦于其长期稳定性、清晰的参考标准和详尽的文档，以促进更广泛的协作并加速创新。

OpenReg 旨在通过利用 PrivateUse1 DispatchKey（已成为外部硬件集成的社区标准方法）来实现这些目标。

OpenReg 提供以下功能

  • 测试后端：作为 PrivateUse1 的内置（in-tree）测试后端，通过 CI/CD 确保质量稳定性。
  • 集成参考：作为新后端集成的参考示例。
  • 集成文档：提供与代码相对应的模块级集成文档。

注意

OpenReg 的目标不是实现一个功能完备、高性能的 PyTorch 后端，而是作为一个用于机制验证的极简参考实现。

架构

OpenReg 的架构

OpenReg 的整体架构可分为三个方面：基于 CPU 的模拟器、CI/CD 集成和配套文档。

模拟器

OpenReg 被设计为一个独立的、自包含的 PyTorch 模拟后端。其核心目的是通过 PrivateUse1 DispatchKey 促进加速器与 PyTorch 计算图及执行环境的集成，利用 PrivateUse1 设备类型进行自定义操作和设备管理。核心组件与特性：

  • 基于 CPU 的加速器模拟器： OpenReg 强大的基于 CPU 的加速器模拟器对于在无法立即访问硬件的情况下开发、测试和调试自定义 PyTorch 操作及后端逻辑至关重要。它确保了功能正确性，支持早期开发，并在部署前提供准确的性能预测和资源利用率分析。
  • PyTorch 加速器集成机制： OpenReg 是将第三方软硬件集成到 PyTorch 中的参考范例。它使用 PrivateUse1 设备类型 处理自定义算子、内核和内存管理，从而避免冲突并简化外部供应商的开发工作，进而扩展 PyTorch 生态系统。

CI/CD 集成

为了确保加速器的稳定性、可靠性和持续的质量，我们正在实施一套全面的 CI/CD 策略。

  • 将所有加速器集成机制测试迁移至 OpenReg： 所有基于 PrivateUse1 的加速器集成测试都将迁移到 OpenReg 模块中，确保所有新集成的一致性、可维护性和单一事实来源（Single Source of Truth）。
  • 将 OpenReg 加入 CI/CD 工作流： OpenReg 已完全集成到现有的 CI/CD 流水线中，实现了每次代码提交时的自动化测试、构建和部署。这确保了代码质量，识别回归问题，并通过单元测试、集成测试和系统级测试保持稳定性和功能性。

文档

全面且易于获取的文档对于 OpenReg 的采用和可用性至关重要。

  • 加速器集成指南： 为开发者集成其自有加速器到 PyTorch 提供详细指导，涵盖了运行时、算子注册及其他高级特性。

实现

OpenReg 的实现位于 pytorch/test/cpp_extensions/open_registration_extension/torch_openreg，其文件结构如下：

torch_openreg/
├── csrc  // Core device implementation
│   ├── amp
│   ├── aten  // Operator registration
│   │   ├── native
│   └── runtime // Implementations for Host memory, device memory, Guard, etc

│        └── ...
├── third_party // A C++ library that simulates a CUDA-like device using the CPU
│   └── openreg
└── torch_openreg // Python interface implementation
    ├── csrc // Python C++ binding code
    └── openreg // Python API
        └── ...

特性

OpenReg 是一个旨在通过 PrivateUse1 将加速器与 PyTorch 无缝集成的参考实现。它提供了一套全面的特性，包括：

  • 算子注册： OpenReg 为注册各类算子提供了稳健的机制：
    • 内置 PyTorch 算子： 支持注册现有的 PyTorch 算子，确保兼容性并利用丰富的 PyTorch 生态系统。
    • 自定义算子： 促进了定制化算子的集成，允许开发者通过专用计算来扩展 PyTorch 的功能。
    • 回退机制： 实现了一套回退系统，以优雅地处理注册算子不可用或不适用时的情况，确保操作的连续性。
  • 运行时支持： OpenReg 提供全面的运行时支持，包含事件管理、流处理、设备管理、执行守卫和随机数生成 (RNG) 等必要功能。这确保了对执行环境的精细控制和优化。
  • 跨平台支持： OpenReg 支持 Linux、macOS (OSX) 和 Windows 操作系统。
  • 高级特性： 
    • 自动混合精度 (AMP)： 提供混合精度便捷方法，其中部分操作使用 torch.float32 数据类型，其他操作使用低精度浮点数据类型：torch.float16 或 torch.bfloat16。
    • 自动加载： 当执行 import torch 时，已安装的加速器（如 torch_openreg）将被自动加载，获得与内置后端相同的体验。

安装与使用

OpenReg 需要从源码构建，但通过运行以下命令即可简单完成：

cd test/cpp_extensions/open_registration_extension/torch_openreg

python -m pip install --no-build-isolation -e . # for develop
python -m pip install --no-build-isolation .    # for install

使用 OpenReg 就像使用 CUDA 设备一样直接——您只需要通过 torch.device("openreg") 声明您的设备，如下面的代码所示：

import torch

if not torch.openreg.is_available():
    print("OpenReg backend is not available in this build.")
    exit()

print("OpenReg backend is available!")

device = torch.device("openreg")

x = torch.tensor([[1., 2.], [3., 4.]], device=device)
y = x + 2
print("Result y:\n", y)
print(f"Device of y: {y.device}")

z = y.cpu()
print("Result z:\n", z)
print(f"Device of z: {z.device}")

下一步

OpenReg 有明确的路线图来增强其在 PyTorch 生态系统中的能力。未来的改进包括：用于大规模模型的分布式计算集成、用于互操作性的 ONNX 支持、用于优化性能的高级编译器特性，以及旨在降低采用门槛的综合文档。

OpenReg 也正被考虑作为关键测试后端纳入持续集成 (CI) 流水线。这将使开发者能够验证硬件能力，确保 PyTorch 的稳定性和可靠性（尤其是对于非内置的 out-of-tree 后端），并通过对变更提供即时反馈来加速开发周期。

结论

通过利用 CPU 作为替代加速器，OpenReg 确保了 PyTorch 的 PrivateUse1 机制保持稳定可靠。它既是生态系统的安全网，也是新硬件集成的参考指南。本质上，OpenReg 是一个具有前瞻性的项目，为不断演进的 AI 领域和 PyTorch 生态系统提供了增强功能和稳健的测试机制。

如果您正在考虑为 PyTorch 添加新的设备后端，OpenReg 是您的起点。它虽然不能取代在性能或可扩展性方面所需的详尽工作，但它为您正确设置设备提供了清晰的路径。

如果您有兴趣与我们联系，并帮助培育一个更繁荣的 PyTorch 加速器生态系统，请考虑加入加速器集成工作组 (Accelerator Integration Working Group)，该工作组隶属于 PyTorch 技术咨询委员会 (TAC)。这是讨论和塑造 PyTorch 加速器集成未来的主要社区论坛。