返回行业动态
企业AI落地

AWS 开源 TorchNeuron:PyTorch 代码改一处设备名就能跑 Trainium

AWS 10 月 20 日以 Apache 2.0 开源 TorchNeuron,让已有 PyTorch 代码在 Trainium2 与 Trainium3 上原生运行,通常只需把设备名改成 neuron。Hugging Face 的 Transformers、Diffusers、PEFT、TRL 已把 Trainium 列为原生设备。AWS 同时公开缺口:约 154,000 项 PyTorch 测试结果中约 69,000 项被跳过,动态输入形状尚不支持。此举直接压缩 CUDA 生态的迁移成本。

来源:Amazon Web Services 官方博客(AWS Insights)原标题:Your last hardware decision shouldn't decide what you build next查看原文
工程师工作台,双屏显示芯片互联拓扑,桌面摆放一块深色 AI 加速卡与线缆,背景机房虚化

图片来源:北京拓实科技原创(AI 生成封面)

发生了什么

AWS 在 10 月 20 日的 PyTorch Conference North America 上把 TorchNeuron 以公开测试版发布。这是一个原生 PyTorch 后端,让已有模型代码在 Trainium2 与 Trainium3 上直接运行,通常只需改一处设备名,不必重写模型或重做验证。代码以 Apache 2.0 许可开源在 aws-neuron/torch-neuronx 仓库,相关改动也已提交回 PyTorch 上游。

AWS 给出的理由是迁移成本:按其说法,85% 的 AI 开发者用 PyTorch 写代码,但换一颗芯片往往意味着数月重写与重新校验,这笔开销从不出现在报价单上。

具体交付了什么

  • 支持 eager 执行(便于逐步调试)、torch.compile(面向生产性能)与 Neuron Kernel Interface(自定义算子);
  • 通过 PyTorch 的 PrivateUse1 设备后端与 torch.accelerator 接口接入,按现代写法编写的代码无需改动;
  • 与 Hugging Face 合作,Transformers、Diffusers、PEFT、TRL 已把 Trainium 作为原生设备,官方称 100 个最常用的文本生成架构开箱可用;
  • 在 Trainium 上训练出的 checkpoint 是标准格式,可迁回其他受支持硬件,选择保持双向。

测试覆盖的缺口

AWS 在技术博客里少见地公开了测试口径:2026 年 8 月的一次 PyTorch 自身测试套件运行中,约 154,000 项结果里有约 69,000 项被跳过;其中约 30,000 项本身只针对 CUDA 或 CPU,约 39,000 项属 Neuron 侧主动跳过,最大的一块(约 22,000 项)是尚无 Trainium 原生实现的算子。

首版不支持动态输入形状、稀疏张量与嵌套张量;经 TorchAO 的 FP8 支持仍在私有测试,自适应 eager 模式被标为实验性。这意味着做变长序列推理的团队需要先跑一遍自己的模型,而不是直接照搬厂商基准。

商业侧的算盘

对 AWS 而言,TorchNeuron 打的是 CUDA 的护城河——真正的锁定往往不在芯片速度,而在围绕 CUDA 积累的库、调试习惯与训练配方。AWS 首席执行官 Matt Garman 的说法是「开发者应该按价格、性能和可获得性选硬件,而不是被已有的代码锁住」。首发客户包括语音领域的 Deepgram 与 Cartesia、脑科学研究的 Lieber Institute。

Deepgram 自评:在同一批基准上,Trainium 每美元处理的音频时长比其原有 GPU 方案多约 60%,但未披露对比的是哪一代 GPU,且这一数字是数月迁移之后的结果。AWS 另称其自研芯片业务年化收入已达 250 亿美元。

不确定之处

性能证据基本来自合作方自评,尚无独立复测;支持范围写成「Trainium」时覆盖多种芯片设计,但官方明确只针对 Trainium2 与 Trainium3,早期芯片不在列;动态形状缺失意味着相当一部分线上推理负载短期内仍无法直接迁移。对计划做国产或非英伟达替代方案评估的企业,TorchNeuron 降低的是试点成本,不等于降低迁移风险。

编辑说明

本文由北京拓实科技根据公开信息整理,核心事实与数据请以原始来源为准。

AWS 开源 TorchNeuron:PyTorch 代码迁移 Trainium 只需一行|北京拓实科技