AWS推出TorchNeuron,为Trainium3带来原生PyTorch支持
要点速览
- •AWS推出了面向Trainium2和Trainium3加速器的原生PyTorch后端TorchNeuron,于2026年10月19日进入公开测试,并于次日在北美PyTorch大会上开源。
- •只需一行代码即可将模型迁移至Trainium,并可在即时执行模式下使用标准PyTorch API,取代了此前需要在训练循环中散布xm.step()等调用、依赖繁重编译器工作流的XLA方案。
- •Neuron SDK 2.33.0增加了多项配套升级,包括用于检测硬件健康问题的Neuron Observer、支持底层定制的改进NKI功能,以及经调优适配TorchNeuron的更新版vLLM插件。
- •AWS表示已有超过30家早期客户在机器人、语音AI和大语言模型开发领域使用TorchNeuron,其中Cartesia和Deepgram报告性能显著提升,Hugging Face也已采用Trainium支持。
- •TorchNeuron仅支持Trainium2和Trainium3,使用Trainium1或Inferentia芯片的团队必须继续使用早期基于XLA的工具。

亚马逊云服务(AWS)推出了TorchNeuron,这是面向其Trainium2和Trainium3 AI加速器的原生PyTorch后端,旨在让研究训练像其他任何PyTorch任务一样顺畅。此举解决了长期存在的痛点:Trainium芯片在运行已完成的AI模型方面表现出色,但在构建新模型这一繁杂且充满实验性的工作上却一直不够顺手。AWS现在希望弥补这一差距。这些加速器是该公司面向机器学习工作负载的自研芯片产品线的一部分。
AWS究竟发布了什么
TorchNeuron于2026年10月19日进入公开测试阶段。AWS在一天后的10月20日,于北美PyTorch大会上将该开源。此次发布还伴随Neuron SDK 2.33.0版本一同推出,这是AWS用于编程其自研芯片的软件工具包。
其核心亮点是简洁性。开发者只需一行代码即可将模型迁移至Trainium,例如.to('neuron'),就像将代码指向任何其他设备一样。这行代码保持了对标准PyTorch API的兼容性,并支持即时执行模式,即操作在代码调用时立即运行,而无需先编译成更大的计算图。
此前的方案依赖于PyTorch/XLA集成。开发者需要在训练循环中散布诸如xm.step()之类的调用,并适应一种定制的、以编译器为中心的编程风格。这种方式虽然可行,但需要额外的代码,且API覆盖范围较窄。
它所消除的摩擦是全行业的普遍问题。PyTorch是AI研究中的主导框架,而NVIDIA的CUDA软件栈仍是运行它的默认环境;其他加速器通常要求开发者在训练开始前先调整代码。将该后端开源也使其开发过程公开化,研究人员可以审查代码、报告问题并贡献修改,而不必只依赖AWS。
Neuron 2.33还带来了多项配套升级:
- Neuron Observer,用于检测硬件健康问题的工具
- 改进的NKI功能,支持更底层的定制
- 更新的vLLM插件,经调优以适配TorchNeuron
不过,对于较旧硬件的用户来说有一个限制:TorchNeuron仅支持Trainium2和Trainium3。Trainium1和Inferentia芯片系列则被排除在外。
为什么推理优先
Trainium已在推理领域站稳脚跟,推理是训练好的模型响应实际请求的阶段。Anthropic在Trainium2和Trainium3上运行生产推理和强化学习部署。
研究训练则是另一回事。科学家们不断更换层、测试变体并进行消融实验——即移除或更改某个组件以观察哪些环节会失效的实验。每一次调整都可能调用PyTorch API的不同部分,在早期基于XLA的工作流下,这种模式需要额外的精力。
这种差异正是此次发布具有重要意义的原因:训练是新模型诞生的环节,原生支持让Trainium从模型的第一次实验开始就参与其中,而不是只在模型完成后才派上用场。
谁已在使用它
AWS表示已有超过30家早期客户在使用TorchNeuron,涵盖机器人、语音AI和大语言模型开发领域。Cartesia和Deepgram这两家从事语音AI的公司报告称,与之前的配置相比性能有显著提升。Hugging Face这一广受欢迎的开源AI模型平台也已采用Trainium支持。
此次更新还将Trainium开放给更广泛的研究用例,药物发现等领域被认为可以从更便捷地获取该硬件中受益。作为公开测试版,TorchNeuron的功能集和硬件支持在项目稳定之前仍可能发生变化;与此同时,使用Trainium1或Inferentia的团队将继续使用早期基于XLA的工具。