PyTorch 给 Trainium 上了原生后端：TorchNeuron 不再走 XLA，代码只改 .to('neuron')

PyTorch 官方博客 10 月 19 日发布《Introducing TorchNeuron》，把 AWS Trainium 注册成 PyTorch 的一等设备：走 PrivateUse1 这个官方扩展点，挂上 neuron 设备后，torch.nn、autocast、distributed、profiler 均按原样工作。改造成本被压到一行 model = MyModel().to('neuron')，编译模式只需先调一次 torch.compiler.set_default_backend('neuron')。官方称以 Apache 2.0 在 GitHub 开源，10 月 20 日进入公开 Beta，覆盖 Trainium2 与 Trainium3。

它改变的是 Trainium 的软件入口。此前这条线走 XLA（torch-neuronx），模型要先过一层图编译，生态库要么适配要么绕开；原生后端意味着按 PyTorch API 写的代码——Hugging Face Transformers、TorchTitan、vLLM——不改就能落到 Trainium 上。这也是 AWS 在软件栈上的一次取舍：让出自建编译器的隔离层，换取 PyTorch 生态的直接接入。

配套给了 vLLM Neuron 与 vLLM Omni Neuron 两个插件（分别服务大模型推理和图像视频扩散）、NKI 内核接口（NKI-lib 已开源 FlashAttention 与 MoE blockwise 内核），以及用同一套 torch.profiler 从框架追到硬件的 Neuron Explorer。分布式侧 DDP、FSDP2、张量并行与流水并行照常，集合通信走 Trainium 的 collective core（节点内）与 EFA（跨节点）。

边界要写清。这条记的是公告，不是「已经能用」：截稿时 github.com/aws-neuron/torch-neuronx 在 GitHub 上返回 404，公开 Beta 落在官方所说的 10 月 20 日（美西时间，约北京时间 21 日凌晨）。它也不是全新项目——AWS 文档记载 TorchNeuron 早在 2025 年 12 月 re:Invent 公布、随 Neuron 2.27.0 首发，此前只对受邀客户开私有 Beta；这一步是开源加公开。另有两处仍是私有 Beta：TorchAO 与 FP8，官方称下个版本才普遍可用；整篇博客一个性能数字都没给，只有接口与架构。

原文链接：PyTorch 官方博客