改一行 .to(tpu) 就能跑：谷歌和 Meta 开源 TorchTPU，TPU 成了 PyTorch 一等公民

让一个 PyTorch 模型跑在 TPU 上，过去是件需要”动手术”的事：装 XLA、改图、处理一堆惰性张量的边界情况，调通之后还要担心下一次框架升级把它弄坏。在圣何塞举行的 2026 年 PyTorch 北美大会上，谷歌和 Meta 一起把这件事压缩成了一行代码——把 model.to(“cuda”) 改成 model.to(“tpu”)。

一分钟速览

• TorchTPU 是谷歌 TPU 的原生 PyTorch 后端，由谷歌与 Meta 联合开源
• 取代 2020 年推出的 PyTorch/XLA，后者需要对源码做侵入式修改
• Hugging Face 模型开箱即用，vLLM、SGLang 可原生运行
• 8192 块 TPU v7x 跑 Qwen3 235B 专家混合模型，模型 FLOPs 利用率 28%
• 同期 AWS 也放出 Trainium 原生后端 TorchNeuron，同样走 PrivateUse1 扩展点

XLA 的问题不是不能用，是太贵

PyTorch/XLA 是一套惰性张量桥接库：它把 PyTorch 的计算图翻译给 XLA 编译器，再由 XLA 生成 TPU 指令。这套设计能跑，但代价是开发者必须理解两套执行模型的差异，代码里到处是针对 XLA 的特殊分支。谷歌云在公告里承认，旧方案”需要对源代码进行侵入性修改”。TorchTPU 换了个思路：直接做进 PyTorch 的底层调度器，让 TPU 成为与 CUDA 平级的目标设备，而不是一个需要绕着走的桥。

这个改动带来的连锁效果比一行代码本身更重要。模型层面，Hugging Face 上的模型可以开箱即用；推理层面，vLLM 和 SGLang 能原生跑在 TPU 上；训练层面，TorchTitan 覆盖稠密、专家混合和扩散三类模型；集群层面，Ray 与 llm-d 可以在 Kubernetes 上跨多机扩展 TPU 工作负载。相关代码库已经开源。

为什么是现在：推理取代了训练

PyTorch 基金会执行董事 Mark Collier 给了一个判断：推理已经取代训练，成为数据中心里最大的工作负载，由此带来”对 token 的巨大需求”。Agentic AI 基金会执行董事 Mazin Gilbert 补充了一组数字：2026 年 2 月 6 日智能体的 token 消耗量首次超过人类，此后智能体消耗增长了 14 倍，人类只增长了 2.8 倍。需要提醒的是，这组数据只来自 OpenRouter 一个平台，拿来当趋势参考可以，当成全行业统计就要打折扣。

对谷歌来说，逻辑很直接：推理是长期、稳定、可预测的负载，而 TPU 是自家最能控制成本和供给的芯片。让 PyTorch 生态”零摩擦”地跑上来，等于把全球最大的开源模型社区接进自己的算力池。Meta 的动机则写在台面上——PyTorch 工程总监 Jana van Greunen 说，推动前沿 AI 需要巨大的计算规模和硬件多样性，”我们不能被单一硬件架构所束缚”。

一行代码只是起点，不是终点

谷歌自己在公告里写得很坦白：仅凭这一行改动还达不到 TPU 的峰值性能，因为模型的结构、分片方式和算子可能更适合 GPU；但模型应该能开箱即用。真正要发挥 TPU 的算力，还得重做分片策略、写定制 kernel、调通信重叠。

规模数据给了一个参照：8192 块 TPU v7x 上跑 Qwen3 235B 专家混合模型，模型 FLOPs 利用率 28%，芯片数量增长时单芯片吞吐基本保持恒定，接近线性扩展。28% 这个数字在 MoE 推理里不算低，但也说明还有大量优化空间。为此谷歌同步放出了 TPU 代理套件，已在 GitHub 开放预览，能接进编码框架，自动完成环境接入、故障诊断和性能调优——本质上是用 AI 来补”一行代码到峰值性能”之间那段人力密集的活。

可移植性是真的，供应商锁定也没消失

值得注意的是，这不是谷歌一家的动作。亚马逊云科技在 10 月 20 日宣布了 Trainium 的原生 PyTorch 后端 TorchNeuron，代码改成 .to(“neuron”) 即可，同样通过 PyTorch 的 PrivateUse1 扩展点注册。两件事放在一起看，PyTorch 正在变成硬件厂商的”通用插头”。

但通用插头不等于没有锁定。有分析指出，PyTorch 的扩展点给了开发者一套通用语言，可优化层和供给层仍然牢牢握在厂商手里：同一份模型代码在 TPU 和 Trainium 上都能跑，跑多快、要改多少东西，取决于你站在谁的云上。TorchNeuron 首个版本目前还不支持动态输入形状，也跳过了约 22000 个 PyTorch 测试用例，这就是”能跑”和”好用”之间的距离。

一句话总结

TorchTPU 真正改变的不是某一行代码，而是议价结构：模型团队第一次可以认真地把”换一家芯片”当成一个季度内可执行的技术选项，而不是一场伤筋动骨的重构。而对正在被 token 需求推着扩容的推理服务商来说，多一个能原生跑的选择，比多一个跑分更高的芯片更实用。