AWS 开源 TorchNeuron：Trainium 成了 PyTorch 原生设备，训练代码改成 .to('neuron') 就能跑，但约 2.2 万个上游测试还没覆盖

从「适配 Neuron」到「换个设备名」

在 AWS 自研的 Trainium 芯片上跑模型，过去要用 Neuron 专用的软件和配置去改模型代码，Hugging Face 在声明里也是这么描述的。AWS 团队在 PyTorch 官方博客发文，10 月 20 日起以 Apache 2.0 开源 TorchNeuron（仓库 aws-neuron/torch-neuronx）并开放公测，可从 AWS Neuron 的 pip 源用 pip install torch-neuronx 安装。

它走的是 PyTorch 官方给第三方硬件留的 PrivateUse1 扩展点，把 Trainium 注册成 neuron 设备。官方示例里，训练循环唯一要改的就是 .to('neuron')；torch.nn、torch.autocast、torch.distributed、torch.profiler 按原样工作，Transformers v5 用 device_map="neuron" 就能加载模型。支持 Trainium2 和 Trainium3。

两种模式、一个算子接口

eager 模式逐个执行算子，Python 控制流和断点都能用，适合研究与调试；新加的自适应 eager（实验性）尝试在不改代码的前提下合并算子、减少启动开销。上生产用 torch.compile 指定 neuron 后端，经 TorchDynamo、AOTAutograd、Torch-MLIR 转成 StableHLO 再交给 Neuron 编译器。需要更底层控制的，可以用 NKI 写自定义算子，官方开源的 NKI-lib 提供了 FlashAttention、MoE 等常用算子。分布式方面支持 DDP、FSDP2、张量并行和流水线并行，vLLM 也有 Neuron 插件。

官方自己列出的缺口

博客没有回避覆盖率：8 月跑 PyTorch 上游测试套件约 15.4 万个结果，跳过约 6.9 万个。其中约 3 万个是 PyTorch 自身判定不适用（如仅限 CUDA 或 CPU），约 3.9 万个是 Neuron 主动跳过，最大一块约 2.2 万个是相关算子还没有原生实现。稀疏张量、NestedTensor、前向模式自动微分、动态与符号形状首版不支持，TorchAO 和 FP8 仍在私测。

客户数据目前只有厂商渠道的说法：Deepgram CTO 称，在他们评测的场景里，Trainium 每美元处理的输入音频时长比原 GPU 基线多约 60%。对已经写好 PyTorch 代码、想比较 GPU 与 Trainium 成本的团队，现在可以用自己的模型直接试；用到动态形状或稀疏算子的，先等后续版本。