PyTorch 终于能原生跑 TPU 了：Google 补了六年的课，一次补完

如果你在过去几年尝试过把 PyTorch 模型搬到 TPU 上，大概率经历过这样的流程：先改代码，把模型拆成可编译的区块，再忍受漫长的首次编译，最后发现某个分布式原语不支持，回滚重来。这套体验让不少团队在选型时直接把 TPU 划掉。现在，Google 想把这页翻过去。

一分钟速览

• Google 与 Meta 合作推出开源项目 TorchTPU，为 TPU 提供原生 PyTorch 后端。
• 迁移只需把 cuda 换成 tpu、nccl 换成 tpu_dist，改动量大幅下降。
• 支持即时执行、编译执行、分布式计算和硬件专属优化。
• 在约 5000 个 Hugging Face 模型、500 多种架构上做了兼容性验证。
• 300 亿参数模型在单枚 TPUv7 上跑到每秒 25700 个 Token。

六年弯路：PyTorch/XLA 到底卡在哪

Google 从 2020 年就开始通过 PyTorch/XLA 支持 PyTorch 跑在 TPU 上。这条路能走通，但代价很高。官方自己总结了几个关键问题：需要对源代码做侵入式修改，显式定义哪些区域可以被编译；编译延迟严重；不支持 PyTorch 原生的命令式执行模式；对 PyTorch 生态里新兴的分布式原语支持滞后。

翻译成工程语言就是：你不是在“用 PyTorch 跑 TPU”，而是在“用一套为 TPU 定制的方言重写 PyTorch 代码”。模型越大、依赖越多，改造和维护成本越高。很多团队算完这笔账，宁可继续买 GPU。

TorchTPU 换了一套思路，把即时执行放到了最高优先级。迁移动作被压缩到几步：把设备从 cuda 替换为 tpu，把分布式后端从 nccl 替换为 tpu_dist，或者直接用 torch.accelerator.current_accelerator 让框架自己识别可用加速器。对绝大多数训练脚本来说，这已经接近“改两行就能跑”。

三种即时执行模式，和那个快 50% 的中间态

TorchTPU 提供 Strict、Debug、DeferAndFuse 三种即时执行模式。Strict 最贴近开发者熟悉的 PyTorch 语义，适合调试；Debug 提供更多可观测信息；DeferAndFuse 则是混合模式，把一部分计算推迟并融合，Google 声称它比 Strict Eager 快 50% 以上，同时降低 HBM 峰值内存占用。

这个设计其实承认了一个现实：纯粹的即时执行好用但不够快，纯粹的编译执行快但不好调试。真实的生产环境需要在两者之间找到一个可切换的中间态，而不是逼用户二选一。

数据支撑：兼容性和性能都拿了实样本

为了验证“不改代码也能跑”，Google 建了一条集成流水线，从 v5e 开始的多代 TPU 上，用即时和编译两种模式评估了约 5000 个 Hugging Face 模型，覆盖 500 多种架构。测试方法采用语料库分层和架构均衡选择，避免反复测试细微变体，同时确保真实生产权重被覆盖。

性能方面，Google 用标准 PyTorch 实现的 300 亿参数 Muse Glimmer 模型做了实测：单枚 TPUv7、序列长度 5120 时推理吞吐为每秒 25700 个 Token，序列长度 10240 时为每秒 22400 个；训练侧用四枚 TPUv7，两个序列长度分别为每秒 5600 和 4600 个 Token。

更大规模的验证用了 Qwen3 MoE 235B，训练从 1024 枚 TPUv7 扩展到 8192 枚。推理侧，Google 表示用 vLLM 跑大规模 MoE 负载，单位成本性能具备竞争力。这句话的潜台词很清楚：TPU 不只是“能跑”，还要在性价比上跟 GPU 掰手腕。

为什么是现在，为什么拉上 Meta

PyTorch 是 Meta 主导的框架，Google 想让 TPU 被 PyTorch 生态接纳，绕不开 Meta。两家在 TorchTPU 上的合作，本质是一次生态互换：Google 得到框架层的原生支持，Meta 得到除 NVIDIA 之外的第二个大规模算力选项。

这个动机在当下格外强烈。过去一年，GPU 供应、交付周期和价格波动一直是训练集群扩张的最大变量。对 Meta 这种体量的公司来说，多一个可用的加速器后端，等于多一层议价能力。对 Google 来说，TPU 长期只在自家云上出租，客户迁移成本高；TorchTPU 把迁移成本打下来，等于给 TPU 开了一扇更大的门。

Hacker News 上的评论很能说明开发者的情绪。用户 tingkuanpei 的留言是：“非常有意思的工作。Google 早就应该这样做了。PyTorch 的生态系统实在太丰富了。”也有人指出，Google 除了通过 GCP 提供 TPU，还会把 TPU 卖给用于本地部署的企业客户，其中包括 Anthropic，Meta 未来也可能跟进。

还差什么

TorchTPU 仍处于早期阶段。官方路线图列了动态形状、编译性能、算子覆盖、强化学习和量化几个方向。这几项恰好也是生产环境最容易踩坑的地方：动态形状决定它能不能跑变长序列和 Agent 类负载，算子覆盖决定冷门模型能不能一次性迁移，量化决定推理成本能不能压到可商用的水平。

换句话说，TorchTPU 解决的是“能不能用”的问题，“好不好用”还要再看一两年。但这一步的意义不小：它把 TPU 从“需要专门适配的特殊硬件”，往“PyTorch 的又一个后端”推了一大步。当硬件差异被框架抹平，竞争就会回到价格、供应和能效这些更实际的地方。

趋势洞察

加速器的竞争正在从芯片层上移到框架层。CUDA 的护城河从来不是某一颗 GPU，而是十几年积累的算子库、编译器和开发者习惯。要撼动它，靠单点性能不够，必须让开发者在几乎无感的情况下完成迁移。TorchTPU 是 Google 在这条路上走得最认真的一次。

对国内做国产算力的团队来说，这也是一个提醒：生态适配的投入产出比，可能比再堆一代硬件参数更高。开发者不会为了一颗芯片重写代码，但会为了“改两行就能跑”试一试。