AI-translated from English; not yet reviewed by a fluent editor.

# AWS 的 TorchNeuron 为 Trainium 带来原生 PyTorch 后端

> AWS 的 Neuron 2.33 版本为 Trainium2 和 Trainium3 引入了处于公开测试阶段的 PyTorch 后端。支持的版本、算子覆盖范围和编译限制，决定了工程师在评估时需要检查哪些内容。

By BIG CHANGE Editorial

Published: 2026-10-20T04:20:56.038Z
Updated: 2026-10-20T04:20:56.038Z
Canonical: https://bigchange.ai/blog/aws-torchneuron-native-pytorch-backend-trainium

![Conceptual charcoal illustration of a code-marked monitor beside an open rack drawer containing an unbranded accelerator board.](https://bigchange.ai/api/media/file/torchneuron-native-device-trainium-hero-v1.png)
AI-generated conceptual illustration by BIG CHANGE.

AWS Neuron 2.33 的[发行说明](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html)将 TorchNeuron 公开测试版的日期标为 10 月 19 日。同样发布于 10 月 19 日的[PyTorch 公告](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/)则称，公开测试版和开源版本于 10 月 20 日上线。`torch-neuronx` 软件包在 PyTorch 中为 Trainium2 和 Trainium3 注册了一个 `neuron` 设备，让受支持的模型和张量操作可以使用熟悉的 PyTorch API。对于考虑采用这些芯片的工程师，有用的变化是：从现有 PyTorch 程序到设备层面试用的路径缩短了。程序能否正确且高效地运行，仍取决于其操作、张量形状和软件环境。

此前，PyTorch 在 Trainium 上使用基于 XLA 的 `torch-neuronx` 库及 `torch_xla` 设备代码。AWS 表示，原生后端使用 PyTorch 的 `PrivateUse1` 扩展，并将 `xm.xla_device()` 替换为 `neuron` 设备放置方式。对于已经使用 Neuron 的项目，这是实质性的接口变化，尽管软件包名称依然熟悉。该版本适用于 Trn2 和 Trn3 实例上的训练与推理；支持矩阵并未列出 Trn1。

## PyTorch 程序发生了哪些变化

在[AWS 发行说明](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html)中，示例使用 `.to("neuron")` 移动模型，并在 `device="neuron"` 上创建输入张量。受支持的 `torch.nn`、autograd 和优化器操作随后通过该后端运行。使用 PyTorch 与设备无关的 `torch.accelerator` API 编写的代码可以发现 Neuron，而无需针对设备重写，参见[TorchNeuron 概述](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html)。

有两条执行路径。Eager 模式在 Python 调用时执行 PyTorch 操作，保留普通控制流、打印输出和断点。每个操作首次使用时会编译，后续匹配的调用可以使用缓存。工程师也可以用 `torch.compile(model, backend="neuron", dynamic=False)` 编译模型或选定区域。原生后端会编译受支持的前向和反向计算图；捕获区域之外的代码继续以 eager 模式运行。[PyTorch 公告](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/)还介绍了在编译前一次性设置 `torch.compiler.set_default_backend("neuron")` 的方法，这样之后的编译调用就可以省略后端参数。

该版本提供了一个 `neuron` 后端，用于 `torch.distributed`，并记录了集合通信和点对点操作。AWS 将 DistributedDataParallel、FSDP2、DTensor、DeviceMesh 和流水线并行列为受支持的 API。对于分布式训练项目，这些路径值得考察，但发行说明并不能证明现有多节点任务无需修改就能运行，或能达到某个特定的吞吐量目标。

## 评估从匹配的软件环境开始

AWS 的[安装快速入门](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/get-started/quickstart-torch-neuron.html)要求使用 Trn2 或 Trn3 实例，以及 Neuron Base DLAMI 或受支持的 Ubuntu AMI。该版本的矩阵指定了运行于 Ubuntu 24.04 或 26.04 的 Linux x86\_64。Ubuntu 22.04 对此版本的支持已经结束。矩阵列出了 Python 3.12、3.13 和 3.14，并要求为 PyTorch 2.13.0 配套使用 `torch-neuronx` 2.13.3.0，或者为 PyTorch 2.14.0 配套使用 `torch-neuronx` 2.14.3.0。软件包系列必须与 PyTorch 版本匹配。

快速入门从 Neuron 和 PyTorch 软件包索引中安装 `torch-neuronx==2.14.*`，同时安装 `neuronx-cc==2.*` 和 `nki`。随后会检查设备可用性，并将 Neuron 上的一个小型矩阵操作与 CPU 结果进行比较。这验证的是 AWS 示例中已安装设备的执行路径，并不能验证读者自己的模型。合理的下一步评估是运行实际模型的前向和反向路径，使用 `torch_neuronx.get_fallback_ops()` 检查记录的 CPU 回退操作，将输出与当前环境进行比较，并分别测量编译和稳态运行。这些后续步骤是 BIG CHANGE 提出的评估问题，并非我们执行过的测试。

## 测试版仍可能怎样影响工作

在 eager 模式下，不受支持的操作可能回退到 CPU，并发出警告。在 `torch.compile` 下，如果计算图中捕获的操作无法由 Neuron 后端降级转换，默认情况下编译就会失败。AWS 建议开发者查阅其[算子兼容性表](https://github.com/aws-neuron/torch-neuronx/blob/main/docs/torch_api_compatibility.md)，了解每项操作的覆盖情况。因此，一次成功完成的 eager 运行，本身并不能证明全部操作都在设备上执行，也不能证明编译模式兼容。

发行说明还指出，`torch.compile` 要求 `dynamic=False`；新的输入形状可能需要再次编译。AWS 建议通过填充或按形状分桶来减少形状变化。如果缓存中没有匹配的产物，首次运行会包含编译过程。为工作负载计时时，这些细节很重要：编译延迟和输入形状的分布模式可能改变评估结果。

其他明确的限制进一步缩小了兼容性声明的范围。该测试版不支持 `torch.sparse`、`torch.func`、`torch.nested` 或 `torch.distributions`。只有部分 `torch.fft` 变换可在 Neuron 上运行，而 `float64` 会以 `float32` 执行，并发出警告。此版本没有集成 TorchAO，也不支持复合 MXFP8 数据类型。使用这些功能的团队，需要先检查自己的代码路径，再把设备放置方式视为迁移计划。

### 重大变化

TorchNeuron 在公开测试版中，将 Trainium 接入了 PyTorch 的原生设备和编译接口。这使受支持的 PyTorch 程序更便于进行首次试用，也为现有 Neuron 用户提供了不同于基于 XLA 代码的路径。是否采用它，决定性证据要来自各个工作负载：算子是否留在预期设备上，数值结果是否仍可接受，以及测得的性能和成本是否足以证明实例和迁移工作值得投入。BIG CHANGE 为本次分析审阅了文档和公告；我们没有运行 TorchNeuron，也没有独立验证 AWS 的性能声明。

## 来源与延伸阅读

- [AWS Neuron 2.33 TorchNeuron 发行说明](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html)：发布日期、测试版状态、版本和操作系统矩阵、API 及已知限制。
- [AWS TorchNeuron 概述](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html)：设备集成、eager 与编译执行、CPU 回退行为，以及通过 Neuron 的软件路径。
- [AWS 安装与验证快速入门](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/get-started/quickstart-torch-neuron.html)：实例和镜像前提条件、软件包安装，以及文档中的设备检查。
- [PyTorch 博客：介绍 TorchNeuron](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/)：从 PyTorch 角度解释设备切换、编译和分布式 API。该文章是 AWS 贡献者在 PyTorch 博客上发布的公告，并非独立基准测试。
- [Amazon 公告](https://www.aboutamazon.com/news/innovation-at-amazon/torchneuron-pytorch-aws-trainium-chips-ai)：AWS 对其目标和早期采用者的说明。关于迁移、性能和经济性的声明，需要针对具体工作负载进行验证。

## Sources

- [AWS Neuron 发行说明：TorchNeuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchneuron.html) — 一手来源；将公开测试版日期标为 2026-10-19，比 PyTorch 公告中所称的公开测试版与开源版本上线日期早一天；包含版本矩阵和明确限制。
- [什么是 TorchNeuron？](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html) — 关于设备与后端使用的一手文档。
- [PyTorch 博客：TorchNeuron，面向 AWS Trainium 的原生 PyTorch 后端](https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/) — AWS 贡献者从 PyTorch 角度发布的公告；日期为 10 月 19 日，并称公开测试版和开源版本于 10 月 20 日上线。当前 URL 通过 Amazon 文章访问；并非独立性能验证。
- [Amazon 公告](https://www.aboutamazon.com/news/innovation-at-amazon/torchneuron-pytorch-aws-trainium-chips-ai) — 供应商公告；迁移和经济性声明出自 AWS，未经过独立验证。
- [安装并验证 torch-neuronx](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/get-started/quickstart-torch-neuron.html) — AWS 快速入门，涵盖前提条件、匹配的软件包安装，以及文档中的 CPU/Neuron 设备检查。
- [TorchNeuron API 兼容性表](https://github.com/aws-neuron/torch-neuronx/blob/main/docs/torch_api_compatibility.md) — 维护者按 API 列出的原生执行、CPU 回退和不支持分类。