亚马逊推出开源工具TorchNeuron,开发者仅需修改单行配置,即可将PyTorch代码无缝迁移至自研Trainium芯片上运行。此举直接挑战英伟达CUDA生态护城河,大幅降低AI硬件切换门槛。
亚马逊发布了一款名为 TorchNeuron 的软件工具。开发者只需修改一项配置,无需重写任何代码,即可将原生 PyTorch 代码直接运行在亚马逊自研的 Trainium AI 芯片上。该工具由亚马逊旗下芯片部门 Annapurna Labs 打造,直接瞄准英伟达凭借软件生态建立的算力壁垒。
TorchNeuron 此前自去年底开展小范围测试,现已面向所有开发者开启公测。亚马逊还计划在北美 PyTorch 大会上,以 Apache 2.0 协议将该软件完全开源,全面支持 Trainium2 及后续芯片架构。

此前,开发者若想在 Trainium 上运行 PyTorch 代码,必须重构代码以调用专门针对亚马逊芯片设计的 API。Annapurna Labs 高管 Colin Brace 表示,TorchNeuron 的目标是让开发者“只要懂 PyTorch,就能直接上手 Trainium”。
TorchNeuron 接入了 PyTorch 内置的硬件注册机制。这意味着开发者现有的测试、调试和扩缩容工具链均可在 Trainium 上无缝复用,同时兼容 PyTorch 的分布式训练框架 TorchTitan。
亚马逊内部每晚会针对 Hugging Face 上的 10 万多个模型对该工具进行自动化测试。此外,亚马逊还向 PyTorch 主干提交了代码修改,推动移除了原有与特定芯片绑定的逻辑。
Brace 透露,目前已有超过 30 家涵盖机器人、世界模型和语音 AI 领域的早期客户采用该软件。语音 AI 初创公司 Deepgram 表示,基准测试显示 Trainium 在每美元处理音频量上比其现有的 GPU 方案提升了约 60%。
由 Meta 主导的 PyTorch 是目前绝大多数 AI 开发者的首选框架。英伟达深耕十余年构建了 CUDA 生态,确保 PyTorch 能够在其 GPU 上高效运行。过去,将代码迁移至非英伟达硬件通常需要耗费数月的工程重构成本,导致大量企业对硬件迁移望而却步。
其他云巨头也在推进类似策略:Google 正在开发用于在其自研 TPU 上原生运行 PyTorch 的 TorchTPU,微软近期也宣布计划将其大部分 AI 服务迁移至自研 Maia 芯片上。
目前,Trainium 芯片的核心客户主要包括与亚马逊有投资合作的 Anthropic 和 OpenAI 等顶尖 AI 机构。亚马逊 CEO 安迪·贾西(Andy Jassy)此前透露,亚马逊自研芯片业务的年化营收已突破 200 亿美元。Brace 表示,TorchNeuron 将成为 Trainium 走向更广泛开发者生态的“加速器”,为行业提供真正可替代的算力选择。
免费获取企业 AI 成熟度诊断报告,发现转型机会
关注公众号

扫码关注,获取最新 AI 资讯
3 步完成企业诊断,获取专属转型建议
已有 200+ 企业完成诊断