The science behind TorchNeuron, Amazon’s new PyTorch backend for AWS Trainium chips

TorchNeuron背后的技术原理：亚马逊为AWS Trainium芯片推出的新型PyTorch后端 - 亚马逊科学

云计算与系统

TorchNeuron背后的技术原理：亚马逊为AWS Trainium芯片推出的新型PyTorch后端

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

作者：

Robert Hundt

2026年10月22日

7分钟阅读

分享

• 复制链接

• 邮件

• X

• 领英

• Facebook

• Line

• Reddit

• QZone

• 新浪微博

• 微信

• WhatsApp

分享到微信

x

PyTorch是大多数AI开发者构建模型时使用的框架。它是一个用于张量（矩阵或二维数字网格的高维类比）操作的Python库。现代AI模型的关键特性，如输入数据、模型权重和注意力机制，几乎都通过张量表示，而PyTorch显著简化了对它们的操作代码编写。

本周，我们宣布推出TorchNeuron，这是一个开源后端，可将PyTorch代码转换为可在AWS Trainium系统上执行的操作。目前在其他加速器上运行的PyTorch代码，只需进行最小或无需代码修改即可在Trainium上运行。

TorchNeuron连接了PyTorch和Trainium，双方都从中受益。PyTorch开发者现在可以将代码运行在专为大规模张量工作负载设计的Trainium上。他们还可以使用AWS Neuron SDK，其中包含一个能够以周期级精度追踪每条指令的性能分析器，以及Neuron内核接口（NKI，发音为"Nikki"），该接口几乎可以实现汇编级别的控制。

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

TorchNeuron包含多项优化功能，使PyTorch代码能够在Trainium上高效运行，包括多级缓存机制和可选的"自适应即时模式"。该模式不会像普通即时模式那样逐个执行操作，而是将操作分组并编译以实现更高效的执行，接近完全编译的性能水平。

然而，TorchNeuron的许多创新之处在于其性能评估方式。AI代理搜索了数十万开源模型，并根据其架构特征对它们进行分类。这使Khaled ElGalaind和Sarang Ozarde领导的Annapurna团队能够识别出841个代表性的模型，涵盖了所有模型架构的完整光谱。从这些模型中，他们还提炼出数万个具有实际意义的张量维度（或"形状"），这些维度随后被用于调整和基准测试TorchNeuron的性能。结果是，PyTorch程序员可以信任TorchNeuron能够准确、高效地执行他们的代码，并且最重要的是，开箱即用。

运行时优化

自适应即时模式

标准PyTorch有两种原生执行模式。一种称为"即时模式"，意味着PyTorch操作会按照程序顺序立即执行。

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

TorchNeuron 在两者之间引入了中间方案，即一种可选设置称为自适应即时模式，它兼具编译模式的部分优势，但作用于更小的子图。此模式在保持即时模式响应性的同时，也实现了编译模式的性能表现，且无需承担编译时间开销。

运行时，运行在CPU上的Python解释器会收集从PyTorch操作符调度中获取的指令并构建小型图。这些图会被发送到TorchNeuron进行编译、缓存并放入执行队列。多个图集合和编译过程会并行进行，进一步减少编译时间开销。

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

在即时模式下运行PyTorch代码的开发者习惯于将硬件上执行的操作与代码中的具体行对应起来；这就是即时模式对调试如此有用的原因。我们不希望破坏这种体验，因此TorchNeuron的自适应即时模式保持了精确的源码对应关系。

缓存

当编译后的图从队列中取出时，TorchNeuron会将其转换为名为NEFF（Neuron可执行文件格式）的可执行文件。即时模式中的单个操作可以生成单个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模块、注意力机制和前馈层，其中注意力机制可能是自注意力、交叉注意力或两者兼具，具体形式包括闪存注意力、环形注意力和增量注意力等；后端则将模型的内部表示转换为可用的输出形式，可能是任何类型的数据，通过任意数量的策略生成。

在所有模型（包括LLM和其他类型）中，我们发现仅有841个代表性模型涵盖了TorchNeuron在实际应用中可能遇到的所有设计特征。

但即使在841个模型上完全评估新代码版本，对于快速回归检测来说仍然耗时过长。因此，每次评估运行时，我们都会记录Trainium需要处理的不同规模的张量操作。最终，我们建立了一个"形状注册表"，记录了数以万计的张量操作及其维度，这些正是TorchNeuron在实践中可能遇到的操作。

基于该注册表，我们生成了单元测试用例，能够以显著加快的速度测试和验证新软件版本。任何通过快速单元测试的版本都会进一步在841个代表性模型上进行测试。如果通过这些测试，还会在完整的模型分布上再次验证。这种分阶段测试大幅提升了开发效率，因为问题可以在开发周期早期就被发现。

所有这些测试和广泛的模型覆盖使我们确信，TorchNeuron能够在Trainium上准确高效地执行PyTorch模型，为客户提供更多选择，帮助他们应对日益拥挤的AI技术环境。

研究领域

• 云计算与系统

标签

• 硬件

关于作者

Robert Hundt 是亚马逊网络服务（AWS）的副总裁兼杰出工程师。