关键要点
- Amazon的TorchNeuron无需重写即可在Trainium上运行现有的PyTorch代码。
- TorchNeuron 将于 10 月 20 日在 PyTorch 大会上开源。
- 机器人、视频、世界模型、自动驾驶汽车、语音、药物发现、大语言模型（LLMs）以及金融领域的初创企业和企业已经在基于此进行构建。
如今，85% 的 AI 研究人员和开发人员使用 PyTorch 编写代码。在此之前，这些代码可运行的硬件选择远比代码本身有限，而扩展到新的芯片意味着需要重写并重新验证那些耗时数月才能完善的工作。
Amazon正在缩小这一差距。借助TorchNeuron，开发人员已编写的PyTorch代码只需进行一项配置更改且无需重写，即可在Amazon的AWS Trainium芯片上原生运行。TorchNeuron将于10月20日在北美PyTorch大会上以开源许可证形式发布。
所谓的“无需更改代码”究竟是什么意思？
将你的 AI 代码想象成一辆你花费多年时间精心调校的汽车。引擎、悬挂系统——所有部件协同工作。现在出现了一款更实惠且马力强劲的新引擎，但安装它意味着要拆解汽车并逐块重建。TorchNeuron 让你无需拆解汽车即可更换新引擎。
在一项并列演示中，Amazon工程师在同一款传统芯片和Trainium上运行了相同的AI训练脚本。唯一的区别是一个告知代码使用哪种芯片的设置。结果显示，AI模型以相同的速度学习了相同的内容。
对开发人员而言，这意义重大。PyTorch 已经内置了新硬件注册自身的机制。TorchNeuron 接入了该标准机制，这意味着开发人员用于测试、调试和扩展其 AI 模型的现有工具均可在 Trainium 上直接运行，无需修改。
为什么 TorchNeuron 对开发人员和行业至关重要
AI的真实成本远超芯片价格。对许多公司（尤其是初创企业）而言，最昂贵的支出项目是工程时间。开发人员每周花费在为新硬件适配软件上的时间，都是未能用于改进产品的时间。对于资源有限的小型团队来说，这种权衡可能决定一家公司的成败。
这不仅仅是初创企业的问题。对于运营大规模AI集群的大型企业而言，能够使用同一硬件同时执行训练和推理工作负载是一项变革性优势。在低峰时段，原本服务于推理工作负载的容量可以无缝转移到训练任务中，从而从同一硬件中获得更多价值，而不是让专用硬件闲置。
通过降低采用新硬件的摩擦，TorchNeuron 引入了 AI 行业长期缺乏的东西：有意义的选择。当开发人员可以无代价地在不同芯片之间切换时，芯片制造商必须赢得开发人员的业务，而开发人员则受益于更低的成本和更快的迭代速度。
从语音 AI 到药物发现：谁已经在基于 TorchNeuron 进行构建
TorchNeuron 的早期采用者涵盖了广泛的 AI 应用领域。从实时语音到分子设计，它们有一个共同点：它们在迁移到 Trainium 时没有重写代码。
Cartesia 正在构建能够像人类一样从世界中学习并与之交互的智能，从对话开始。借助 Trainium 上的原生 PyTorch 支持，Cartesia 对其标准推理管道进行了极少量的 PyTorch 代码更改即可运行。应用场景也延伸至医疗保健领域。Lieber 大脑发育研究所是致力于针对精神分裂症等主要神经精神疾病开发新医疗干预措施的最大研究机构，它在药物靶点和药物发现中使用了生物、化学和语言 AI 模型。分子动力学模拟功能强大，但计算成本极高，新一代 emerging 的 AI 工具有望加速或取代它们。TorchNeuron 允许 Lieber 研究所在 Trainium 上探索这些策略，从初步的模型探索到针对瓶颈优化性能。
开放芯片生态系统
亚马逊的方法不仅限于其自有硬件。该公司正在向 PyTorch 本身贡献上游更改，与 PyTorch 社区合作移除将框架绑定到特定芯片的代码。目标是使 PyTorch 对整个开发者社区（而不仅仅是 Trainium 用户）更加硬件无关化。
为了保持质量，亚马逊在每个构建版本中都会使用数十万个 AI 模型对 TorchNeuron 进行测试，从而在任何客户工作负载运行之前捕获问题。
Hugging Face 已在其流行的 AI 开发工具中直接启用了原生 Trainium 支持，使开发人员无需任何亚马逊特定代码即可在 Trainium 上进行原生开发。
TorchNeuron 与 PyTorch 的原生分布式训练框架 TorchTitan 集成，为开发人员提供了一条标准路径，以便使用完全分片数据并行（FSDP）在 Trainium 实例上扩展模型训练——这与他们在任何其他 PyTorch 设备上使用的工具和模式相同。
对于希望深入探索的开发人员，TorchNeuron 提供了跨多个芯片扩展训练的工具，以及编写自定义低级代码以实现最大性能的能力。亚马逊还向 vLLM（一种用于加速推理的流行开源工具）做出了重大开源贡献，并开源了其完整的性能内核集，为开发人员构建自己的解决方案提供了起点。
下一步计划
TorchNeuron支持Trainium2和Trainium3。它将于10月20日在北美PyTorch大会上以Apache 2.0许可证形式作为开源软件发布。
TorchNeuron是Amazon十年来在定制硅片领域持续投资的最新篇章。自2015年收购芯片设计公司Annapurna Labs以来，Amazon从零开始构建了四个芯片系列，每个系列都专为特定的云工作负载而设计。借助TorchNeuron，该公司现在正使任何使用PyTorch进行开发的开发人员都能访问这些硬件。
了解更多关于TorchNeuron的信息：AWS Trainium的原生PyTorch后端，以及为何AI初创企业选择AWS Trainium来训练世界模型。