精选项目
前沿模型涵盖语言、视频和音频，其架构从稠密与混合专家 Transformer 到扩散模型、状态空间模型、世界模型和结构预测模型。PyTorch 是 AI 研究发生的地方，是定义模型的地方，也日益成为生产发生的地方。研究人员希望进行实验和创新，客户则希望 AWS Trainium 提供的性能能帮助他们实现性能和规模目标。训练和服务库都是基于 PyTorch 自身的 API 编写的，新库和新版本持续发布。过去几年里，PyTorch 使这些 API 与加速器无关，并提供了用于注册自定义硬件的官方扩展点。
今天，为了满足开发者与研究社区对 Trainium 日益增长的需求，我们很高兴在 10 月 20 日于北美 PyTorch Conference 开源发布之前，分享一种更简单的方法，将芯片作为标准 PyTorch 设备进行编程。AWS Trainium 是一系列 AI 芯片，专为下一代云规模 AI 训练和推理从零设计。我们宣布推出 TorchNeuron，这是一个面向 Trainium 的开源原生 PyTorch 后端，使你已在使用的模型代码、库和工具能够在 Trainium 上运行。该集成直接构建在 PyTorch 社区创建的与加速器无关的扩展点之上，包括 PrivateUse1 设备后端和 torch.accelerator API。
TorchNeuron 在 Trainium2 和 Trainium3 上支持用于研究和调试的即时执行（eager execution），包括性能提升的全新自适应即时模式、用于生产性能的 torch.compile，以及用于自定义内核的 Neuron Kernel Interface（NKI）。此外，借助 Neuron Explorer，性能工程师使用相同的 torch.profiler 来追踪从框架到硬件的执行过程，将每个事件关联回其源代码行，从而精准定位瓶颈并高效优化。TorchNeuron 将以 Apache 2.0 许可证在 github.com/aws-neuron/torch-neuronx 开源发布，并于 10 月 20 日作为公开测试版（public Beta）上线。
“PyTorch 的愿景很明确：开发者应在不同硬件平台上获得一致的体验。对 AWS Trainium 的原生支持扩展了这一选择，使研究人员和工程师能够使用他们已熟悉的即时执行和分布式 API 进行实验和迭代。团队可以充分利用 Trainium 的能力，而无需重塑其工作流程，也无需离开 PyTorch 生态系统。” – Jana van Greunen，Meta PyTorch 工程总监
开始使用：轻松切换设备
入门很简单。使用指针在 Trainium 上运行你的代码：.to('neuron')。
基于 `torch.accelerator`（PyTorch 推荐的设备无关 API）编写的现代代码可无需修改运行。要编译模型，需要指定 Neuron 后端：通过 `torch.compiler.set_default_backend("neuron")` 一次性指定，之后 `torch.compile(model)` 会采用它；或者在每次调用时以 `torch.compile(model, backend="neuron")` 指定。
能力
TorchNeuron 通过 PrivateUse1 注册 Trainium，PrivateUse1 是 PyTorch 面向自定义硬件的官方扩展点，因此标准 API 可按原样工作：torch.nn、torch.tensor、torch.autocast、torch.distributed 和 torch.profiler 都能在 neuron 设备上运行。PyTorch 的两种执行模式——eager 和 torch.compile——均可原生工作。
TorchNeuron 的高层架构图（详情）。
Eager 会立即运行操作，因此 Python 控制流、打印和断点均可在无需编译或跟踪步骤的情况下正常工作。这是用于研究和调试的模式。Eager 模式会逐个分派操作。直接执行时，每个操作都会编译自己的 Neuron 可执行文件格式（NEFF）文件，在设备上启动它，并通过 HBM 移动其中间结果。我们还将引入自适应 Eager 模式（实验性），以降低启动成本，并提供更好的操作融合机会，从而在不更改代码的情况下接近编译图的性能，同时保留交互式体验。
torch.compile
torch.compile(model, backend="neuron") 会将模型捕获为 Neuron 编译器可优化的图，是稳定、高性能模型投入生产时所采用的模式。torch.compile 可用于优化整个模型或模型的一部分，并与 eager 模式无缝协作。原生支持 FlexAttention。TorchDynamo 将模型捕获为 FX 图，AOTAutograd 提前追踪反向图，Neuron 后端通过 Torch-MLIR 将结果降为 StableHLO，Neuron 编译器生成可在硬件上执行的 NEFF 文件。
Neuron 推理栈构建在 TorchNeuron 中相同的 torch.compile 支持之上。vLLM Neuron 是 vLLM 的插件，用于为 LLM 提供服务；vLLM Omni Neuron 是 vLLM Omni 的插件，用于为图像和视频生成提供扩散模型服务。
运算和数据类型
广泛的 aten 算子覆盖，涵盖文本、多模态与视觉语言、图像和视频生成与理解、状态空间模型以及世界模型中出现的算子。模型可使用 BF16、FP32 和 FP16，以及整数和布尔类型，torch.autocast 和混合精度代码可按原样运行。TorchAO 和 FP8 目前以私有 Beta 形式提供支持，并将在下一版本中正式开放。
分布式 API
分布式推理和训练可使用标准 PyTorch API，高层并行模块可与按原样编写的模型组合使用：DDP、FSDP2、SimpleFSDP、张量并行和流水线并行，以及对于训练大模型至关重要的常见组合，如 FSDP2 与张量并行、ZeRO。支持 DTensor 和 DeviceMesh，以处理分片和多维进程组。支持集合通信 API（例如：all_gather、reduce_scatter、all_reduce）和 P2P API（send、recv）。集合通信在节点内通过 Trainium 专用集合通信核心执行，在节点间通过 Elastic Fabric Adapter (EFA) 执行，因此通信传输是硬件原生的，无需更改 PyTorch API。
流与事件
原生支持 Torch 流和事件。torch_neuronx.Stream 和 torch_neuronx.Event 将 PyTorch 的流语义映射到 Trainium 的硬件队列上。在一个流内，数据移动和计算会自动保持顺序。不同流之间没有隐式依赖关系，因此在跨流时，你需要通过在生产者流上记录事件，并在消费者流上等待该事件后再使用结果，来显式建立这种顺序。与设备无关的 torch.Stream、torch.Event 和 torch.accelerator 入口点会到达同一实现，因此针对它们编写的代码无需修改即可运行。错误会在下一个同步点浮现，而不是在导致该错误的具体操作处出现；TorchNeuron 会跟踪最早的异常，并在该处使用其原始上下文重新抛出。在执行时，非默认流会复用到单个硬件队列上，因此它们会保留顺序、优先级和事件语义，但不会各自获得专用引擎；并且当主机集合通信被禁用时，在非默认流上发出的集合通信会保留其在流执行顺序中的位置，但运行在默认硬件队列上。
Neuron 内核接口 (NKI)
Neuron 内核接口 (NKI) 使研究人员和开发者能够为 eager 和 torch.compile 开发、优化并执行自定义算子。研究人员可以 prototyping 需要底层硬件控制的想法，包括新的模型架构。开发者可以通过在模型代码中直接调用性能关键的 NKI 内核来替代标准 PyTorch 操作（ATen 算子）序列，从而获得超越 Adaptive Eager 和 torch.compile 的更精细性能控制。
TorchNeuron 还提供由 NKI 实现并以 NKI-lib 开源发布的 Trainium 优化高性能算子，覆盖稠密模型和 MoE 模型的主流 LLM 与多模态场景，例如 FlashAttention 和 MoE 分块内核，开发者可用其在 Trainium 上提升性能。以下示例使用 PyTorch 推荐的 register_autograd 风格定义并注册自定义 NKI 算子：
使用 TorchNeuron 进行性能分析
TorchNeuron 与标准 torch.profiler.profile API 集成，并扩展了 NeuronConfig 对象，用于控制 Neuron 特定数据的收集。这使用户能够捕获并关联来自框架级 CPU 操作（PyTorch 调度、Python 栈）以及 Neuron 运行时和 Trainium 活动的跟踪记录（由 NeuronConfig 控制），从而实现端到端性能分析。
捕获的性能分析可在 Neuron Explorer UI 中可视化，用户可以在其中跟踪跨框架、运行时和设备的执行过程，从任意事件跳转到准确的源代码行（源代码查看器），通过调用层级导航工作负载（层级查看器），并沿依赖链查看每个停顿正在等待什么（依赖链查看器）。
Neuron Explorer 还会输出包含逐算子和逐引擎指标的 Parquet 输出，因此 UI 背后的同一数据可直接查询，或由智能体用于自动化瓶颈分析。对于程序化访问，用户可使用内置的 torch.profiler 工具。例如，key_averages 表让用户可以深入查看特定算子，并使用 sort_by 参数显示每个算子的设备端开销，这对于识别热点很有用。
以下代码片段展示了一个示例（详情）：
有关更多能力和功能，请参阅 Neuron Explorer 文档。
基于 TorchNeuron 构建的客户与合作伙伴
TorchNeuron 已被初创公司、企业、机器学习行业合作伙伴和学术机构使用：涵盖自动驾驶、机器人、视频生成、语音、电信、医疗保健和地理空间建模。以下是各团队如今正在构建的内容概览。
“Cartesia 正在构建能够像人类一样从世界中学习并与之交互的智能，并从对话开始。我们在 Mamba 等状态空间架构上的开创性工作，为 Sonic（文本转语音）和 Ink（语音转文本）提供支持，这些语音 AI 模型支撑着真正像人类一样的对话体验。真实对话同时要求三点：质量（自然度和准确性）、效率（大规模下的低延迟）和可控性。大多数解决方案只能满足其中一两点。我们从第一性原理出发设计研究和产品，以同时实现这三点，这意味着我们需要能够在不重写整个技术栈的情况下，从每一层榨取最大性能的基础设施。借助 AWS Trainium 上的 PyTorch 原生支持，我们只需对 PyTorch 代码做极少更改，就能运行标准推理流水线。小型、快速的操作通过原生路径高效映射，因此我们不必为它们手写内核。这让我们团队能够专注于为那些被反复调用并决定端到端延迟的层编写高性能 NKI 内核。对于实时语音而言，每一毫秒的网络距离都至关重要。在 Trainium 上运行有助于让 Sonic 和 Ink 贴近全球客户。”
Arjun Desai，Cartesia 联合创始人
“语音正在成为 AI 的一等模态，Deepgram 每年为超过 200,000 名开发者和 1,400 个组织处理数千万亿音频 token。为满足这一需求，我们需要高效扩展，同时继续交付更好的模型。在为期数月的合作中，AWS 生成式 AI 创新中心帮助我们将其模型移植到 AWS Trainium 上的 PyTorch 原生支持。在我们评估的基准测试中，Trainium 每美元可处理的输入音频秒数比我们现有的 GPU 基线多约 60%，这有助于我们以更低价格向客户提供更好的模型。我们期待在我们的解决方案中支持未来的 AWS AI 芯片。”
Adam Sypniewski，Deepgram CTO
“Hugging Face 是面向 AI 构建者的领先开放平台，通过由 PyTorch 驱动的开源库（包括 Transformers 和 Diffusers），免费提供对超过 600 万个开放模型、数据集和应用的访问。在此之前，要利用 AWS Trainium，需要使用 Neuron 特定的软件和配置来适配 Hugging Face 模型；TorchNeuron 使 Trainium 成为 Transformers 和 Diffusers 的原生后端，因此开发者可以通过完全开源的栈使用熟悉的模型定义——这是我们致力于芯片中立性和兼容性的实际体现。如今，在 AWS Trainium 上部署最先进的开放大语言模型以及最新的开放图像和视频生成模型比以往任何时候都更容易；TRL 支持微调和强化学习，Neuron 内核通过我们的库自动提供优化性能。”
Jeff Boudier，Hugging Face 产品副总裁
除了这些客户之外，我们还一直与许多正在使用 TorchNeuron 构建并增强 Trainium 生态系统的合作伙伴合作。这些合作伙伴包括 Anyscale、Decart、Kernelize、Loka、Luminal、Makora、Modular、Nota AI、Radix、Wafer 和 Yasp。
TorchNeuron 如何实现更广泛的兼容性
我们实现原生 PyTorch 支持的目标之一，是让开发者能够利用 Trainium 的性能和成本优势，同时继续使用他们已经熟悉的模型、库和工作流。这有助于保护他们在 PyTorch 上的投资，包括其模型代码、推理和训练基础设施、调试实践，以及他们所依赖的社区项目。
运行 AI 工作负载的成本包括让其正常运行、优化其性能，以及在模型和库演进过程中维护它们所需的工程时间。我们希望开发者花更多时间试验架构和改进模型质量，花更少时间让软件适配底层硬件。他们应当能够从熟悉的模型或训练配方开始，将其适配到自己的工作负载，并以最少的 Neuron 特定更改在 Trainium 上运行。
TorchNeuron 通过兼容更广泛的 PyTorch，为这一体验提供了基础。开发者可以使用 TorchTitan 等库，该库汇集了 PyTorch 的分布式和编译能力以支持大规模训练；还可以使用 Hugging Face Transformers，它提供对社区模型实现、预训练检查点以及用于将模型适配到其应用的 API 的访问。
TorchTitan 示例
TorchTitan 是 PyTorch 的原生库，用于快速实验和大规模训练。其模块化设计让定制训练、测试新想法和扩展实验变得更加容易。TorchTitan 目前可在 Trainium 上进行训练。我们持续将改进上游合入主线 TorchTitan，而 torchtitan-neuron 会在这些改进合入之前就为你提供它们，并在每一项改进落地时逐步向主线收敛。Trainium 专用的制品也随 torchtitan-neuron 一同发布，并通过 TorchTitan 自身的扩展点注册；这些制品包括 NKI 内核、面向 Trainium 调优的模型配置和训练配方，以及面向 Neuron Explorer 的性能分析集成。
以下是使用 TorchTitan，在 8 台 trn2.48xlarge 上以张量并行（TP）和混合分片数据并行（HSDP）训练 Llama 3 8B 的示例损失曲线
Hugging Face Transformers v5
Transformers v5 为互操作性而构建：一套模型实现，可在任何平台和任何硬件上运行，并可被下游训练和服务栈直接使用。
Trainium 是 Hugging Face 各库上游支持的设备，包括 Transformers、Diffusers、trl 和 kernels。Transformers v5 支持文本和音频模型的训练与推理，Diffusers 支持图像和视频生成，trl 支持大型语言模型的强化学习，NKI 内核可通过 Hugging Face kernels 库获取。AWS 和 Hugging Face 为每个库添加了支持，因此现有的 Hugging Face 工作流无需任何设备专用代码即可在 Trainium 上运行。
训练。Transformers v5 增强了完整预训练和微调的就绪程度，具备更好的初始化，并与并行化和工具生态系统兼容——TorchTitan 用于规模化，Unsloth 和 TRL 用于后训练。
推理。Transformers v5 在各模型实现之间统一了注意力和内核。
扩散。Diffusers 中面向 Neuron 的张量并行已合入上游，因此图像和视频生成可通过 diffusers API 原样运行。
兼容性。Transformers v5 是 TorchTitan、vLLM 和 SGLang 等运行时库的原生后端，将其模型覆盖范围扩展到 Hugging Face 上约 300 万个公开模型中的大多数。我们正在努力在 Neuron 上启用并验证这一集成，将扩展后的模型支持带到 Trainium 实例。
TorchNeuron 通过 PyTorch 的 PrivateUse1 扩展点注册 Trainium，这使 Neuron 进入 PyTorch 的设备枚举。device_map="auto" 无需任何配置即可找到它，device_map="neuron" 则显式指定它。无论哪种方式，权重都会加载到设备上进行训练或推理，脚本的其余部分保持不变。
测试方法与结果
我们的测试方法包括在 HuggingFace 和 TorchTitan 模型上对 TorchNeuron 进行广泛的模型覆盖验证，以及在 eager 和 torch.compile 两种模式下运行 PyTorch 上游测试套件。
模型覆盖
盲扫（大规模广泛测试）。我们使用 Hugging Face 模型库对开源模型进行广泛扫描。我们持续运行 10 万个 Hugging Face 模型，覆盖 99.9% 的真实使用场景（基于下载量）。我们在 Trainium 上以缩减但具代表性的配置运行每个模型，并记录发生的情况：它能否加载、前向和反向传播能否运行、能否编译，以及哪些算子在该设备上执行、哪些回退到 CPU。
下图展示了截至 2026 年 9 月，在 eager 和 torch.compile 两种模式下运行 10 万个模型的结果。我们正在努力弥合剩余的差距，目标是在未来版本中弥合大部分乃至全部差距。