一个后缀改写的东西

2026 年 11 月 3 日的 AI 洞察日报头条,属于一个看起来很小的新东西:Google 与 Meta PyTorch 团队联合开源的 TorchTPU。它做的事情用一句话就能说清——让 PyTorch 原生跑在 TPU 上,现有工作负载只需把 .to("cuda") 改成 .to("tpu"),就能开箱运行。

一次设备字符串的替换,凭什么上头条?

因为在这行代码背后,是 AI 算力市场一条延续了近十年的铁律正在松动:买 GPU,买的不只是芯片,是 CUDA 生态。过去任何一家想挑战 NVIDIA 的芯片厂商,最终都会撞上同一堵墙——开发者早已被锁死在 CUDA 和 PyTorch 的 GPU 后端里,迁移成本高到让再强的硬件也卖不动。而 TorchTPU 的意义恰恰在于:它把迁移成本压缩到了改一个后缀的程度。

这件事值得深挖。因为它不只是 Google 的一次开源动作,而是整个算力竞争的游戏规则正在发生位移的信号——护城河正在从芯片本身,迁移到软件栈的开放程度与社区运营能力上。

PyTorch/XLA 的失败遗产:为什么这次不一样

要看懂 TorchTPU,得先看它的"前任"是怎么败下阵来的。

2020 年,Google 推出了第一代 PyTorch-on-TPU 方案 PyTorch/XLA,基于 Lazy Tensor 机制。从纸面上看,它打通了 PyTorch 和 TPU,但实际用过的团队几乎都体验过它的三宗罪:

第一,侵入式改码。 开发者必须在源代码里显式标记哪些区域可以编译,这意味着一份模型代码要维护 GPU 和 TPU 两个分支,工程团队最恨的就是这种事。

第二,编译延迟严重。 图编译要等,改一个 shape 重新编译又要等。在快速迭代的研究场景里,这种延迟是致命的。

第三,跟不上 PyTorch 的演进。 PyTorch 是社区驱动的、以月为单位狂奔的项目,而 PyTorch/XLA 始终像个外挂适配层,eager 模式支持残缺,DDP/FSDP/DTensor 这些新的分布式习语迟迟跟不上。

后来 Google 又试过 TorchAX——把 PyTorch"塞进" JAX 的运行方式里。本质上,这两代方案的思路都是一样的:让用户迁就我的基础设施。

TorchTPU 的根本转向在于设计原则的倒置。Google 在项目文档里明确列出了所谓"PyTorch 公民身份"(PyTorch Citizenship)原则:运行时不依赖任何 JAX 组件;eager 执行是一等公民;torch.compile 基于 Torch Dynamo(而不是 Lazy Tensor);DDP、FSDP、DTensor 这些分布式习语通过 PyTorch 栈最底层的 Process Group API 原生支持。翻译成人话就是:这次是 TPU 来迁就 PyTorch 生态,而不是让 PyTorch 用户迁就 TPU。

数字里的工程诚意

判断一个开源项目是真投入还是发布会工程,看细节最准。TorchTPU 交出的几组数据,能看出这不是玩具:

算子覆盖:约 450 个 Aten 算子,覆盖主要深度学习架构所需的全量 Core ATen 算子的前向和反向传播。450 这个数字听起来不大,但 Aten 里剩下的多是结构性或可组合委托的算子——真正挡住模型运行的坑,基本都填了。

生态兼容性:与 Hugging Face Transformers v5、Diffusers、Kernels 库内置兼容,在 HF 下载量最高的约 5000 个模型上做了开箱验证——LLM(Qwen、DeepSeek、Gemma、Llama、Mistral)、语音(Whisper、Wav2Vec2)、视觉(DINOv2、SigLIP、ViT)全在内。最终结果:4917 个模型端到端测试,Tier 1 通过率 98%,Tier 2 架构数值一致性(对 GPU 基线)97%。

性能不再打折:这是最关键的一条。传统认知是"跨硬件后端=性能妥协",但 TorchTPU 在纯 PyTorch 代码(不写任何自定义 kernel)上,单颗 TPU v7 跑 30B 模型 prefill 达到 25.7k TPS、59% MFU;用 TorchTitan 在整整 8192 颗 TPU v7 上训 Qwen3 MoE 235B,从 1 个 slice 扩到 8 个 slice,TPS 每 TensorCore 几乎持平(2236 → 2104),MFU 稳定在 27-29%——教科书级的弱线性扩展。64 卡规模上,PyTorch/TorchTitan 跑出 1828 TPS,对比 JAX/MaxText 的 1937 TPS,差距只有 5%。也就是说,Google 自家王牌的 JAX 栈,在自家芯片上对自家新后端的优势已经聊胜于无。

第三方背书也有了:SemiAnalysis 的 InferenceX 基准里,TPU v7 + vLLM-TorchTPU 组合在 MoE 推理上拿到了每百万 token 2.01 美分的最低成本(9 月首测时还是 3.6 美分),AgentX 智能体负载基准上更是 0.83 美分、全帕累托前沿领先。

这些数字合起来说明一件事:迁移的心理成本和工程成本被同时打掉了。

更大的图景:这不是 Google 一家的动作

把镜头拉远,会发现 TorchTPU 并非孤立事件。就在同一周,AWS 发布了 TorchNeuron——Trainium 芯片的原生 PyTorch 后端;AMD 在 MI350 系列上推进 FlyDSL 集成;此前还有 ROCm 对 PyTorch 的长期投入。“为非 CUDA 硬件做原生 PyTorch 后端"正在成为一个集体动作。

这个集体动作背后是同一个判断:AI 算力市场进入了买方市场阶段。当训练一家前沿模型的算力采购动辄以十亿美元计,云厂商和模型公司第一次有了真实的动机和议价能力去认真评估第二供应商。Anthropic 同期披露的"收入增 12 倍、经营亏损 230 亿美元"侧面印证了算力开支的量级——这么大的盘子,没有买家愿意单点依赖。

但第二供应商成立的前提,是软件栈不能拖后腿。于是我们看到了逻辑的闭环:

  1. 硬件层面,TPU v7、Trainium 2、MI350 在实测中已经能与同代 GPU 打到同一量级;
  2. 软件层面,CUDA 之外最大的生态资产其实是 PyTorch 本身——而 PyTorch 的治理权在 Meta 和社区手里,不属于 NVIDIA;
  3. 于是所有挑战者的最优策略趋同:与其自建生态,不如把 PyTorch 变成中立的公共基础设施,把"生态锁死"这个 CUDA 最大的护城河变成公共品。

Google 这次的选择尤其值得玩味。JAX 是 Google 亲生的、深度绑定的、随 TPU 一起长大的框架,Google 内部几乎所有大模型都在 JAX 上训练。把 TPU 的 PyTorch 后端做成熟、开源、性能追平 JAX,等于主动稀释自家框架的独占价值。为什么?因为 Google 清楚地算过账:TPU 的外部收入天花板,卡的不是硬件产能,是开发者到达率。JAX 再优秀,社区体量与 PyTorch 差着数量级;而外部客户(尤其是 Google Cloud 最大那批客户)的存量代码 99% 是 PyTorch。与其逼客户换框架,不如让芯片去适配客户。“买芯片送生态"的时代结束了,现在得靠软件栈去赢。

对工程团队意味着什么

落到一线团队,这件事的实际影响分三层:

短期(现在就能做):训练与推理负载的可迁移性显著提高。.to("tpu") 级别的迁移成本意味着团队可以在不重写代码的前提下做真实的跨硬件基准测试——注意 TorchTPU 文档也坦承,改个后缀能跑起来,但峰值性能仍需针对 TPU 架构调整 sharding 和 kernel。即便如此,“能跑"和"要重写"之间隔着的正是采购议价权的分界线。

中期(未来 6-12 个月):TorchTPU 路线图上的几项值得跟踪——FP8/FP4 量化支持、bounded dynamism 生产化(解决动态 shape 重编译)、端到端 RL 训练闭环(TorchTitan + vLLM/SGLang + Ray + VeRL)。如果这些兑现,“在 TPU 上复现一套 GPU 上的 post-training 流水线"将从工程项目变成配置问题。

长期(结构性):当每个芯片阵营都有一个合格的 PyTorch 原生后端,硬件之间的比较将被迫回归到最诚实的维度——单位成本的实际吞吐和 TCO。芯片厂商的竞争重心从"生态锁死"转向"软件栈开放度 + 社区运营 + 性价比”,这对所有买家都是好消息。

当然,冷静的另一半也要看到:CUDA 十几年积累的 kernel 库、工具链、长尾算子覆盖,不是一两个 98% 通过率能完全对冲的;TorchTPU 今天仍处于"生产可用但快速演进"的阶段,FlexAttention 这样的新特性还未支持。真实世界里,性能帕累托前沿的领先往往依赖 Pallas 自定义 kernel 这类专家工作——“开箱即跑"和"开箱即最优"之间仍有距离。

结语:软件栈成为主战场

TorchTPU 最深层的信号,是 AI 基础设施竞争的叙事换了主角。上一个十年,这个故事讲的是谁家的芯片更快;下一个十年,决定份额分配的将是谁家的软件栈更开放、谁能让开发者以最低摩擦用上自己的硬件。硬件性能的差距在收窄,而软件栈的开放度差异正在放大。

.to("tpu") 这一行代码的真正含义是:迁移成本作为一种商业武器,正在被开源瓦解。当每一颗芯片都说着 PyTorch 的母语,算力市场才第一次真正接近"性能与价格说话"的理想市场形态。对被锁定太久的买家来说,这大概是目前最值得期待的结构性变化。


参考:PyTorch Blog: Introducing TorchTPU · Google Cloud Blog · InfoQ · SemiAnalysis InferenceX · AI洞察日报 2026-11-03