谷歌 Meta 开源 TorchTPU：PyTorch 原生跑 TPU

TorchTPU 是谷歌 2025 年中启动、与 Meta 的 PyTorch 团队合作的原生 TPU 后端，10 月 20 日随 PyTorch 官方博客开源。过去在 TPU 上跑 PyTorch 只能走 PyTorch/XLA 的惰性张量：要在源码里显式划出可编译区、编译延迟高、不支持 eager，也跟不上 PyTorch 新的分布式原语。TorchTPU 的原则是：运行时不依赖 JAX、torch.compile 走 Dynamo 而非惰性张量，并在 PyTorch 最底层（Process Group API）提供 TPU 集合通信，从而直接支持 DDP、FSDP、DTensor。

覆盖面的口径是「按下载量前 5,000 个 Hugging Face 模型做持续回归」：已测 4,917 个模型、500 多个架构族，支持约 450 个 Aten 算子，并验证了对 vLLM、SGLang、TorchTitan 的集成。性能拿 FLUX 举例，单步从开箱 eager 的 11.49 秒压到 2.78 秒（约 4.1 倍，1.08 samples/s）：DeferAndFuse eager 9.95 秒、加 SimpleFSDP 分片 6.52 秒、开 torch.compile（backend=tpu）3.02 秒。

第三方口径也在文里：SemiAnalysis 的 InferenceX 上，TPU v7 配 vLLM-TorchTPU 在 8K/1K 场景做到每百万 token 2.01 美分（9 月首秀是 3.6 美分），AgentX 场景 0.83 美分、4 颗芯片即达帕累托前沿。意义不在单点性能而在位置：加上 10 月 19 日 PyTorch 给 AWS Trainium 的原生后端 TorchNeuron，非英伟达加速器正在被 PyTorch 原生接纳，CUDA 的护城河从框架层退到生态与工具链。边界同样写清：官方说项目仍在活跃开发，.to(tpu) 只保证能跑，峰值性能仍需按 TPU 的形状与分片重新调优。

原文链接：PyTorch 官方博客