今天，AWS 推出了 TorchNeuron 的公开测试版，这是一个面向 AWS Trainium 的开源原生 PyTorch 后端。Trainium 是一系列专为深度学习 AI 训练和推理工作负载设计的高性能且具成本效益的 AI 芯片组合。TorchNeuron 使 AI 科学家和研究人员能够直接在 Trainium 上使用原生 PyTorch 进行模型实验与构建。借助 TorchNeuron，开发者只需将 PyTorch 代码切换至 .to('neuron')，即可在 Trainium 上运行，同时保持与 eager execution、torch.compile 及标准 PyTorch API 的完全兼容性。此外，TorchNeuron 集成了 Neuron Kernel Interface（NKI），以支持自定义算子并为 Trainium 硬件实现性能优化。TorchNeuron 作为开源项目，可通过 TorchNeuron GitHub 仓库获取，适用于 Amazon EC2 Trainium2 实例和 Amazon EC2 Trainium3 实例。如需进一步了解 TorchNeuron 及 AWS Trainium 加速器，请访问 AWS Neuron 页面。