从 .to("cuda") 到 .to("tpu")：Google 与 Meta 开源 TorchTPU，意味着什么？

10 月 20 日，在 PyTorch Conference 2026 上，Google 与 Meta 的 PyTorch 团队联合开源了 TorchTPU：一个为 Google TPU 打造的原生 PyTorch 后端。过去想在 TPU 上跑 PyTorch，要么重写成 JAX，要么忍受 PyTorch/XLA 的种种限制。现在，官方给出的迁移方式只有一行：把 model.to("cuda") 换成 model.to("tpu")。

为什么不是 PyTorch/XLA 的升级版

PyTorch/XLA 诞生于 2020 年，基于 Lazy Tensor。官方博客直言它的问题：需要侵入式修改代码、编译延迟严重、不支持 eager 执行，对新的分布式原语跟进缓慢。

TorchTPU（2025 年中启动）奉行"PyTorch 公民"原则：运行时不依赖 JAX；torch.compile 基于 TorchDynamo；集合通信下沉到 Process Group 层，直接支持 DDP、FSDP、DTensor。它通过 PrivateUse1 机制注册 C++ 后端，张量就是普通 torch.Tensor。

技术亮点：DeferAndFuse 与编译缓存

TPU 的 TensorCore 本质是单线程处理器，XLA 又偏好静态形状，逐算子下发性能不佳。TorchTPU 提供三种 eager 模式：

• Strict（默认）：逐个 ATen 算子异步下发；
• Debug：同步执行，便于排错；
• DeferAndFuse：延迟记录算子，遇到 .to("cpu")、.item() 等物化点时融合成更大的 StableHLO 图再编译。官方称可比 Strict 快 50% 以上，并降低 HBM 峰值占用。

配套三级编译缓存（进程内存、同机共享内存、远端文件系统）避免重复编译。视图语义方面，PyTorch 允许共享内存和原地修改，而 StableHLO 是纯函数式的，TorchTPU 用"视图分解与逆转"桥接。

实际表现

• 约 4,917 个 Hugging Face 模型验证：Tier 1 通过率 98%（4,821 个）；Tier 2 数值对齐 GPU 的为 638/657（97%）。
• 64 颗 TPU v7 训练 Qwen3 MoE 235B：TorchTitan+TorchTPU 1828 TPS，JAX/MaxText 1937 TPS。
• 1024 至 8192 颗 TPU v7，每 TensorCore 吞吐基本持平（约 2,100–2,236 TPS），MFU 约 27%–29%。
• SemiAnalysis InferenceX：vLLM-TorchTPU 跑 Qwen3.5，每百万 token 2.01 美分（9 月为 3.6 美分）。

注意：以上均为官方公布，第三方复现尚待观察；仅改设备名也无法达到峰值性能。

为什么重要

降低迁移成本：同一份 PyTorch 代码可在 GPU 与 TPU 间切换，TPU 与 GPU 的每 token 成本成为真正的比较项。Hugging Face 同步支持（device_map="tpu"），Google 还发布了 vLLM、SGLang 的 TPU 引擎。

软件生态去 CUDA 化：据 Morning Tokens 报道，亚马逊前一天刚推出 TorchNeuron。

局限：暂不支持 FlexAttention；TPU 同一时刻只允许一个用户进程访问，RL 共置不能开箱即用；动态形状仍在完善；默认的 skinny collectives 会带来约 30%–40% 端到端性能损失。

结语

TorchTPU 的价值不在某个跑分，而在于把"硬件选择"从框架层的沉没成本变回工程与成本问题。不妨拿一个小推理任务，在 TPU 与 GPU 上各跑一遍，比较每百万 token 的成本。