尝试一款新 AI 芯片曾经是一件很费劲的事。你往往需要重写代码、重新验证，耗费数周时间，赌上这个新芯片就是你工作负载的正确选择。

这既费时间又费钱，还让团队疲惫不堪。你最优秀的人员被淹没在底层适配工作里，而不是去构建他们本想构建的东西。

TorchNeuron 将在很大程度上解决这一问题，只需一次配置更改、无需重写，就能让你现有的 PyTorch 代码原生运行在 Trainium 上。

凭借 Trainium 极具吸引力的性价比，TorchNeuron 将让更多公司更容易地以更低成本使用 AI。https://aboutamazon.com/news/innovation-at-amazon/torchneuron-pytorch-aws-trainium-chips-ai…

Amazon 的 TorchNeuron 让 PyTorch 运行在 AWS Trainium 芯片上

来源：https://www.aboutamazon.com/news/innovation-at-amazon/torchneuron-pytorch-aws-trainium-chips-ai

核心要点

• Amazon 的 TorchNeuron 无需重写即可在 Trainium (https://www.aboutamazon.com/stories/ai-chips-aws-Trainium2-explain) 上运行现有的 PyTorch 代码。
• TorchNeuron 将于 10 月 20 日在 PyTorch Conference 上开源。
• 机器人、视频、世界模型、自动驾驶、语音、药物发现、LLM 和金融领域的初创企业和大型企业已经在其上进行构建。

如今，85% 的 AI 研究人员和开发者使用 PyTorch 编写代码。此前，这些代码能运行的硬件远比代码本身受限得多，扩展到新芯片意味着要重写和重新验证那些耗费数月才打磨完善的工作。

Amazon 正在缩小这一差距。通过 TorchNeuron (https://github.com/aws-neuron/torch-neuronx)，开发者已有的 PyTorch 代码只需一次配置更改、无需重写，即可原生运行在 Amazon 的 AWS Trainium 芯片 (https://www.aboutamazon.com/news/aws/aws-trainium-graviton-ai-chips-explained) 上。TorchNeuron 将于 10 月 20 日在 PyTorch Conference North America 上以开源许可发布。

Amazon Network AWS (https://www.aboutamazon.com/stories/inside-aws-networking-lab-cupertino)

“无需改代码“究竟意味着什么？

把你的 AI 代码想象成一辆你花数年时间精心调校的汽车。引擎、悬挂——每个部件都协同工作。现在来了一台更便宜、动力强劲的新引擎，但安装它意味着要把整辆车拆开再逐个零件重建。TorchNeuron 让你无需拆解汽车就能换上新引擎。

在一次对比演示中，Amazon 工程师在传统芯片和 Trainium (https://aws.amazon.com/ai/machine-learning/trainium/) 上运行了相同的 AI 训练脚本。唯一的区别是一个告诉代码使用哪块芯片的设置项。结果表明，AI 模型以相同的速度学到了相同的内容。

对开发者而言，这意义重大。PyTorch 内置了让新硬件自我注册的机制。TorchNeuron 正是接入这一标准机制，这意味着开发者已有的用于测试、调试和扩展 AI 模型的工具都能在 Trainium 上直接使用，无需修改。

Amazon AI 创新场景，包括量子计算硬件、数据中心内部、工程师团队协作以及定制芯片 (https://www.aboutamazon.com/news/aws/amazon-ai-chips-business-history)

TorchNeuron 对开发者和行业为何重要

AI (https://www.aboutamazon.com/what-we-do/artificial-intelligence-ai) 的真实成本远不止芯片本身的价格。对于许多公司，尤其是初创企业 (https://www.aboutamazon.com/news/aws/why-ai-startups-choose-amazon-trainium-chips)，最昂贵的开支是工程时间。开发者每花一周适配新硬件，就意味着少花一周改进产品。对于资源有限的小团队来说，这种取舍可能直接决定公司的成败。

而且不仅仅是初创企业。对于运行大规模 AI 集群的大型企业，能够在同一硬件上兼顾训练和推理 (https://www.aboutamazon.com/news/aws/ai-chip-terms-explained) 是一个改变游戏规则的能力。在非高峰时段，服务推理工作负载的容量可以无缝转移到训练任务，让同一硬件创造更多价值，而不是让专用硬件闲置。

通过降低采用新硬件的门槛，TorchNeuron 带来了 AI 行业长期缺乏的东西：真正的选择权。当开发者可以无痛地在芯片之间迁移时，芯片制造商就必须凭实力赢得开发者的选择，而开发者则受益于更低的成本和更快的迭代速度。

从语音 AI 到药物发现：谁已经在 TorchNeuron 上构建

TorchNeuron 发布——面向 AWS Trainium 的原生 PyTorch 后端

TorchNeuron 的早期采用者涵盖广泛的 AI 应用领域。从实时语音到分子设计，它们有一个共同点：迁移到 Trainium 时无需重写代码。

Cartesia 正在打造能够像人类一样学习和与世界交互的智能体，从对话开始。借助 Trainium 上的 PyTorch 原生支持，Cartesia 几乎无需修改 PyTorch 代码即可运行其标准推理流程。其应用场景也延伸到了医疗健康领域。Lieber Institute for Brain Development 是一家专门致力于为精神分裂症等重大精神疾病开发新医学干预手段的最大研究机构，他们在药物靶点和药物发现中使用生物、化学和语言 AI 模型。分子动力学模拟虽然强大但计算成本极高，而新一代 AI 工具有望加速甚至替代它们。TorchNeuron 让 Lieber Institute 得以在 Trainium 上探索这些策略，从初始模型探索到针对性能瓶颈进行优化。

Amazon Bio Discovery (https://www.aboutamazon.com/news/aws/aws-amazon-bio-discovery-ai-drug-research)

开放芯片生态系统

Amazon 的策略不止于自身硬件。该公司正将变更贡献给 PyTorch 上游，与 PyTorch 社区合作移除将框架绑定到特定芯片的代码。目标是让 PyTorch 对整个开发者社区更加硬件无关，而不仅仅是 Trainium 用户。

为了保证质量，Amazon 在每次构建时都会用数十万个 AI 模型对 TorchNeuron 进行测试，确保在任何客户工作负载运行之前就捕获问题。

Hugging Face 已在其广受欢迎的 AI 开发工具中直接启用了 Trainium 原生支持，让开发者无需编写任何 Amazon 专属代码即可在 Trainium 上原生工作。

TorchNeuron 与 TorchTitan（PyTorch 的原生分布式训练框架）集成，为开发者提供了一条标准路径，可通过全分片数据并行（FSDP）在 Trainium 实例上扩展模型训练——使用的是他们在任何其他 PyTorch 设备上所使用的相同工具和模式。

对于想要深入探索的开发者，TorchNeuron 提供了跨多芯片扩展训练以及编写自定义底层代码以获得极致性能的工具。Amazon 还向广受欢迎的开源推理加速工具 vLLM 做出了一项重大的开源贡献，公开了其完整的性能内核套件，为开发者构建自己的内核提供了起点。

下一步计划

TorchNeuron 支持 Trainium2 (https://aws.amazon.com/ec2/instance-types/trn2/) 和 Trainium3 (https://www.aboutamazon.com/news/aws/trainium-3-ultraserver-faster-ai-training-lower-cost)。它将于 10 月 20 日在 PyTorch Conference North America (https://events.linuxfoundation.org/pytorch-conference-north-america/) 上以 Apache 2.0 许可证开源发布。

TorchNeuron 是 Amazon 十年来在定制芯片 (https://www.aboutamazon.com/news/aws/amazon-ai-chips-business-history) 领域持续投入的最新篇章。自 2015 年收购芯片设计公司 Annapurna Labs 以来，Amazon 已从零开始构建了四大芯片系列 (https://www.aboutamazon.com/what-we-do/amazon-custom-chips)，每一款都为特定的云端 (https://www.aboutamazon.com/news/tag/cloud-computing) 工作负载而设计。通过 TorchNeuron，该公司正在让所有使用 PyTorch 的开发者都能便捷地使用这些硬件。