一个框架如何支撑日益多样化的 AI 硬件？

在一篇新的 PyTorch Foundation 博客中，来自 @IBM 的贡献者分享了如何利用 PyTorch 现有的设备、流、内存管理和编译等抽象将 Spyre 集成进来。

其成果是更无缝的开发者体验、更健壮的 Spyre 运行时，以及一个无需为每种设备单独开发框架就能支持更广泛硬件的 PyTorch 生态。

了解更多：https://bit.ly/4y7RkKU

将 Spyre 构建为 PyTorch 原生设备 - PyTorch

来源：https://pytorch.org/blog/building-spyre-as-a-native-pytorch-device/

精选项目

• PyTorch logo (https://pytorch.org/projects/pytorch/)

一句话总结

Spyre 通过 torch-spyre（https://github.com/torch-spyre/torch-spyre）将 PyTorch 现有的设备、分配器、流和编译器抽象连接到 Spyre 运行时和固件，从而成为 PyTorch 的原生设备。PrivateUse1（https://docs.pytorch.org/tutorials/advanced/privateuseone.html）赋予 Spyre 真正的设备标识，PyTorch 分配器存储使张量驻留在 device="spyre" 上，流提供熟悉的顺序语义，运行时事件则在跨流的依赖工作之间进行连接。FX 图保留在 Inductor 编译器路径中。本文聚焦于 PyTorch 到 Spyre 的设备/运行时集成：设备标识、分配器支持的存储、流、事件以及已编译程序的启动如何映射到 Spyre 的运行时约束上。一旦编译产物生成完毕，启动过程就是将类型化操作按序提交到队列中，而非基于图的运行时接口。最终结果是即时模式和编译模式共用同一条路径，更低的启动开销，以及能够干净映射到 Spyre 硬件的 PyTorch 接口。

Spyre 是 IBM 的数据流 AI 加速器，针对推理进行了优化。它专为在 IBM Z、LinuxONE 和 Power 系统上与应用和数据一起运行 AI 的企业团队而设计。其低精度计算非常适合语言生成和嵌入模型中密集的矩阵运算。

对于开发者而言，挑战在于如何通过熟悉的 PyTorch 代码使用该硬件，同时跟上新模型和推理技术的步伐。高效执行不仅需要一个编译器：张量需要能在操作之间保留在设备上，启动需要足够轻量，主机工作和数据传输应在依赖关系允许的情况下与计算重叠。本文介绍了 IBM 的 torch-spyre 团队如何基于 PyTorch 社区的接口，在数据流加速器上支持这些工作流。

PyTorch 接口背后的数据流约束

该卡有 32 个核心，通过高带宽环形网络互连。每个核心拥有 2 MB 的本地暂存存储器和处理单元阵列。最多 128 GB 的 LPDDR5 提供张量和程序的存储空间。

这两级内存有不同的所有者。运行时栈分配并管理 LPDDR5。在程序运行之前，编译器已经生成了在 LPDDR5 与每个核心的暂存存储器之间搬移数据块的加载和存储指令。数据以 128 字节的“stick“为单位沿该路径移动，因此内存布局和对齐方式会影响取数效率。运行时不直接分配或调度暂存存储器。

原生 PyTorch 设备仍然必须遵循底层硬件的执行模型。对 Spyre 来说，这意味着将 PyTorch 的设备、分配器、流、事件和启动抽象映射到一个围绕已编译程序、有序队列、显式数据搬移和固定布局约定构建的运行时上。四个约束驱动着这一设计。

计算由数据触发，而非由单一的内核指令流触发。 GPU 内核通常围绕调度好的线程组执行内核指令序列来组织。已编译的 Spyre 内核则以不同方式编排：它包含多个功能单元的程序——如处理单元、特殊函数流水线和加载/存储单元——每个单元运行为该单元编译的程序。硬件队列和本地内存将这些单元连接在一起，因此进展由操作数在消费单元处可用所驱动。Spyre 编译器构建设备端的生产者/消费者调度；运行时将已编译的内核连同其张量参数作为一个完整的设备计算操作来提交。

运行时重叠来自独立的流水线，而非并发的计算启动。 Spyre 拥有计算流水线和数据搬移流水线，因此在已编译程序进行计算时，传输可以并行运行。这些流水线不可互换：一个已编译程序可以并行使用多个核心，但独立提交的计算操作共享同一个运行时计算队列，彼此排队等待。因此，运行时级别的重叠来自在计算的同时搬移数据，而非同时启动两个计算操作。

有序队列保留每个流内的完成顺序。 工作以一系列类型化操作的形式到达——把这个搬进来，运行这个，把那个搬出去——每个流按顺序完成这些操作。为了在不削弱该保证的情况下使用独立的硬件流水线，运行时将独立的传输和计算工作分配到不同的流上，并在一方消费另一方结果时通过事件进行连接。将所有工作放在一个流上虽然保持了正确性，但会串行化这些流水线。

程序针对固定布局进行提前编译。 编译器确定值如何组合成设备原生的固定大小块以及维度出现的顺序，它还可以为属于该作业自身的内存（中间结果和编译器拥有的临时缓冲区）确定地址。但它无法确定的是由调用方拥有的地址：张量输入和输出保持符号化，因为分配器尚未放置它们。运行时必须在启动时提供这些地址。

在本文描述的实现中，PyTorch 不直接与卡通信。torch-spyre 提供 PyTorch 设备集成，而 Spyre 运行时和固件管理设备工作和通信。这种分层正是本文将张量、流和事件等 PyTorch 概念与设备区域、DMA 可见的主机暂存缓冲区和显式工作提交等底层关注点分开讨论的原因。

Spyre 映射到的 PyTorch 抽象

PyTorch 已有可以映射到这些约束上的概念。

设备标识。 PrivateUse1（https://docs.pytorch.org/tutorials/advanced/privateuseone.html）注册始于两个调用：

torch.utils.rename_privateuse1_backend("spyre") torch._register_device_module("spyre", make_spyre_module())

这些调用将 PrivateUse1 命名为 spyre 并注册其设备模块。有了分配器、拷贝内核和调度器注册作为该标识的支撑，tensor.to("spyre") 即可工作，调度器将 Spyre 张量路由到 Spyre 内核。

设备管理 API。 后端特定的控制位于 torch.spyre 下，而 torch.accelerator（https://docs.pytorch.org/docs/main/accelerator/index.html）为可移植代码提供了有关可用性、设备选择、流、同步和内存报告的通用接口。张量分配和内核启动仍然通过张量、调度器、设备分配器和后端运行时进行。

流抽象。 实现框架要求的设备守卫钩子，就能让原生的 torch.Stream(device="spyre") 和 torch.stream(...) 在你的设备上工作。框架原生的接口之所以能工作，是因为这些钩子的存在，而非后端重新实现了 PyTorch 的流 API。

设备分配器接口。 其核心契约返回设备存储及其释放方式。额外的钩子将后端连接到 PyTorch 的内存统计、缓存控制和流关联。

引用计数的张量存储。 张量的生命周期成为框架的职责：当对 Spyre 张量的最后一个引用消失时，PyTorch 会回调后端分配器来释放或回收设备分配。

通过 Inductor 的编译器路径。 torch.compile(backend="inductor") 使 FX 图保持在 PyTorch 的编译器流水线中。相同的降低、缓存和启动路径可以同时服务于已编译模型和注册的即时操作。

事件模型。 PyTorch 将事件定义为表达流之间顺序关系的方式。Spyre 内部通过运行时事件和派生的依赖边使用相同的模型。

将二者并列对比，对应关系非常紧密：

本文接下来将深入介绍四个需要真正设计工作的映射。

映射 1 —— 执行流水线到流

运行时需要一个有序、异步且按设备组织的工作放置位置。在 PyTorch 中，这个抽象就是流。对 Spyre 来说，有用的映射不是硬件行为与 PyTorch 流完全一致，而是流排序的工作可以被降低为类型化的运行时操作，同时运行时保留这些操作所需的依赖关系。真正需要判断力的是决定什么应放在该边界之上。

与程序相关的工作保持在流边界之上。在 Spyre 上，面向框架的层将编译产物转化为预备好的启动方案：一个记录运行该产物所需的主机工作、数据搬移、设备计算和同步的“配方“。该层可以携带张量布局和符号化参数元数据，并绑定由 PyTorch 传入的实际张量的启动时值。一旦工作通过流提交，底层运行时看到的是具有明确操作数和依赖关系的类型化操作，而不是需要重新解释的程序描述。

底层永远不会误解程序描述，因为它根本不会收到程序描述。不同的提交机制也可以置于其下，而无需更改面向框架的代码。

**

图 1：边界应该放在哪里

将程序知识保留在队列之上。 所有理解“编译程序是什么“的部分都属于面向框架的层。面向运行时的层接受具有明确操作数和依赖关系的类型化操作，涵盖主机工作、设备计算、数据搬移、同步和其他运行时任务。

预备和启动发生在不同时刻。将编译产物翻译为可提交的操作是每个产物一次的工作，而非每次启动的工作：

``

仅在首次看到编译产物时执行一次

job_plan = prepare_kernel(spyrecode_dir)

每次调用时执行

launch_jobplan(job_plan, args) ``

预备阶段解析产物、分配并传输程序二进制，将执行计划翻译为有序的类型化步骤列表。启动阶段遍历这些步骤并入队操作。任何在产物首次被看到时就能确定的事情都应在那时确定，使得启动路径只做操作的构造和入队。

为什么不在运行时再画一张图？

早期的集成同样从一个 PyTorch FX 图和自定义后端实现开始。FX 不是问题；图属于编译器边界。问题出在此交接之后，当启动、拷贝和张量参数被翻译成第二张后端特定的运行时图时。

运行一个已编译程序意味着读取和反序列化该运行时图，为张量参数添加节点和边，加载和解析结果后才发出设备操作。主机到设备的拷贝同样会围绕数据转换和传输构建一个小图。当运行时需要完整图时，该接口是有用的，但它为在启动时已完全确定的操作增加了额外的图处理工作。

将整个编译图复制到我们自己的抽象中还导致了核心运行时与 PyTorch 之间通信的断裂。PyTorch 无法得知张量是驻留在设备上还是主机上。结果，PyTorch 认为每个张量都是 CPU 张量，即使核心运行时实际上已将其移到了设备上。

新的边界将 FX 图保留在 Inductor 中，并只预备一次编译产物。此后，运行时需要的是有序的类型化操作配方，而非每次调用时需要重建或遍历的另一张模型图。

映射 2 —— 设备内存到分配器

PyTorch 原生设备不能将每个操作都视为“将输入拷贝到加速器、运行、再将输出拷贝回来“。位于 device="spyre" 上的张量需要真实的设备存储，其生命周期由 PyTorch 通过分配器和存储契约管理。早期的自定义 torch.compile 后端可以通过专用的运行时路径让特定的模型状态（如权重和 KV 缓存条目）驻留在设备上，但这不等同于 PyTorch 原生的张量驻留。一个普通的 PyTorch 值并不是简单地成为一个存储由 PyTorch 管理、分配可被下一个操作复用的 device="spyre" 张量。这对即时执行尤其昂贵，因为每个小操作都会额外产生一个传输边界。

Spyre 在该基本分配器契约之上添加了一个硬件约束。框架分配器通常将设备内存呈现为一个以指针寻址的平坦内存池。在该接口之下，Spyre 以区域（region）为单位管理内存——每个区域是由句柄标识的一段连续设备内存。一个区域可以包含多个张量的分配。由于句柄预算有限，尤其在一张卡被多个租户共享时，分配器必须将许多张量压缩到少量区域中。

Spyre 支持 PF 模式（卡专用于一个租户）和 VF 模式（多个租户共享）。VF 模式的句柄预算更为紧张。分配器获取少量大区域，并在其中子分配对齐的块，使多个存活张量可以共享少量句柄。

无论哪种模式，每次分配都使用相同类型的描述：一个或多个“区域加偏移“片段。只有对区域标识的解释会有所不同，因此分配器之上的层不需要根据部署模式进行分支。分配器返回该描述及释放方式；PyTorch 的引用计数在张量存活期间将其保留在设备上。

更完整的分配器接口还添加了 PyTorch 的标准内存管理 API，包括使用统计、缓存控制和流关联。这些有用，但它们并非使张量驻留的原因。

这里的分配不仅仅由一个地址描述：基本可寻址单元是一个区域标识、区域内的偏移量和长度。交错分布在多个内存域中的张量由多个这样的片段组成，而非一个平坦指针，PyTorch 为此预留了空间：分配可以在指针旁边携带一个不透明上下文（opaque context），该描述就存放其中，引用计数仍然负责清理。

图 2：设备分配的真实含义

分配是一种描述，而不仅仅是一个地址。 绑定到单个内存域的张量是一个区域加偏移的片段；交错分布在多个域上的张量则有多个片段。PyTorch 将该描述存放在不透明上下文中，