Google 与 AWS 为 PyTorch 开辟 Nvidia 之外的新路径
这件事讲了什么
在 2026 年北美 PyTorch 大会上，Google 开源了 TorchTPU，让 PyTorch 代码能直接跑在 Google 的 TPU 芯片上，已在 GitHub 发布；AWS 则推出 TorchNeuron 公开测试版，让 PyTorch 模型能跑在其 Trainium 芯片上。过去 PyTorch 开发者基本绑定 Nvidia 的 CUDA 生态，换芯片往往要大量改代码，这两个项目把迁移简化到改几行代码甚至只改一个配置。不过官方也说明，想达到最高性能仍需针对新芯片专门调优，例如 TorchTPU 默认分布式方案可能比优化方案慢约 30–40%。对开发者来说，这多了可选的硬件路线，但现阶段迁移容易、跑出极致性能仍不简单。
可以使用的工具
文中可用工具
- TorchTPUGoogle 开源的 PyTorch 后端，让现有模型不经 JAX 直接跑在 TPU 上
- TorchNeuronAWS 开源的 PyTorch 后端，让模型跑在 Trainium 芯片上，目前为公开测试版
- vLLM原文提及的推理框架，与 TorchTPU 有集成
- SGLang原文提及的推理框架，与 TorchTPU 有集成
- llm-d原文提及的集成项目，具体用途原文未说明
- TorchTitan原文提及的 PyTorch 训练项目，可用于 TorchTPU 的生产级方案
- Ray原文提及的分布式计算框架，与 TorchTPU 有集成
- TorchRec原文提及的推荐系统库，与 TorchTPU 有集成
- LiteRTGoogle 贡献给 PyTorch 生态的运行时，具体用途原文未说明
- Hugging Face原文提及与 Google 有集成，具体形式原文未说明
- NKIAWS Trainium 和 Inferentia 的底层编程语言与编译器，用于写自定义算子
- Pallas原文建议用于 TorchTPU 高性能场景的 kernel 方案
讨论 0
登录后参与讨论，阅读资讯无需注册。
还没有评论，来分享第一个看法。