开源项目

AWS开放TorchNeuron公测,让PyTorch原生接入Trainium

TorchNeuron: native PyTorch backend public beta

AI 基础设施模型部署模型服务框架分布式训练基础设施

概述

TorchNeuron是AWS为Trainium2与Trainium3提供的原生PyTorch后端。本次按AWS宣布正式开放源码及公测的阶段收录,项目容器和内部代码提交更早,不据提交日期认定当时已公开。它通过PrivateUse1扩展把Trainium注册为neuron设备,兼容的模型和输入可移到该设备,使用PyTorch张量、模块、自动混合精度与分布式接口。 后端提供逐算子执行的eager路径,以及torch.compile的neuron编译后端;训练可继续使用标准优化器和分布式抽象。NKI自定义算子可注册为PyTorch操作,纳入执行或编译图。仓库公开Python/C++扩展与构建入口,采用Apache2.0;需要另外配置Trainium实例、Neuron驱动、运行时、编译器与匹配的PyTorch/torch-neuronx软件版本,不能直接套用旧XLA版同名包。 公测存在明确边界:编译图只支持静态形状,新的输入形状可能重新编译;非连续张量会产生连续化复制,部分API依赖CPU回退或尚不支持。float64按float32执行可能改变精度结果,TorchAO集成尚未在该公开版本提供。首次编译开销也不能当成稳定运行时延。使用者应按官方兼容性表验证目标工作负载;本次仅静态核对源码、构建和运行入口,未在Trainium上进行迁移或性能实测。