Google 开源 TorchTPU，AWS 推出 TorchNeuron 公测
Google and AWS open new paths for PyTorch beyond Nvidia
在 PyTorch Conference North America 2026 上，Google 开源了面向 TPU 的原生 PyTorch 栈 TorchTPU，AWS 则发布面向 Trainium 的开源后端 TorchNeuron 公测版，两者均已在 GitHub 提供。TorchTPU 由 Google 与 Meta 的 PyTorch 团队合作开发，运行时不再依赖 JAX，现有负载只需少量代码改动即可迁移；TorchNeuron 支持 Trainium2 和 Trainium3 实例，迁移时改用 .to('neuron') 即可进行 eager 执行。PyTorch 方面指出，TorchTPU 默认的 Skinny Collectives 方案在分布式负载下相比优化方案约有 30–40% 的端到端性能损失，追求峰值性能仍需 Pallas kernel 及 vLLM-TorchTPU、SGLang-TorchTPU 等生产级栈。
推荐理由
AI 评分 67
两家云厂商同时为 PyTorch 提供非 Nvidia 芯片的原生后端，但原文指出默认配置下分布式负载仍有约 30–40% 性能损失。
主题
正文 · AI 翻译
在本周于北美举办的 PyTorch Conference 2026 上，Google 开源了 TorchTPU，这是一个面向 Google TPU 的原生 PyTorch 栈，现已在 GitHub 上提供。AWS 也推出了 TorchNeuron 公测版，这是其面向 AWS Trainium 的开源原生 PyTorch 后端，同样已在 GitHub 上提供。
这两个发布似乎都旨在让开发者更容易摆脱 Nvidia GPU——其 CUDA 软件生态长期主导着 AI 开发。不过，尽管 Google 和 AWS 的新产品确实缓解了切换硬件时的一些难题，但要获得顶级性能仍然并非易事。
对于渴望更多计算选择的开发者来说，这令人期待，但也有些令人失望。
Google 希望让 TPU 对 PyTorch 更友好
Google 于今年 4 月面向部分客户推出了 TorchTPU 预览版（PyTorch 表示 Google 自 2025 年年中以来一直在推进该项目），预告了一种在 Google AI 芯片上运行 PyTorch 工作负载的更简便方式。长期以来，对 PyTorch 的支持一直是一件颇为别扭的事，因此新的后端有望让现有模型更容易在 TPU 上运行。
Google 表示，它希望让 TPU 软件生态对 PyTorch 社区更加开放、对开发者更友好，同时努力让这些芯片成为已经在使用 PyTorch 的开发者的更实用选择。
TorchTPU 由 Google 与 Meta 的 PyTorch 团队合作开发，允许开发者在运行时无需依赖 JAX 即可运行现有的 PyTorch 工作负载。据 PyTorch 称，这样一来，开发者无需进行大量代码修改，只需“少量代码改动”即可迁移现有工作负载。
具体而言，开发者可以使用 eager 执行来调试模型，并使用 torch.compile 来优化 TPU 利用率并降低内存开销。
Google 与 Meta 推出的这一新 PyTorch 栈，与大会上宣布的一系列其他产品和集成一同亮相。Google 称这些产品和集成“旨在帮助你在 Google Cloud 或自有基础设施上高效运行 PyTorch”。Google 还宣布了涉及 Hugging Face、vLLM、SGLang、llm-d、TorchTitan、Ray 和 TorchRec 的集成，同时向 PyTorch 生态贡献了 LiteRT。
AWS 也想加入
AWS 的发布将同样的 PyTorch 原生方式扩展到了自家的 AI 芯片上。
与 TorchTPU 类似，TorchNeuron 为开发者提供了一种用 PyTorch 原生地实验和构建模型的方式——这次是针对 Trainium，即 AWS 专为深度学习训练和推理工作负载设计的 AI 芯片产品组合。该公测版支持 Trainium2 和 Trainium3 实例。
AWS 表示，将模型迁移到 Trainium 可以很简单，只需改用 .to(‘neuron’) 进行 eager 执行即可。对于编译后的模型，还需要选择 Neuron 后端，可以将其设为默认后端，或在 torch.compile 调用中指定。
TorchNeuron 还与 Neuron Kernel Interface（NKI）集成——这是一种用于为 AWS Trainium 和 Inferentia NeuronDevices 编程的裸机语言和编译器——因此开发者可以创建自定义算子并优化 Trainium 上的性能。AWS 和 PyTorch 已就该后端的执行模式和性能权衡发布了更详细的指南。
两者都在接近目标，但尚未完全达成
尽管 TorchTPU 和 TorchNeuron 在为 PyTorch 开发者开放更多硬件选择方面是值得注意的进步，但两者都不是万能灵药。
PyTorch 表示，简单的代码改动就能让模型在 TPU 上运行，但要达到峰值性能可能需要针对 TPU 的专门优化。
对于分布式工作负载，PyTorch 表示，TorchTPU 默认的 Skinny Collectives 方案相比优化方案可能会带来约 30–40% 的端到端性能损失，因为它会阻止 XLA 编译器重叠计算与通信。
对于追求峰值性能的工作负载，PyTorch 建议进行额外调优并使用生产级栈：
“当然，对于追求峰值性能的生产工作负载，我们绝对建议采用基于 Pallas kernel 的最优方案，以及 vLLM-TorchTPU、SGLang-TorchTPU 或 TorchTitan-TorchTPU 等生产级栈。”
TorchNeuron 目前处于公测阶段。在另一篇博客文章中，AWS 表示，团队“在大多数情况下”可以通过更改配置而非重写代码或模型来试用 Trainium，但并未说明哪些工作负载需要更多改动。
在一个案例研究中，AWS 描述了其生成式 AI 创新中心如何与 AI 研究与技术公司 i14 合作，将其 Laplace 3B 模型适配到 Trainium2。该实现需要三项改动：1）为 GLA 注意力内核使用原生 PyTorch；2）重新设计 Muon 的通信方式；3）编译顺序反向传播。
这些发布为 PyTorch 开发者提供了 Nvidia GPU 之外的更多选择。就目前而言，迁移现有代码比在陌生芯片上获得峰值性能要容易。
文章《Google 和 AWS 为 PyTorch 开辟 Nvidia 之外的新路径》首发于 The New Stack。