产品简介
TorchNeuron 是 AWS 于 2026 年 10 月 19 日发布、面向 AWS Trainium 芯片的原生 PyTorch 后端，并于 2026 年 10 月 20 日在 PyTorch Conference North America 以 Apache-2.0 协议开放公开 Beta。其代码仓库位于 github.com/aws-neuron/torch-neuronx。该工具通过 PyTorch 的 PrivateUse1 机制为上游 PyTorch 增加一个 neuron 设备，开发者只需将模型搬到 neuron 上即可进行训练或推理，并沿用已有的 PyTorch API。TorchNeuron 默认采用 eager 模式，算子逐个在设备上执行，无需先 trace 计算图；在需要性能时，可使用 torch.compile 的 neuron 后端对模型或模型的一部分进行编译。它与基于 XLA 的旧版 torch-neuronx 是两条不同的技术路径，本仓库定位为 PyTorch 原生后端。
核心功能
- 原生 PyTorch 设备接入：通过 PyTorch 的 PrivateUse1 机制为上游 PyTorch 增加 neuron 设备，将模型迁移到 neuron 上即可训练或推理，沿用已有 PyTorch API。
- 默认 eager 执行模式：默认采用 eager 模式，算子逐个在设备上执行，无需先 trace 计算图，便于调试与快速迭代。
- torch.compile 编译加速：需要性能时，可使用 torch.compile 的 neuron 后端编译模型或模型的一部分，在保留 PyTorch 编程体验的同时提升执行效率。
- 分布式训练支持：支持 torch.distributed、DTensor 与 functional collectives，覆盖数据并行、张量并行和全分片并行等分布式训练场景。
- NKI 自定义 kernel 调用：可从 PyTorch 代码中调用 NKI 自定义 kernel，在 eager 模式和 torch.compile 内均可用，便于针对特定算子进行底层优化。
- 预装运行环境：Neuron DLAMI 与 Deep Learning Container 已预装 TorchNeuron，其他 Linux 实例需先安装 Neuron 驱动与运行时；运行环境要求 Trn2 或 Trn3 实例，Python 3.12/3.13/3.14。
适合谁用
以下类型的用户会特别受益于 TorchNeuron：
- 在 AWS Trainium 芯片上使用原生 PyTorch API 进行模型训练
- 在 AWS Trainium 芯片上使用原生 PyTorch API 进行模型推理
- 利用 torch.compile 的 neuron 后端对模型或模型的一部分进行编译以提升性能
- 基于 torch.distributed、DTensor 与 functional collectives 开展数据并行、张量并行和全分片并行训练
- 在 PyTorch 代码中调用 NKI 自定义 kernel 进行算子级优化
优点亮点
- 以 Apache-2.0 协议开源，代码仓库公开可访问
- 通过 PrivateUse1 机制实现原生 PyTorch 设备接入，沿用已有 PyTorch API，迁移成本较低
- 默认 eager 模式无需先 trace 计算图，便于调试与快速验证
- 支持 torch.compile 的 neuron 后端，可按需对模型或模型的一部分进行编译
- 分布式训练覆盖数据并行、张量并行和全分片并行，支持 torch.distributed、DTensor 与 functional collectives
- 可从 PyTorch 代码调用 NKI 自定义 kernel，且在 eager 模式和 torch.compile 内均可用
价格与方案
TorchNeuron 以 Apache-2.0 协议开源，不单独收费，使用成本主要是 Trn2 或 Trn3 实例及配套 AWS 服务。Neuron DLAMI 与 Deep Learning Container 已预装，其他 Linux 实例需先安装 Neuron 驱动与运行时。具体实例价格以 AWS 官网为准。
总结与建议
TorchNeuron 是 AWS 面向 Trainium 芯片推出的原生 PyTorch 后端，以 Apache-2.0 协议开放公开 Beta。它通过 PrivateUse1 机制增加 neuron 设备，默认 eager 执行并支持 torch.compile 编译，分布式训练覆盖数据并行、张量并行和全分片并行，还可调用 NKI 自定义 kernel。运行环境要求 Trn2 或 Trn3 实例及指定 Python 版本，与基于 XLA 的旧版 torch-neuronx 是两条不同路径。
版本演进
- AWS 发布 TorchNeuron：Trainium 的原生 PyTorch 后端 （2026-10-19）：AWS 于 2026-10-19 发布 TorchNeuron，面向 AWS Trainium 的原生 PyTorch 后端，Apache-2.0 协议，2026-10-20 在 PyTorch Conference North America 开放公开 Beta。它通过 PyTorch 的 PrivateUse1 机制增加 neuron 设备，默认 eager 执行，可用 torch.compile 编译，支持 torch.distributed、DTensor 与 functional collectives 的分布式训练，并可调用 NKI 自定义 kernel。运行需 Trn2 或 Trn3 实例与 Python 3.12 到 3.14。