Amazon 发布 TorchNeuron：面向 AWS Trainium 的 PyTorch 后端

云与系统

Amazon 面向 AWS Trainium 芯片的新 PyTorch 后端 TorchNeuron 背后的科学

运行时优化以及由智能体式 AI 驱动的全面评估流程，有助于确保用 PyTorch 编写的模型能在 Trainium 系统上高效、可靠地执行。

本周，我们宣布 TorchNeuron 正式可用。这是一个开源后端，可将 PyTorch 代码转换为可在 AWS Trainium 系统上执行的操作。已在其他加速器上运行的 PyTorch 代码只需极少改动或无需改动即可在 Trainium 上运行。

TorchNeuron 在 PyTorch 与 Trainium 之间架起桥梁，双方都从中受益。PyTorch 开发者现在可以在 Trainium 上运行他们的代码，而 Trainium 正是为大规模张量工作负载从头设计的。他们还可以使用 AWS Neuron SDK，其中包含一个性能分析器，能以周期级精度追踪每条指令，以及 Neuron 内核接口（NKI，读作“Nikki”），它几乎能在汇编级别实现控制。

TorchNeuron 实现了 PyTorch 模型、AWS Neuron SDK 与 AWS Trainium 芯片之间的无缝集成。

TorchNeuron 包含多项优化，使 PyTorch 代码能在 Trainium 上高效运行，包括多级缓存，以及一个可选的“adaptive eager”模式。该模式不像普通 eager 模式那样逐个执行单个操作，而是将操作分组并编译，以实现更高效的执行，性能接近完全编译模式。

然而，TorchNeuron 背后的许多创新在于其性能的评估方式。AI 智能体搜索了数十万个开源模型，并根据其架构特征对它们进行分类。这使得由 Khaled ElGalaind 和 Sarang Ozarde 领导的 Annapurna 团队得以识别出 841 个代表性模型，覆盖了模型架构的完整谱系。他们还从这些模型中提炼出数万个实际重要的张量维度，即“shape”，随后用于调优和基准测试 TorchNeuron 的性能。其结果是，PyTorch 程序员可以放心地让 TorchNeuron 准确、高效地执行他们的代码，而且最重要的是，开箱即用。

运行时优化

Adaptive eager 模式

标准 PyTorch 有两种原生执行模式。一种称为“eager”，即 PyTorch 操作按程序顺序立即执行。

另一种是编译模式，程序员选择应按传统意义编译的代码块，即由编译器构建张量操作的控制流图并对其进行优化。目前已有许多优化技术，例如算子融合，以减少内存流量并提高局部性。开发者通常在调试模型时使用 eager 模式，而在确信代码正确并希望其尽可能高效执行时使用编译模式。

TorchNeuron 增加了一种介于两者之间的方式，即一个名为 adaptive eager 模式的可选设置，它提供编译的部分优势，但作用于更小的子图。该模式兼顾了 eager 模式的响应性和编译模式的性能，同时没有编译时间的开销。

在运行时，运行在 CPU 上的 Python 解释器会累积从 PyTorch 算子分发中获得的指令，并构建小型图。这些图被发送到 TorchNeuron，由其编译、缓存并放入执行队列。多个这样的图收集和编译过程并行进行，进一步将编译时间开销降至最低。

这里的一个关键创新在于，累积的图、其中的各个操作及其输入张量 shape 与编译时间直接相关。更大的图作用于更大的张量，编译时间更长。因此，一种智能的自适应算法会动态选择合适的子图，以最大化编译时间与已编译代码片段执行的重叠，从而为多种模型有效隐藏全部编译时间。

以 eager 模式运行代码的 PyTorch 开发者习惯于将硬件上执行的操作与单行代码对应起来；这正是 eager 模式对调试如此有用的原因。我们不想破坏这种体验，因此 TorchNeuron 的 adaptive eager 模式保持了准确的源码对应关系。

缓存

当编译后的图从队列中出来时，TorchNeuron 会将其转换为一种名为 NEFF 的可执行文件，即 Neuron 可执行文件格式。eager 模式下的单个操作可以生成单个 NEFF，同样的机制也适用于完整图或部分图。从这个意义上说，Neuron 软件栈默认采用提前（AOT）编译。

将操作转换为 NEFF 的成本很高，因此 TorchNeuron 会尽量避免不必要的转换。它生成的每个 NEFF 都会在多个层级进行缓存。内存缓存保留最近生成的 NEFF，而磁盘上的持久缓存则使 NEFF 能够被复用于快速重启和后续工作。默认采用 AOT 的方式让这一切变得简单直接。

Amazon EC2 UltraServer 最多由四台独立的 Graviton 服务器组成，根据配置不同，包含 16 到 144 个 Trainium 加速器芯片。对于非常大的模型，即使一台 UltraServer 也不足以容纳完整模型，模型会被分片到多台服务器上，涉及数千个加速器芯片。由于它们都运行同一模型的相同部分，因此都需要相同的 NEFF。

TorchNeuron 运行时协调这些加速器，使它们不必各自编译相同的代码；相反，由一个系统编译操作和图，其他系统则复用编译后的代码。一个共享编译服务将编译器池化，为服务器上的所有加速器服务，从而避免主机 CPU 不堪重负。在将大型语言模型分布到多个加速器上的测试中，这些优化将首个训练步骤的时间缩短了约三分之一。

评估

我们构建了 AI 智能体，对 Hugging Face 仓库进行了广泛扫描，排除了因配置损坏而无法在任何加速器上运行的模型，并筛选出数十万个具有可运行、可执行代码的模型。

在所有那些模型上评估 TorchNeuron 代码的每一个修订版本将极其耗时。然而，而且此外，其中许多模型执行相似的任务或处理相似类型的数据。

因此，我们团队的成员构建了另一个智能体，根据一组定义的架构属性对模型进行分类。例如，大型语言模型有三个主要组件：前端通常是编码器，针对不同数据类型（如文本、音频、图像、视频或其某种组合）进行专门化。主干包含主要的 Transformer 块、注意力机制和前馈层，但注意力机制可能是自注意力、交叉注意力或两者兼有，并且有多种变体，如 flash attention、ring attention 和 delta attention。最后，后端将模型的内部表示转换为可用的输出形式，这可能是任何类型的数据，通过任意数量的策略生成。

在所有模型中，无论是 LLM 还是其他模型，我们发现仅 841 个代表性模型就覆盖了 TorchNeuron 在实际中可能遇到的所有设计特征。

但即使在 841 个模型上完整评估新的代码修订版本，对于快速回归检测来说仍可能耗时过长，因此在每次评估运行时，我们都会记录 Trainium 被要求处理的不同大小的张量操作。最终，我们得到了一个“形状注册表”，编目了数万个操作及其张量维度，这些是 TorchNeuron 在实践中可能遇到的。

从该注册表中，我们生成了单元测试，能够以显著更快的方式测试和验证新的软件版本。任何通过快速单元测试的修订版本随后会针对 841 个代表性模型进行测试。如果也通过了，我们会在完整的模型分布上再次测试。这种分阶段方法极大地提高了开发速度，因为问题可以在开发生命周期中更早地被发现。

所有这些测试以及对模型的广泛覆盖让我们相信，TorchNeuron 将在 Trainium 上准确高效地执行 PyTorch 模型，在客户面对日益拥挤的 AI 格局时为他们提供更多选择。

Robert Hundt 是 Amazon Web Services（AWS）的副总裁兼杰出工程师。