亚马逊正式将其用于在Trainium AI芯片上运行PyTorch代码的软件TorchNeuron向公众开放。该工具此前处于私有测试阶段，亚马逊定于10月20日在圣何塞PyTorch大会上将其作为开源项目发布，旨在降低开发者适配门槛，直接挑战英伟达的软件生态壁垒。
亚马逊旗下芯片设计部门Annapurna Labs技术软件副总裁Colin Brace表示，以往开发者需在代码中调用针对亚马逊芯片的特定编程接口，而使用TorchNeuron后，仅需进行基本配置更改即可直接运行原有PyTorch代码。“如果你懂PyTorch，你就懂Trainium。”Brace称。
技术原理与兼容性
TorchNeuron通过PyTorch框架用于自定义硬件的官方扩展点PrivateUse1，将Trainium注册为标准设备。开发者可通过.to('neuron')将模型移至芯片。后端支持用于研究和调试的即时执行（eager execution），包括实验性的“自适应即时”模式；同时支持用于生产负载的torch.compile以及面向Trainium2和Trainium3的自定义内核Neuron Kernel Interface (NKI)。
DDP和FSDP2等标准分布式API可原生运行，性能分析通过常规的torch.profiler进行，追踪数据可在亚马逊Neuron Explorer中查看。该代码已在GitHub以Apache 2.0许可证发布。亚马逊表示，每晚针对来自Hugging Face的超过10万个模型对TorchNeuron进行测试。Hugging Face也确认，Trainium已成为其Transformers和Diffusers库的原生后端。
意在突破英伟达软件护城河
英伟达在AI芯片领域的主导地位部分得益于围绕其CUDA平台构建的庞大软件生态，代码适配其他硬件往往耗时数月。谷歌正开发类似项目TorchTPU以在其TPU上原生运行PyTorch，微软亦计划在其自研Maia芯片上运行大部分AI服务。
据报道，亚马逊正洽谈收购初创公司Decart，后者软件可帮助开发者将基于英伟达芯片的AI模型迁移至Trainium等替代方案，这笔交易可能使Decart估值达到约70亿美元。面对与英伟达的关系，Brace回应称亚马逊“始终致力于提供选择权”。
在私有测试期间，已有超过30家公司及亚马逊内部团队使用TorchNeuron。语音AI初创公司Deepgram基准测试显示，Trainium每美元处理的输入音频秒数比其现有GPU基线高出约60%，目前已在生产环境中部署。亚马逊Prime Video团队则利用该工具构建了Prime 360技术，可将NBA比赛动作重建为3D格式，并应用于本赛季部分转播。
目前，Trainium的最大用户是与亚马逊有紧密财务联系的研究实验室。Anthropic使用这些芯片训练和运行Claude模型；OpenAI今年2月同意使用2吉瓦的Trainium容量，这是亚马逊向该公司投资500亿美元协议的一部分。亚马逊CEO Andy Jassy在4月股东信中表示，包括Trainium和Graviton在内的芯片业务年营收运行率已超过200亿美元。