Trainium

torch-neuronx：AWS Trainium 的原生 PyTorch 后端，支持 eager 执行、torch.compile、分布式训练与 NKI 自定义内核

aws-neuron/torch-neuronx · A native PyTorch backend for AWS Trainium with support for eager execution, torch.compile, distributed training, and custom NKI kernels

torch-neuronx 是 AWS Trainium 的原生 PyTorch 后端，通过 PrivateUse1 机制在 PyTorch 中新增 neuron 设备，让用户可以用熟悉的 PyTorch API 把模型和输入迁移到 neuron 设备上训练或推理。项目默认以 eager 模式逐算子执行，也可通过 torch.compile 的 neuron 后端编译模型或模型区域以提升性能；支持 torch.distributed、DTensor 与函数式集合通信，覆盖数据并行、张量并行和完全分片并行；并支持在 eager 模式和 torch.compile 编译图内调用 NKI 自定义内核。项目当前为 Beta，需要 Trn2 或 Trn3 实例及对应 Python 版本，并列出静态形状、连续张量、float64 降级等已知限制。

Google 与 AWS 为 PyTorch 开辟 Nvidia 之外的新路径

在 2026 年北美 PyTorch 大会上，Google 开源了 TorchTPU，让 PyTorch 代码能直接跑在 Google 的 TPU 芯片上，已在 GitHub 发布；AWS 则推出 TorchNeuron 公开测试版，让 PyTorch 模型能跑在其 Trainium 芯片上。过去 PyTorch 开发者基本绑定 Nvidia 的 CUDA 生态，换芯片往往要大量改代码，这两个项目把迁移简化到改几行代码甚至只改一个配置。不过官方也说明，想达到最高性能仍需针对新芯片专门调优，例如 TorchTPU 默认分布式方案可能比优化方案慢约 30–40%。对开发者来说，这多了可选的硬件路线，但现阶段迁移容易、跑出极致性能仍不简单。