谷歌开源TorchTPU：PyTorch代码一行不改跑上TPU，CUDA独占时代松动了

2026年10月20日，PyTorch Conference 2026上，谷歌联合Meta的PyTorch团队正式开源TorchTPU。这套软件栈让PyTorch模型直接跑在谷歌TPU上，TPU从此与GPU、CPU并列为PyTorch的一等设备。项目采用Apache 2.0协议，代码已在GitHub公开，PyPI上的版本为0.2.0，要求Python 3.11以上。

老方案有多痛

过去想在TPU上跑PyTorch，开发者只有PyTorch/XLA一条路。XLA编译器把PyTorch算子转成TPU可执行代码，问题一大堆：编译开销大、调试链路复杂、算子覆盖不全。

更麻烦的是，PyTorch/XLA只支持纯SPMD代码，也就是每个节点跑同一份程序。开发者还得手动标记可编译区域，调试一个算子经常要等完整编译周期。对习惯即时反馈的PyTorch用户来说，这套流程足够劝退。

结果就是：绝大多数AI团队守着PyTorch，而谷歌自己的芯片生态一直围着JAX和TensorFlow转。硬件再强，软件门槛摆在那里。

TorchTPU改了什么

TorchTPU通过PyTorch预留的PrivateUse1接口接入，这是框架专门留给外部硬件后端的通道。TPU上的张量就是普通PyTorch张量，没有包装类，没有子类。每个设备一个进程，张量视图和原地修改都保留。

官方推荐的工作流分两步：调试用eager模式逐算子执行，上生产用torch.compile。编译路径上，TorchDynamo抓取计算图，交给XLA生成TPU机器码，编译产物跨进程、跨主机、跨任务重启都能复用。

eager模式有三个档位：

Strict（默认）：逐算子异步分发，标准PyTorch语义
Debug：同步执行，带张量边界检查
DeferAndFuse：自动融合连续算子，减少HBM访存，论文数据显示比Strict快50%以上

自定义内核也顺了。Pallas内核可以直接当标准PyTorch函数调用，Tokamax和Helion的预置内核同样能用。分布式这块支持DDP、FSDP2和DTensor，配合感知拓扑的mesh处理TPU互连。和只支持SPMD的前辈比，TorchTPU还能跑分歧式MPMD执行。

性能数字要分开看

谷歌自报的成绩：用TorchTitan训练Qwen 3 8B，模型FLOPs利用率（MFU）达到52%，扩展到8192颗Ironwood TPU（TPU v7x）仍保持近线性。

推理侧同样有动作。vLLM和SGLang都拿到了原生后端，上面再叠一层llm-d 1.0做前缀缓存路由。谷歌举的例子是Zenteiq.ai：用GKE Inference Gateway加llm-d，每1000颗TPU v6e的日推理吞吐从14.7亿Token涨到59.7亿Token，翻了四倍多。

数字背后有注脚。一个微调案例里，开箱即用的eager模式每步要11.49秒，调优后压到2.78秒，差距四倍。序列长度或批次大小一变就触发重编译，动态形状支持还不完整，这些都被写进了官方路线图。性能数据全部来自谷歌自家测试，第三方还没验证。

为什么这一步重要

英伟达的护城河从来不只是芯片本身，CUDA加PyTorch的组合才是真正的壁垒。TPU在性价比上早就能打，谷歌宣称过每美元性能最高领先50%，卡住开发者的始终是软件。

TorchTPU把这堵墙拆掉了一角。PyTorch代码从GPU迁到TPU，理论上改一个device字符串就行：`torch.device(“tpu”)`。训练Qwen 3 235B、Llama 3 405B、FLUX.1 12B的TorchTitan配方已经备好。

Meta的态度很说明问题。PyTorch工程负责人Jana van Greunen在公告里直言，公司不能被单一硬件架构绑死，Meta已经在语言模型、广告和推荐系统上把PyTorch负载部署到TPU。Meta本来就是TPU大客户，供应链多元化对它是刚需。

Hacker News上的讨论集中在算力生态格局：GPU供应紧张、Blackwell租金两个月涨48%的背景下，开发者多了一个按成本和性能灵活选硬件的选项。

冷静看三点

第一，TorchTPU只在谷歌TPU上运行，TPU只能通过谷歌云访问。它带来的是GPU与TPU之间的可迁移性，跨云迁移免谈。

第二，PyTorch/XLA没有被正式废弃，谷歌只建议新工作负载直接上TorchTPU。存量项目迁移节奏自己把握。

第三，v0.2.0还很年轻。重编译开销、动态形状、预编译内核库，每一项都是硬骨头，XLA系编译器的老毛病需要时间消化。

对Python开发者来说，务实的做法是拿TorchTitan的小模型配方在v5e或v6e上试一轮，把torch.compile的收益和重编译成本都量一遍，再决定生产负载要不要跟。算力成本高企的当下，多一种硬件选项，议价能力就多一分。