从一个 .to("tpu") 说起
上周的 PyTorch Conference North America 2026 上,Google 与 Meta 联合开源了 TorchTPU——一个面向 Google TPU 的原生 PyTorch 后端。官方的口号异常克制:把 model.to("cuda") 改成 model.to("tpu"),你的模型就能在 TPU 上跑起来。
对一个不熟悉基础设施的人来说,这条新闻听起来像是无数个"XX 硬件支持 PyTorch"公告中的一个。但把它放进过去几年 AI 算力格局的脉络里看,它的分量远超表面:这是 CUDA 生态锁定被撬动的过程中,迄今最实质性的一块砝码落地的声音。
这篇文章想回答三个问题:TorchTPU 到底解决了一个什么问题?为什么出手的是 Google 和 Meta 这对奇怪的组合?以及,“CUDA 锁定松动"对普通的 AI 工程团队意味着什么?
一、被 XLA 挡在门外的 PyTorch 用户
要理解 TorchTPU 的价值,得先理解它取代了什么。
TPU 是为机器学习定制的芯片,软件层之下由 XLA 编译器栈驱动。历史上,TPU 与 TensorFlow、JAX 这两个 Google 亲儿子框架深度绑定。PyTorch 用户想用 TPU,走的是 2020 年 Google 推出的 PyTorch/XLA——一条基于 Lazy Tensor 的间接编译路径。
问题在于,PyTorch/XLA 本质上是把 PyTorch “硬塞进” JAX 的工作模式:
- 要求侵入式修改源码,显式圈出可编译区域;
- 编译延迟严重,调试体验糟糕;
- 不支持 PyTorch 原生的 eager 模式——而 eager 恰恰是 PyTorch 之所以是 PyTorch 的原因,是无数研究者选择它的第一理由;
- 对 PyTorch 生态里快速演进的分布式原语(DDP、FSDP、DTensor)跟进迟缓。
结果是:过去几年,TPU 虽然在性价比和能效上有实打实的优势,但对 PyTorch 社区——也就是今天深度学习社区的主体——始终隔着一道玻璃门。门外的人看得见省钱的可能,却不愿付出改写训练代码的代价。于是大家继续买 GPU,CUDA 的护城河越挖越深。
TorchTPU 的做法完全不同。它不再把 PyTorch 塞进 JAX 的模具,而是把 TPU 做成 PyTorch 的一等公民后端:
- 运行时零 JAX 依赖;
- eager 模式是一等优先,现有 PyTorch 工作负载开箱即跑;
torch.compile基于 Torch Dynamo(而非 Lazy Tensor),编译性能对齐 JAX;- TPU 集合通信下沉到 PyTorch 栈的最底层(Process Group API 层),DDP/FSDP/DTensor 原生可用;
- 已支持约 450 个 Aten 算子,覆盖主流深度学习架构前向与反向所需的全部核心算子;Hugging Face 下载量最高的 5000 个模型已验证开箱可跑;
- 自定义算子可通过 JAX/Pallas 和 Helion DSL 编写;
- 推理侧(vLLM-TorchTPU、SGLang-TorchTPU)与训练侧(TorchTitan)的配套仓库同批开源,多机推理和微调训练均有验证。
还有一个容易被忽略的工程细节:TorchTPU 引入了 “bounded dynamism”(有界动态性)来规避 XLA 因 shape 变化反复重编译的老毛病,并探索可预编译的动态 TPU kernel。这些不是发布会上的形容词,而是踩过 PyTorch/XLA 每一个坑之后才知道要修的地方。
二、Google 和 Meta,各自的账本
为什么是这两家?
Google 的账本很直白。 TPU 对外销售的主要载体是 Google Cloud。TPU 再好,如果世界上最大的深度学习用户群体(PyTorch 社区)迁移成本高到无人问津,TPU 就只能给自家 JAX 团队和少数大客户用,云上算力卖不出规模。TorchTPU 由 Google 2025 年年中启动、与 Meta PyTorch 团队联合开发,且"已用于部分最大 Google Cloud 客户的生产推理与训练负载”——这说明它不是实验品,是 TPU 商业化的基础设施投资。Meta 公开表示计划把内部 TPU 训练负载迁移过来,等于第一个企业级客户已经签约。
Meta 的账本更有意思。 Meta 既是 PyTorch 的亲爹,又是全球最大的 AI 算力买家之一,而它的算力盘子高度依赖 NVIDIA。过去几年 Meta 持续采购 TPU 做多元算力布局——对 Meta 来说,让 PyTorch 在非 NVIDIA 硬件上跑得尽可能顺畅,是在与 GPU 供应商议价时最硬的筹码。框架中立性,就是采购议价权。
所以 TorchTPU 不是慈善,而是两家巨头利益精确重合的产物:Google 要把 TPU 卖给 PyTorch 世界,Meta 要用 TPU 压 GPU 的价。而这两个动机叠加的客观结果,是整个行业对单一硬件生态的依赖被削弱了。
三、CUDA 锁定到底在锁什么
“打破 CUDA 锁定"喊了很多年,多数尝试失败了。复盘一下,CUDA 的护城河从来不是某一项技术,而是三层结构的叠加:
- 编译器与算子层:cuDNN/cuBLAS 十几年积累的算子库与调优,覆盖了你能想到的所有算子和硬件代际;
- 框架路径层:PyTorch 默认的执行路径以 CUDA 为第一公民,其他硬件走"适配"路线,永远慢半拍;
- 人力与惯性层:全世界会写 CUDA kernel、会调 GPU 集群参数的工程师,比所有其他加速器的工程师加起来还多一个量级。
过去 AMD 的 ROCm、各家国产算力栈的困境在于:总在第一层上追赶(“我们又支持了 XX 个算子”),却在第二层被卡死——只要 PyTorch 里 CUDA 是原生路径、你是二等路径,工程团队的重心就永远在 GPU 上,第一层的差距就永远补不上。
TorchTPU 的真正意义在第二层:它把 TPU 从"适配方"变成了"原生方”。当 torch.compile、DTensor、Process Group 这些分布式与编译基础设施在最底层就对 TPU 打通,PyTorch 社区的千行代码就能自然地流向 TPU。Meta 的背书则顺手解决了第三层的一角——Meta 内部训练负载迁移,意味着大批 PyTorch 核心工程师的日常工作就在 TPU 路径上,社区惯性开始转向。
当然,要冷静地补一句:官方自己也承认,改一行 .to("tpu") 只能"跑起来",拿不到峰值性能——模型的 shape、切分和算子写法往往带着 GPU 架构的偏置,深度调优仍需针对 TPU 重做。CUDA 的护城河没有干涸,但水位第一次出现了肉眼可见的下降。
四、同一周的另一个信号:开源正在变成进攻武器
把镜头拉远,TorchTPU 并不是孤立事件。同一周里:
- Google 开源了 ML Drift——LiteRT 背后的端侧 GPU 推理引擎,覆盖 Vulkan/Metal/WebGPU,Apache 2.0 协议;
- NVIDIA 开源了 Flora——GR00T-2 机器人模型的数据标注引擎,把机器人视频标注成本压到每小时 0.72 美元,是用 GPT-6 Astra 标注的约 1/125。
三家把过去视为压箱底的基建——训练硬件后端、端侧推理引擎、机器人数据管线——在同一周开源。这个模式上的共同点值得单独指出:当"工具链锁定"这条护城河被对手用开源进攻时,最高效的防御是把护城河本身变成所有人的公共设施,把竞争拉回自己占优的战场。 对 Google 是模型与云的规模,对 NVIDIA 是硬件与互联的代差。
工具链开源之后,竞争重心向"模型能力 + 数据规模 + 硬件能效"上移。对中小团队,这意味着自建与混合算力、端侧部署的迁移成本在系统性下降。
五、给工程团队的三条务实建议
落到执行层面,我认为 AI 工程团队现在应该做三件事:
第一,把"可迁移性"写进选型权重。 新的训练与推理基建选型时,明确要求支持多硬件后端(TorchTPU、ROCm 等),哪怕你今天 100% 用 GPU。可迁移性是未来的议价权,今天它几乎免费,等你被锁定之后它贵得惊人。
第二,跑一个 TPU 的真实 PoC。 TorchTPU 刚好提供了最低成本的路径:从 Hugging Face 下载量前 5000 的模型里挑一个你实际在用的,用 vLLM-TorchTPU 跑一次推理基准,用 TorchTitan 跑一次微调。不为马上迁移,为的是建立自己对"TPU 到底便宜多少"的第一手数据——公开 TCO 计算永远代替不了你自己的 workload 实测。
第三,训练代码守住 PyTorch 标准路径。 尽量使用 DTensor/FSDP 等标准分布式原语,避免私有的 CUDA 特化代码。这不是抽象的正确,而是 TorchTPU 这类原生后端的存在方式决定的:越是走标准路径,越能在硬件代际和供应商之间自由移动。
结语
CUDA 锁定的黄昏不会在一夜之间降临——NVIDIA 的硬件代差、互联优势和工程师存量依然是真实的。但 TorchTPU 标志着打破锁定的路径终于从"另起炉灶"(让用户改用 JAX)变成了"釜底抽薪"(让硬件适配用户)。当硬件厂商开始争相成为 PyTorch 的好公民,而不是要求用户成为自家框架的移民,权力就开始从芯片厂商向框架社区回流。
对开发者来说,这是难得的好消息:你不需要站队,你需要的是让每一家都来争取你。
参考:PyTorch 官方博客:Introducing TorchTPU、TorchTPU 论文 (arXiv:2610.24161)、google-pytorch/torch_tpu