AWS Neuron 2.33 的发行说明将 TorchNeuron 公开测试版的日期标为 10 月 19 日。同样发布于 10 月 19 日的PyTorch 公告则称，公开测试版和开源版本于 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 发行说明中，示例使用 .to("neuron") 移动模型，并在 device="neuron" 上创建输入张量。受支持的 torch.nn、autograd 和优化器操作随后通过该后端运行。使用 PyTorch 与设备无关的 torch.accelerator API 编写的代码可以发现 Neuron，而无需针对设备重写，参见TorchNeuron 概述。

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

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

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

AWS 的安装快速入门要求使用 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 建议开发者查阅其算子兼容性表，了解每项操作的覆盖情况。因此，一次成功完成的 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 发行说明：发布日期、测试版状态、版本和操作系统矩阵、API 及已知限制。
  • AWS TorchNeuron 概述：设备集成、eager 与编译执行、CPU 回退行为，以及通过 Neuron 的软件路径。
  • AWS 安装与验证快速入门：实例和镜像前提条件、软件包安装，以及文档中的设备检查。
  • PyTorch 博客：介绍 TorchNeuron：从 PyTorch 角度解释设备切换、编译和分布式 API。该文章是 AWS 贡献者在 PyTorch 博客上发布的公告，并非独立基准测试。
  • Amazon 公告：AWS 对其目标和早期采用者的说明。关于迁移、性能和经济性的声明，需要针对具体工作负载进行验证。