AREX Feed Article
Amazon 推出 TorchNeuron:PyTorch 代码一处配置更改即可原生运行 AWS Trainium,定于 10 月 20 日开源
10 月 19 日,Amazon 在宣布推出 TorchNeuron:开发者已有的 PyTorch 代码只需一处配置更改、无需改写,即可原生运行在 Amazon 自研的 AI 芯片 AWS Trainium 上。这套软件定于 10 月 20 日在硅谷举行的 PyTorch Conference North America 大会上以 Apache 2.0 许可证开源发布,支持 Trainium2 与 Trainium3 两代芯片。
把 TorchNeuron 的推出定性为意在「侵蚀支撑 NVIDIA 主导地位的优势」。Amazon 芯片部门 Annapurna Labs 软件副总裁 Colin Brace 告诉路透社,此举旨在减少开发者试用自研芯片所需的步骤——「到客户所在之处迎接他们」。他称已有 30 余家早期客户,并表示切换到 Amazon 的芯片可带来显著的成本节省。
TorchNeuron 接入的是 PyTorch 预留的硬件注册机制
Amazon 用一张四层架构图说明这条路径:最上层是开发者的 PyTorch 代码,向下依次是 TorchNeuron 原生后端、AWS Neuron SDK 与编译器,最底层是 Trainium 芯片;图中 PyTorch 层的示意代码写明了那处更改落在哪——model = model.to('neuron'),把模型搬进 'neuron' 设备。
技术上的关键是 PyTorch 本身就为新硬件预留了注册机制。Amazon 称 TorchNeuron 接入的正是这套标准机制,因此开发者原本用来测试、调试和扩展模型的一整套工具,无需修改就能在 Trainium 上运行。博客还给出一场对照演示:Amazon 工程师把同一份 AI 训练脚本分别跑在常规芯片和 Trainium 上,唯一差别是告诉代码使用哪块芯片的那个设置;按 Amazon 的说法,结果显示模型「以相同的速率学到了相同的东西」。
Amazon 在博客里把原有迁移方式比作给一辆调校多年的车换引擎:即便新引擎更便宜、马力更大,通常也得把车拆开逐件重装;TorchNeuron 的方案是不拆车、直接换引擎。
面向大规模训练,TorchNeuron 集成了 PyTorch 原生分布式训练框架 TorchTitan,开发者可以沿用 Fully Sharded Data Parallelism(全分片数据并行,FSDP)的工具与模式,把训练扩展到多台 Trainium 实例上;需要更深一层优化的开发者,还能用它提供的工具做多芯片扩展、编写底层定制代码。
为什么是 PyTorch:主流框架与 CUDA 的十余年积累
PyTorch 是由 Meta 托管的开源框架,路透社称大多数 AI 开发者都在使用;Amazon 给出的数字是,85% 的 AI 研究者和开发者用 PyTorch 写代码。NVIDIA 则花了十余年时间打造 CUDA 软件栈,让 PyTorch 代码在其芯片上高效运行——在路透社的表述里,NVIDIA 主导地位所倚仗的正是这套积累。
路透社还指出,同样拥有自研 AI 芯片的 Google 等云计算厂商,也在用更好的 PyTorch 工具争夺 NVIDIA 的开发者。
早期客户:从对话式语音 AI 到药物发现
Brace 口中的早期客户横跨机器人、世界模型和语音 AI;Amazon 博客的要点栏把范围列得更宽:机器人、视频、世界模型、自动驾驶、语音、药物发现、大语言模型(LLM)和金融领域的初创公司与企业都已在 TorchNeuron 上构建。博客具体展开了两家。
Cartesia 构建以类人方式学习和与世界交互的智能,从对话做起。借助 Trainium 上的 PyTorch 原生支持,这家公司只做了极小的改动,就跑通了自己标准的推理流水线。
Lieber Institute for Brain Development 的场景更偏研究。按 Amazon 博客的介绍,这是致力于为精神分裂症等重大神经精神疾病开发新医学干预手段的最大研究机构;研究所把生物、化学和语言 AI 模型用于药物靶点与药物发现,而传统的分子动力学模拟虽然强大、计算开销却极大,新一代 AI 工具有望加速甚至替代它们。借助 TorchNeuron,研究所得以在 Trainium 上探索这些路线,从最初的模型探索一直做到针对瓶颈的性能优化。
上述客户数量、成本节省的幅度与对照演示的结果,均出自 Amazon 一方口径,尚无独立性能基准或第三方验证。
向 PyTorch 上游贡献「去芯片绑定」的改动
Amazon 的动作没有停在自家工具上。按官方博客的说法,公司正在向 PyTorch 上游贡献改动,与社区合作移除把框架绑在特定芯片上的代码,目标是让 PyTorch 对整个开发者社区——而不只是 Trainium 用户——更加硬件无关(hardware-agnostic)。
生态一侧,Hugging Face 已在其流行的 AI 开发工具中启用原生 Trainium 支持,开发者无需写入任何 Amazon 专有代码;Amazon 还向流行的开源推理加速工具 vLLM 做了一项重大开源贡献,并开源了整套性能内核(performance kernels)。质量方面,Amazon 称每次构建都会用数十万个 AI 模型测试 TorchNeuron,让问题在任何客户工作负载运行之前先被抓住。
Amazon 算的账:工程时间和闲下来的算力
按 Amazon 博客的说法,AI 的真实成本不止芯片价格。对许多公司、尤其是初创公司,最贵的一行开支是工程师的时间——花一周把软件适配到新硬件上,就少一周改进产品;对资源有限的小团队,这笔账足以影响成败。对运行大规模 AI 集群的大企业,Amazon 的卖点是训练与推理共用同一批硬件:非高峰时段,原本服务推理负载的算力可以平滑转给训练任务,而不是让专用硬件闲置。
这笔账背后是 Amazon 十年的自研芯片积累:2015 年收购芯片设计公司 Annapurna Labs 之后,Amazon 从零建起四个芯片家族,各自对应一类云端工作负载。TorchNeuron 是把这套硬件开放给所有 PyTorch 开发者的最新一步。
截至发稿,Apache 2.0 许可证下的 TorchNeuron 代码尚未发布,按计划它将在 10 月 20 日的 PyTorch Conference North America 大会现场开源。届时,「无需改写」在真实工作负载下的边界,才轮得到 Amazon 之外的检验。