Spyre 通过 torch-spyre 将 PyTorch 现有的设备、分配器、流和编译器抽象连接到 Spyre 运行时和固件,从而成为原生 PyTorch 设备。PrivateUse1 赋予 Spyre 真正的设备身份,PyTorch 分配器存储使张量驻留在 device="spyre" 上,流提供熟悉的排序语义,运行时事件连接跨流的依赖工作。FX 图保留在 Inductor 编译器路径中。本文聚焦于 PyTorch 到 Spyre 的设备/运行时集成:设备身份、分配器支持的存储、流、事件和编译程序启动如何映射到 Spyre 的运行时约束。一旦编译产物存在,启动就是一组在有序队列上执行的类型化操作的预备配方,而非基于图的运行时接口。其结果是 eager 和编译执行共用一条路径、更低的启动开销,以及一个能干净映射到 Spyre 硬件的 PyTorch 接口。
Spyre 是 IBM 的数据流 AI 加速器,针对推理进行了优化。它专为在 IBM Z、LinuxONE 和 Power 系统上运行 AI 应用及数据的企业团队而设计。其降低精度的计算非常适合语言生成和嵌入模型中的矩阵密集型工作。
对开发者而言,挑战在于通过熟悉的 PyTorch 代码使用该硬件,同时跟上新模型和推理技术。高效执行需要的不仅仅是编译器:张量必须能够在操作之间驻留在设备上,启动需要轻量化,主机工作和数据传输应在依赖允许的情况下与计算重叠。本文解释了 IBM 的 torch-spyre 团队如何基于 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 可见的主机暂存缓冲区以及显式工作提交等更低层关注点区分开来。
PyTorch 已经有一些概念可以映射到这些约束上。
设备标识。PrivateUse1 注册从两个调用开始:
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 为可移植代码提供通用接口,用于可用性、设备选择、流、同步和内存报告。张量分配和内核启动仍然通过张量、调度器、设备分配器和后端运行时进行。
流抽象。实现框架所需的设备保护钩子,才能使原生 torch.Stream(device="spyre") 和 torch.stream(...) 在你的设备上工作。框架原生接口之所以可用,是因为这些钩子存在,而不是因为后端重新实现了 PyTorch 的流 API。
设备分配器接口。其核心契约返回设备存储以及释放它的方式。额外的钩子将后端连接到 PyTorch 的内存统计、缓存控制和流关联。
引用计数的张量存储。张量生命周期成为框架的问题:当对 Spyre 张量的最后一个引用消失时,PyTorch 会回调后端分配器以释放或回收设备分配。
通过 Inductor 的编译器路径。torch.compile(backend="inductor") 将 FX 图保留在 PyTorch 的编译器流水线中。相同的 lowering、缓存和启动路径可以服务于编译模型和已注册的 eager 操作。
事件模型。PyTorch 将事件定义为表达流之间顺序的一种方式。Spyre 在内部通过运行时事件和派生的依赖边使用相同的模型。
并排来看,对应关系非常接近:
| 运行时/硬件需求 | PyTorch 抽象 | | --- | --- | | 有序的类型化操作队列 | Stream | | 传输流水线,与计算分离 | 多个 stream | | 跨流水线排序 | Events,或运行时推导出的边 | | 按区域预算的设备内存 | 设备分配器 | | 驻留设备的值 | Tensor 存储及其生命周期 | | 提前编译的程序 | torch.compile 背后的编译产物 | | 运行时管理的主机工作 | 与 stream 关联的有序主机操作 |
本文的其余部分将逐一介绍四个需要真正设计工作的映射。
运行时需要有地方来放置有序、异步且按设备划分的工作。在 PyTorch 中,这个抽象就是 stream。对于 Spyre 来说,有用的映射并不是硬件的行为与 PyTorch stream 完全一致;而是 stream 有序的工作可以被降级为类型化的运行时操作,同时运行时保留这些操作所需的依赖关系。需要判断的是如何决定这个边界两侧各自放置什么。
程序特定的工作留在 stream 边界之上。在 Spyre 上,面向框架的层将编译产物转换为准备好的启动计划:一份记录运行该产物所需的主机工作、数据移动、设备计算和同步的配方。该层可以携带张量布局和符号参数元数据,并为 PyTorch 传入的实际张量绑定启动时的值。一旦工作通过 stream 提交,下层运行时看到的是具有显式操作数和依赖关系的类型化操作,而不是需要重新解释的程序描述。
下层无法误解程序描述,因为它根本不会收到程序描述。不同的提交机制也可以位于其下,而无需更改面向框架的代码。
图 1:边界应该放在哪里
将程序知识保留在队列之上。所有理解编译程序是什么的内容都属于面向框架的层。面向运行时的层接受具有显式操作数和依赖关系的类型化操作,用于主机工作、设备计算、数据移动、同步和其他运行时任务。
准备和启动也发生在不同的时间。将编译产物转换为可提交的操作是按产物进行的工作,而不是按启动进行的工作:
# once, when the compiled artifact is first seen job_plan = prepare_kernel(spyrecode_dir) # per invocation launch_jobplan(job_plan, args)
准备阶段解析产物,分配并传输程序二进制文件,并将执行计划转换为有序的类型化步骤列表。启动阶段遍历这些步骤并将操作入队。任何在首次看到产物时就能解析的内容都应该在那里解析,这样启动路径就只是操作构造和入队。
早期的集成已经通过自定义后端实现从 PyTorch FX 图开始。FX 本身不是问题;图本就属于编译器边界。不匹配出现在那次交接之后,当时启动、拷贝和张量参数被转换为第二个后端特定的运行时图。
运行一个已经编译好的程序意味着读取并反序列化该运行时图,为其张量参数添加节点和边,并在发出设备操作之前加载和解析结果。主机到设备的拷贝同样围绕数据转换和传输构建了一个小图。当运行时需要整个图时,该接口是有用的,但它为那些在启动时已经完全确定的操作增加了图处理工作。
将整个编译后的图复制到我们自己的抽象中也导致了核心运行时与 PyTorch 之间的通信中断。PyTorch 无法知道张量是驻留在设备上还是主机上。因此,即使核心运行时实际上已将张量移动到设备上,PyTorch 仍认为每个张量都是 CPU 张量。
新的边界将 FX 图保留在 Inductor 中,并一次性准备编译后的产物。之后,运行时需要的是一个有序的类型化操作配方,而不是另一个在每次调用时都要重建或遍历的模型图。
一个 PyTorch 原生设备不能把每个操作都视为“将输入拷贝到加速器,运行,再将输出拷贝回来”。一个位于 device="spyre" 上的张量需要真实的设备存储,其生命周期由 PyTorch 通过其分配器和存储契约来管理。早期的自定义 torch.compile 后端可以通过定制的运行时路径保持特定模型状态(如权重和 KV-cache 条目)常驻,但这与 PyTorch 原生张量驻留并不相同。一个普通的 PyTorch 值并不只是一个 device="spyre" 张量,其存储由 PyTorch 管理,且其分配可以被下一个操作重用。这对于 eager 执行尤其代价高昂,因为否则每个小操作都会创建另一个传输边界。
Spyre 为这一基本分配器契约增加了一个硬件约束。框架分配器通常将设备内存呈现为由指针寻址的扁平池。在该接口之下,Spyre 以区域(region)为单位管理内存——每个区域是一块由句柄标识的连续设备内存。一个区域可以包含许多张量的分配。由于句柄预算有限,尤其是在一张卡被多个租户共享时,分配器必须将许多张量装入少量区域中。
Spyre 支持 PF 模式,即一张卡专用于一个租户,以及 VF 模式,即多个租户共享一张卡。VF 模式的句柄预算更紧张。分配器获取少量大区域,并在其中子分配对齐的块,从而让许多活跃张量共享少量句柄。
在任一模式下,每次分配都使用同一种描述:一个或多个区域加偏移的片段。只有区域标识符的解释会改变,因此分配器之上的各层不会根据部署模式进行分支。分配器返回该描述以及释放它的方式;PyTorch 的引用计数在张量存活期间将其保留在设备上。
更完整的分配器接口增加了 PyTorch 的标准内存管理 API,包括使用统计、缓存控制和流关联。它很有用,但并不是使张量驻留的原因。
这里的分配不仅仅由地址描述:基本的可寻址片段是区域标识符、该区域内的偏移量和长度。一个跨内存域交错的张量是多个这样的片段,而不是一个扁平指针,而 PyTorch 为此留有余地:分配可以在其指针旁携带一个不透明上下文,因此描述就放在那里,引用计数仍然负责清理。
分配是一种描述,而不仅仅是一个地址。绑定到单个内存域的张量是单个区域加偏移片段;跨域交错的张量则是多个。PyTorch 在不透明上下文中携带该描述,其引用计数会告知后端分配器何时可以释放存储。
相同的表示方式可以覆盖具有非统一内存的硬件。在此类设备上,内存被划分为多个域,核心访问某些域的效率高于其他域。张量可以绑定到使用它的核心附近,或跨域交错以利用其组合带宽。这是单个设备内的局部性,而非多卡分布。
放置必须显式且绑定,因为编译后的代码可能依赖它。因此,每一层携带分配信息时,也必须携带该分配是绑定到一个域还是分布在多个域上。
驻留还使布局成为显式契约。模型适配器将权重和 KV 缓存放置在编译内核期望的布局中,而编译器 pass 为中间张量插入合法的布局转换。工作并未消失;它从隐藏的运行时图转移到了模型准备、分配元数据和编译器中。
设备放置和卸载是两个独立的问题。放置决定张量在加速器内的位置。卸载将当前不需要的状态(如 KV 缓存页)移至二级存储层,并在使用前恢复。该层可以是主机内存,或在服务系统层面是 SSD 或网络附加存储。运行时接口应描述传输过程,而不暴露该层背后的物理介质。
每个提前编译加速器都必须在某处解决这个问题。编译器知道计算的结构,但不知道调用者的数据将在哪里,两者必须在程序运行前协调一致。随着运行时的演进,Spyre 使用了三种方法。运行时现在专门使用第三种方法:运行时补丁。
早期的图中介路径让编译后的作业拥有地址。输入被复制到作业拥有的设备缓冲区中,结果被复制回来。这些缓冲区可以在多次启动之间保持具体地址,因为它们的生命周期和放置属于作业。这仍然适用于作业完全拥有的内存,如中间结果和编译器拥有的临时缓冲区,但它不能推广到存储由 PyTorch 分配器选择的驻留张量。
一个过渡运行时将程序的地址窗口指向张量。设备转换表(xlat 表)暴露一小组地址窗口,因此运行时将每个窗口绑定到一个由分配器放置的张量。无需修补任何内容,张量也不需要复制到编译器拥有的位置。限制在于输入参数数量而非分配:每个窗口一个张量且只有少量窗口,大型融合内核无法接受其所有张量。
当前运行时在启动时对程序进行补丁。在分配器放置张量之前,地址保持符号化。这避免了重复的输入和输出复制以及窗口数量上限,使其成为分配器放置张量的通用路径。在 Spyre 上,编译器生成的产物布局固定,但张量起始地址可能是符号化的;二进制文件在地址将要去的位置携带占位符变量。替换真实值是与其他所有操作相同的计划中的三个普通步骤:
主机回调在 CPU 上运行,将解析后的地址写入固定主机缓冲区。
传输将该缓冲区复制到程序自身的分配中,位于编译器指定的偏移处。
计算运行,读取这些值,并修补其操作数。
在单个流上,完成顺序足以防止这三个步骤发生竞争。在接下来描述的初始双流设计中,主机回调和传输在一个准备流上保持有序。计算操作在设备流上运行,并等待传输完成后记录的事件。这保持了相同的顺序,同时允许准备与先前的计算操作重叠。
队列顺序控制操作何时被发出。它不控制超前运行的硬件。Spyre 的程序分发正是如此:馈送指令缓冲区的单元急切地获取,而不是等待当前程序完成,因此它可以在补丁到达内存之前将即将被补丁的程序拉入缓冲区,然后运行过时的字节。补救措施是在程序分发层面低一级的屏障,保持获取直到补丁落地。队列对提交进行排序;低级屏障防止读后写竞争。
启动时修补包含三个普通步骤和两项保证。主机回调将解析后的地址写入固定缓冲区,一次传输将它们移动到程序分配中编译器指定的偏移处,设备读取它们并修补其操作数。准备流对主机回调和传输进行排序,设备流在计算前等待传输完成。程序分发中的屏障也防止取指引擎在修补落地之前读取程序。
初始的双流设计可以复用一个固定的主机暂存缓冲区,因为对该缓冲区的每次写入和从该缓冲区的每次传输都保持在同一个准备流上。完成顺序防止下一个主机回调覆盖主机缓冲区,直到前一次传输已消费它。这一保证不会自动扩展到多个准备流或更深的主机预运行。这些情况需要为每个并发传输使用单独的主机缓冲区,或使用完成门控池,在传输完成前防止复用。
设备修正区域是独立的。传输写入它,计算读取它,因此一个事件跨流对这些操作进行排序。如果下一次启动复用同一个设备区域,其传输也必须等待前一次计算完成读取。主机准备仍可与该计算重叠;只有当传输的目标与计算正在使用的内存无关时,传输才会与计算重叠。
启动时修补是当前的通用解决方案。其主机准备可以使用接下来描述的多流与设备计算重叠,并且当传输针对独立的设备存储时,传输也可以重叠。编译器拥有的地址和地址窗口绑定可能仍适用于特定情况,但运行时不会为普通张量输入在三种方案中进行选择。
未来的支持将允许运行时通过与地址相同的启动时修正机制,提供可变张量维度的实际大小。当编译时将这些维度保留为符号形式时,一个产物就可以在其支持的边界和布局约束内服务于多种输入大小,而无需为每种大小重新编译。
传输和计算使用独立的硬件流水线,但单个流仍可能让其中一个等待另一个。在单流序列中,主机准备修正数据,传输将其移动到设备,计算在下一次迭代开始前消费它。这是正确的,但它也会在设备计算时让主机和传输流水线等待。
在一次迭代内,主机回调必须在传输读取修正缓冲区之前完成,传输必须在计算使用修正后的程序之前完成。这防止传输读取不完整的缓冲区,也防止计算读取未修补的程序。跨迭代时,迭代 N+1 的准备可以在迭代 N 计算时开始,因为其地址在启动时已经已知。
运行时保持所需的顺序,并将独立工作移到另一个流上。准备流可以在设备流计算迭代 N 时为迭代 N+1 生成修正数据。一个事件在实际依赖处连接这些流:计算 N+1 等待其传输,但无关的主机准备不等待计算 N。当传输针对独立的设备存储时,传输也可以与计算重叠。复用同一个设备修正区域会增加一个反向依赖,因此该特定传输会等待较早的计算完成读取。
事件使用 PyTorch 读者已经熟悉的一种原语连接队列:一个流记录标记,另一个流等待它。准备流发出传输完成的信号;设备流在计算前等待该信号。顺序在重要之处保持全序,但表达为队列之间而非单个队列内部。
在运行时层,事件和同步操作与 CUDA 紧密对应:
| CUDA | 此处 | | --- | --- | | cudaEventRecord(event, stream) | 入队一个事件信号操作;只有当该流上已排队的所有内容都完成后,信号才会触发 | | cudaStreamWaitEvent(stream, event) | 入队一个事件等待操作;队列在该等待处停止,直到事件被发出信号,如果生产队列先终止,则直接失败 | | cudaStreamSynchronize() | 同步队列:阻塞直到其排空,并重新抛出记录在其上的第一个错误 |
在 CUDA 中,调用方显式地记录和等待。而在这里,运行时可以为那些设备内存读写足迹已知的操作推导出依赖边。框架层将准备工作与设备工作路由到相应的队列;当这些设备内存范围发生冲突时,运行时会插入信号和等待操作。程序与张量分配也参与该模型。主机暂存缓冲区的复用则遵循独立的生命周期规则,由同流完成顺序或完成门控池来保护。底层运行时支持显式的信号和等待操作,用于处理无法从设备内存足迹推断出的依赖关系。这些操作目前是内部的;通过 PyTorch 进行用户可调用的事件记录与等待已在计划中,但尚未实现。目前,用户可以通过 stream.synchronize() 或 torch.accelerator.synchronize() 等待完成,这些调用会阻塞主机线程。
初始方案采用生产者/消费者关系中的两个队列:一个承载主机工作及其传输的准备队列,一个承载计算的设备队列。该机制并不止步于两个队列,也没有强制要求这种拆分——如果一次启动没有主机工作要做,就完全不需要它。同样的原语可以扩展到任意数量的准备队列,因此没有理由串行化的主机工作可以并行进行;而且由于传输和计算是独立的流水线,针对独立设备存储的传输可以在当前迭代计算的同时进行。额外的准备队列必须使用独立的暂存缓冲区或完成门控池,以免一个队列覆盖另一个传输仍在读取的缓冲区。
随着启动时工作量的增长和设备计算的加快,主机准备工作变得更加显著。在一个队列上解析维度和地址,会将本可并发运行的工作串行化。
补丁并非唯一涉及的主机工作。以设备期望的布局准备张量也可以与设备计算并行运行,并通过事件连接到其消费者。
这些事件是由调度器检查的软件信号。它们是单次使用的,因此每条边都由一个新事件支撑。该接口以后可以使用硬件事件支持,而无需更改其上方的层。
用户接触到的接口是他们已经熟悉的。tensor.to("spyre")、torch.compile 和 torch.Stream 使用与其他设备相同的 PyTorch 概念。这并不会自动提供算子覆盖、性能分析、分布式集合通信、CUDA 特定扩展或模型特定的内核质量;这些仍然是后端的工作。
对于编译支持的即时执行操作,即时模式并不是第二套实现。Torch-spyre 围绕编译期间使用的相同分解注册了一个 PrivateUse1 内核。在 torch.compile 跟踪内部,包装器直接调用该分解,以便 PyTorch 将其捕获到周围的 FX 图中。在跟踪之外,首次即时调用会编译并缓存该分解;后续调用则复用已编译的入口点。因此,这些操作复用了相同的 Inductor 和启动机制,而不是在其旁边维护一个手写的运行时。
这种选择仍然会暴露首次使用时的编译;不支持的操作可能会失败或回退,不同的张量范围可能需要另一个编译产物。好处是只有一条执行路径,其修复和回归同时适用于两种模式。
编译支持的即时执行操作也是可能的最小编译程序:一个操作、真实的张量、真实的设备内存和真实的启动。如果模型失败但相应的即时操作通过,那么故障更可能出在组合上,而不是操作本身。针对单个操作的测试会锻炼整个编译与启动栈,因此逐操作的测试套件同时也是一个集成测试套件。
而且,由于驻留是通用的,而不是针对每种情况手写的,张量可以留在设备上,无需为每种张量编写专门的驻留代码。这正是即时模式变得实用的原因。
当设备的行为与普通 PyTorch 设备无异时,新的模型实现基本上就只是一个模型实现。它可以直接被拿来运行,而无需移植,因为其操作已经能够正确分发,其张量也已经位于正确的位置。这一点在更高一层同样成立。那些重构模型驱动方式的推理技术——比如推测解码(speculative decoding),即由一个小型草稿模型提出 token,再由一个更大的模型进行验证——都是基于流、张量驻留和普通控制流编写的。如果这些是框架自身的抽象,那么该技术就可以直接在设备上尝试,而无需向编译器添加任何内容。
编译器让算术运算变快,但推理栈的许多性能来自于工作如何被调度、批处理、重叠和推测。一个能够干净地映射到框架原语上的后端,将这片空间留给了那些从不接触编译器的人。
性能方面的论点在于启动路径上省去了什么。将类型化操作提交到队列,而不是构造图来描述它们,从每次启动中消除了磁盘读取、反序列化、逐参数拼接以及编译和解析。一次启动变成了构建几个操作并将它们入队。
在真实模型上,这种效果会叠加,因为这些成本原本是按每次启动来支付的。在 Granite 3.3 8B 上,批大小为 1、序列长度为 1024 时,仅将图构造从传输路径上移除,就在 prefill 上带来了 9.9% 的提升,在 decode 上带来了 13.6% 的提升——而这仅仅改变了字节的移动方式。将其从计算分发中也移除后,总计 prefill 快了 1.7 倍,decode 快了 2.4 倍,其中约四分之三的收益来自计算路径。
Mapping 4 中描述的并发是另一个机会:每一块从计算路径上移走的主机准备工作都减少了设备空闲时间。
PyTorch 有设备事件抽象,但上述跨流排序必须在操作实际所在的位置强制执行——即在运行时调度器中,这是能够看到其设备内存占用的层。因此,运行时首先构建软件事件原语和派生依赖。通过 PyTorch 的公共事件钩子暴露等效的用户创建事件则是单独的集成工作。
在框架的扩展点上构建并不是单向的交易。PyTorch 的核心开发者维护着 openreg,这是上游参考后端和加速器模拟器,展示了树外设备如何通过 PrivateUse1 注册而无需修改 PyTorch 核心,硬件生态合作伙伴也有贡献。在那里贡献改进,正是像 torch-spyre 这样的后端工作为下一个设备团队缩短路径的方式,而参与 PyTorch 技术咨询委员会则是让像这样的设备的约束传达给设计下一个扩展点的人的方式。我们与 PyTorch 的跨仓库 CI 中继(CRCR)的集成,针对 Spyre 测试上游变更并将结果报告到 PyTorch 的 CI 仪表板,帮助在框架演进过程中捕获兼容性回归。我们还构建了用于选择和适配上游测试的可复用工具,并计划将其贡献到上游,以便其他树外加速器也能受益。一个只消费接口的后端,不会让接口变得比它发现时更好。
好处在于,Spyre 和 PyTorch 都不需要假装成别的东西。Spyre 保留其数据流执行模型,而用户获得一个原生的 PyTorch 设备,具有熟悉的张量放置、流和编译。
对于 Spyre,PyTorch 的接口提供了通用的张量驻留、面向 eager 和编译执行的统一编译与启动路径、更低的启动开销,以及一种清晰的方式将主机侧准备与设备侧工作重叠。它们还减少了需要维护的后端特定机制。调度、存储生命周期和流管理遵循框架约定,使运行时更易于演进,也更不容易产生不一致的执行路径。
对于 PyTorch,Spyre 证明了这些扩展点可以支持并非围绕传统 GPU 执行模型组织的硬件。通过 openreg 和 PyTorch 技术咨询委员会反馈回来的经验,也使这些接口对下一个加速器后端更有用。最终结果是更健壮的 Spyre 运行时,以及一个能够容纳更广泛硬件、而无需为每个设备单独构建框架的 PyTorch 生态系统。
转载声明:本文为转载整理,原文版权归原作者及原站点所有。
原文来源:PyTorch Blog
原创地址:https://pytorch.org/blog/building-spyre-as-a-native-pytorch-device/
原文作者:Ebba Simpson;如有侵权,请联系管理员删除。