谷歌开源TorchTPU，TPU成PyTorch原生设备，CUDA最深的护城河开始松动
在PyTorch Conference 2026于圣何塞举办期间，谷歌正式对外公开了TorchTPU项目，这是一套完整的软件堆栈，旨在让PyTorch能够直接在TPU上运行。用户现在只需在Hugging Face Transformers v5中设置参数device_map="tpu"，就能将模型部署到TPU上。
这一参数的变更背后，折射出谷歌在软件策略上的调整，同时也标志着英伟达CUDA生态系统在过去近十年里遭遇的最直接挑战。
AI芯片领域的竞争远不止于硬件本身。PyTorch作为当前应用最广泛的人工智能框架，自2016年诞生以来便与CUDA形成了紧密联系：英伟达的工程师长期投入资源，确保PyTorch在其GPU上表现最优，以至于有分析师将CUDA视为英伟达抵御竞争对手的最强壁垒。相比之下，谷歌的TPU在过去十年主要服务于自家的TensorFlow和JAX框架。PyTorch用户若想使用TPU，必须通过PyTorch/XLA这类“中转层”，这会带来大量额外的工程开销。
TorchTPU的推出改变了这一局面。按照谷歌官方的表述，TPU从此与GPU、CPU并列，成为PyTorch的原生选项。
该项目由谷歌与PyTorch的创始方Meta联合开发，同时vLLM、SGLang、Hugging Face、Ray、TorchTitan等主流开源项目也已完成对接。一份包含23页内容的技术报告已发布在arXiv上。
报告标题：TorchTPU – A Hybrid Eager/Compiled Architecture for Native PyTorch Execution on TPU
报告地址：https://arxiv.org/pdf/2610.24161
这一时间节点值得关注。过去一年间，Anthropic签署了最多100万颗TPU的使用协议，Meta也被曝与谷歌达成了价值数十亿美元的TPU租用合同。硬件已经具备市场，软件成为最后的短板，谷歌正着手弥补这一缺陷。
不是又一个“转接头”
PyTorch在TPU上运行并非新鲜事。技术报告指出，谷歌早在2020年就推出了基于惰性张量的PyTorch/XLA，并坦率地指出了其问题：需要侵入式修改源代码来划定可编译区域，编译延迟显著，不支持PyTorch原生的即时执行模式，对分布式新原语的跟进也较为滞后。后来的TorchAX则将PyTorch算子逐一翻译成JAX执行。报告认为，这些方案本质上都是让PyTorch去适应JAX的运行方式。
TorchTPU在2025年年中立项时便确立了若干原则：运行时不依赖JAX，即时执行必须可靠，将代码从GPU迁移到TPU应当像把.to("cuda")改成.to("tpu")一样简单。
在实现上，它基于PyTorch预留给第三方硬件的PrivateUse1接口，在PyTorch分发器中注册为一个新的C++后端。开发者获得的是普通的PyTorch张量，没有子类或包装器。
即时执行被划分为三个等级：
默认的Strict模式逐个分发算子，与CPU异步执行，对应GPU上的常规体验；
Debug模式同步执行，用于排查形状错误和NaN；
第三档DeferAndFuse是TorchTPU新引入的即时与编译混合模式，先将一串算子累积起来，待需要取值时再自动融合成大块计算交给TPU，速度比Strict模式快50%以上，还能降低峰值显存占用。
三种模式共享一套分层编译缓存，相同的算子序列只需编译一次即可重复使用。
追求极致性能时，开发者可以切换到torch.compile。谷歌在此处做出了刻意取舍：使用TorchDynamo捕获计算图后，不走PyTorch默认的Inductor编译器，而是将算子映射到StableHLO，交给XLA生成TPU程序。
报告给出了理由：一是能复用为即时模式搭建的基础设施，二是XLA经过多年检验足够稳定，三是Inductor的成功很大程度上建立在Triton之上，而Triton难以移植到TPU架构。这句话从侧面说明，PyTorch与GPU的绑定已深入至编译器层面。
在TPU上调用torch.compile只需指定backend=“tpu”
分布式方面，TorchTPU在Process Group层实现了TPU集合通信，DDP、FSDP2和DTensor开箱即用，还支持rank 0进程多做日志等额外工作的“不对称”写法，这是仅支持SPMD的PyTorch/XLA无法做到的。根据PyTorch/XLA官方文档的说法，TorchTPU公开后将取代PyTorch/XLA。
技术报告里的成绩单
覆盖面是第一份答卷。谷歌持续在TPU上测试Hugging Face下载量最高的约5000个模型，截至报告撰写时共4917个，横跨500多个架构族。其中98%能够运行不崩溃，在657个代表性架构中，97%的输出与GPU结果数值一致。
第二份答卷是不做专门优化时的性能。谷歌仅使用PyTorch标准模块编写了一个30B参数的MuseGlimmer模型，不调用任何模型专属的定制内核：单颗TPU v7上，5120序列长度的预填充达到每秒2.57万token，MFU为59%；在4颗TPU v7上训练，MFU为38%。将TorchTitan中原本为GPU编写的同款实现稍加修改搬过来，训练MFU也有36%。
第三份答卷是规模。在Qwen3 235B MoE模型上，64颗TPU v7时PyTorch方案的吞吐达到谷歌自家JAX方案MaxText的约94%。从1024颗芯片扩展到8192颗，即一整个TPU v7 Pod，每个TensorCore的吞吐仅下降约6%，MFU稳定在27%至29%之间，损失曲线在各个规模下几乎重合。官方博客另称，Qwen 3 8B的训练MFU最高可达52%。
换句话说，PyTorch在TPU上与“亲儿子”JAX之间的差距已被压缩至个位数百分比。
推理侧的数据同样出色。据报告援引的SemiAnalysis InferenceX结果，TPU v7搭配vLLM-TorchTPU在8K输入、1K输出的场景下，每百万token成本降至2.01美分，比9月首次亮相时的3.6美分低了约四成。
报告还给出了几组案例：FLUX训练从开箱即时模式到完整编译优化，单步时间从11.49秒降至2.78秒；Meta的DLRM v2推荐模型在32颗TPU v7x上与谷歌高度优化的JAX基线完全持平。按报告的说法，TorchTPU已在部分谷歌云大客户的生产环境中承担推理和训练任务。
一次“全家桶”式的开源
比技术本身更有分量的是生态的完整度。
推理侧，与Inferact合作的vLLM后端、与RadixArk合作的SGLang后端都保留了上游原有的调度器和API服务器，现有推理栈无需改动；谷歌还联合Red Hat、IBM发布了llm-d 1.0，在GKE上让GPU与TPU共用同一个推理入口。
训练侧，TorchTitan提供稠密模型、MoE和扩散模型的预训练配方，强化学习框架VeRL合入了TPU后端，Ray支持将同一工作流的不同阶段分别调度到CPU、GPU或TPU。
端侧，前身为TensorFlow Lite的LiteRT正式加入PyTorch生态。TorchTPU、vLLM-TorchTPU、SGLang-TorchTPU和TorchTitan四个核心仓库同步开源。
站台方的表态同样值得玩味。Meta PyTorch工程总监Jana van Greunen直言，Meta不能被单一硬件架构束缚，TorchTPU已支撑起其LLM、广告和推荐系统的关键生产场景。要知道，Meta可是英伟达的超级大客户。
从内部代号到GitHub
回顾TorchTPU的路径，节奏相当紧凑。2025年年中，项目在谷歌内部启动；同年12月，路透社首次披露这一代号，称谷歌正与Meta密切合作，并考虑开源部分代码。2026年4月，谷歌开发者博客公开技术架构，项目负责人在巴黎PyTorch Conference Europe上登台。9月，TorchTPU进入私有Beta，并在SemiAnalysis的推理基准中首次亮相。到今天，整套栈正式开源。
从“考虑开源部分代码”到四个核心仓库全部公开、各大项目上游合入，谷歌给出的答案比外界预期的更彻底。
CUDA护城河会被填平吗
一个更进一步的问题是：框架层的壁垒被拆除后，英伟达还剩下什么？
首先要看到，TorchTPU抹平的是框架层，而不是硬件差异本身。报告自己也承认，仅改一行.to("tpu")无法获得峰值性能，因为模型的形状、分片和算子选择可能本就偏向GPU；TPU的TensorCore本质上是单线程处理器，形状要足够宽才能喂饱矩阵单元，这与GPU靠并行跑大量小内核的思路截然不同。
一些细节差异同样需要迁移者留意：TorchTPU默认采用BF16乘法加FP32累加，而PyTorch在CUDA上默认是完整FP32；默认的集合通信编译方式会带来30%到40%的端到端损失，需要额外标注来优化；每颗TPU只能被一个进程访问，强化学习中训练与推理共置的做法暂时无法开箱支持。“能跑”和“跑得好”之间，仍然隔着调优的工作量。
其次，这并非谷歌一家的动作。PyTorch Conference 2026上，AWS也在讲Trainium走向原生PyTorch的历程，其Neuron SDK已于今年停止支持用PyTorch/XLA做训练，转向原生的TorchNeuron。
PyTorch正在变成一层事实上的硬件抽象层，PrivateUse1则成了各家芯片接入的标准插槽。当所有硬件都能以原生方式运行同一份PyTorch代码时，竞争焦点就会从“生态能不能用”转向“每美元能换来多少吞吐”，这恰恰是谷歌在发布稿里反复强调的指标。
当然，变量依然存在。报告中的大部分性能数据由谷歌自己测得，还需要社区在更多模型和场景中独立验证；TPU能否以足够的产能和价格触达普通开发者，同样决定着这次开源的最终影响。谷歌也表示，未来6到12个月的重点是生产环境加固、动态形状支持以及面向下一代TPU的FP8、FP4量化。
但方向已经清晰。过去十年，选择PyTorch几乎就等于选择了CUDA。从今天起，这道等式不再成立。
如果你想了解更多关于AI硬件生态的最新动态，不妨访问爱游戏网站获取更多资讯。
参考链接
https://cloud.google.com/blog/topics/ai-infrastructure/torchtpu-is-now-open-source-run-native-pytorch-on-tpus https://developers.googleblog.com/torchtpu-running-pytorch-natively-on-tpus-at-google-scale/ https://www.denvergazette.com/2025/12/17/google-works-to-erode-nvidias-software-advantage-with-metas-help/ https://pytorchconferenceeu2026.sched.com/event/2Ikir 本文来自微信公众号“机器之心”（ID：almosthuman2014），作者：Panda，36氪经授权发布。