据SiliconANGLE报道，亚马逊10月20日宣布推出工具包TorchNeuron，使开发者能在AWS Trainium芯片上运行PyTorch工作负载，并计划将其开源。该工具旨在减少在Trainium实例上运行AI项目所需的适配工作。

PyTorch是当前使用最广的人工智能开发框架之一。此前，在基于Trainium的AWS实例上运行由PyTorch支持的项目需要大量手动工作，这增加了开发成本，也给用户设置了进入门槛。亚马逊表示，使用TorchNeuron后，开发者无需进行大规模代码修改，只需配置一项设置，告知工作负载应使用哪种芯片。

TorchNeuron支持Trainium2和Trainium3，这两款加速器均可用于训练和推理。Trainium3拥有8个计算核心，每个核心包含4个模块，分别针对AI模型处理数据时不同类别的计算进行优化；芯片在片上内存中保留模型数据的内存池容量为154GB。AWS将Trainium3部署在定制服务器中，每台服务器最多可容纳144个加速器。亚马逊称，单套系统在以MXFP8格式处理数据时，每秒可完成最多362千万亿次计算。MXFP8格式被AI模型广泛使用。

Trainium系列还配有名为NKI的定制编程语言。该语言带有编译器，允许开发者对AI模型进行底层优化，以降低模型对硬件的占用并提升速度。TorchNeuron包含一组预先打包的NKI算子，即供客户开发AI模型使用的软件构建模块。部分算子可简化混合专家架构的实现，另一些则用于降低神经网络的内存占用，并加速量化等相关任务。

TorchNeuron除支持PyTorch外，还支持软件团队常与该框架搭配使用的开源工具，其中包括AI训练引擎TorchTitan。TorchTitan通过完全分片数据并行方法将工作负载拆分到不同芯片上，不仅拆分正在训练的模型，也拆分优化器状态和梯度。优化器是提升模型能力的程序，梯度则在这一过程中发挥重要作用。亚马逊称，TorchNeuron包含的功能可让开发者利用Trainium芯片进一步优化训练流程。

目前已有AWS客户采用该工具包。获风险投资支持的语音智能体初创公司Cartesia AI正使用TorchNeuron在Trainium芯片上运行推理工作负载；Lieber Institute for Brain Development则在探索如何优化其由AI支持的医学研究工作流程。