谷歌联手Meta开源TorchTPU：PyTorch代码一行迁TPU，98%热门模型开箱即跑

10 月 20 日，在圣何塞举办的 PyTorch Conference North America 上，谷歌宣布 TorchTPU 正式开源：一个由谷歌与 Meta PyTorch 团队联合开发的原生 PyTorch TPU 后端。对数百万 PyTorch 开发者来说，这意味着把训练和推理负载从英伟达 GPU 搬到谷歌 TPU 的代码成本，理论上已经降到了改一个字符串——把 model.to("cuda") 换成 model.to("tpu")。项目以 Apache 2.0 协议发布在 GitHub，同时开源的还有 vLLM-TorchTPU、SGLang-TorchTPU 和 TorchTitan 三个配套仓库。

为什么要重做一遍

TPU 并不缺 PyTorch 支持——谷歌 2020 年就推出过 PyTorch/XLA。但官方博客毫不客气地列出了老方案的短板：基于 Lazy Tensors 的架构需要侵入式修改源码来显式定义可编译区域，编译等待时间长，不支持 PyTorch 原生的即时执行（eager）模式，对分布式计算新原语的支持也总是慢半拍。另一条路线 TorchAX 则是把 PyTorch 硬塞进 JAX 的运行方式。

TorchTPU 于 2025 年年中在谷歌启动，定位是给 PyTorch 做一个「有原则的」TPU 后端，立了几条规矩：运行时不依赖 JAX（PyTorch Citizenship）；扎实的 eager 模式；torch.compile 基于 Torch Dynamo 而非 Lazy Tensor，性能对标 JAX 编译；在 PyTorch 栈的最底层（Process Group API）提供 TPU 集合通信，让 DDP、FSDP、DTensor 等分布式习语直接可用。

设计哲学浓缩成一句话：把工作负载迁到 TPU 应该像改一个设备名一样简单。

实力如何：数字说话

谷歌给出的验证范围相当大：

  • 支持约 450 个 Aten 算子，覆盖主流深度学习架构所需全部 Core ATen 算子的前向与反向传播；
  • 在下载量最高的 5000 个 Hugging Face 模型上做过开箱验证：其中 4917 个受测模型里 4821 个（98%）可以无崩溃运行；657 个代表性架构里 638 个（97%）与 GPU 执行结果数值一致；
  • 覆盖 Qwen、DeepSeek、Gemma、Llama、Mistral、Whisper、CLIP、ViT、BERT 等语言、语音、视觉全系热门模型；
  • 提供 Strict、Debug、DeferAndFuse 三种 eager 模式，其中 DeferAndFuse 相比 Strict Eager 性能提升 50% 以上；
  • 编译栈基于 XLA 和 StableHLO，还支持通过 JAX/Pallas 和 Helion DSL 编写自定义内核。

成本端引用了 SemiAnalysis InferenceMax 基准的数据：TPU v7 配 vLLM-TorchTPU 栈，在 8K/1K 测试中推理成本做到每百万 token 2.01 美分，比 9 月首秀时的 3.6 美分又降了四成多。训练侧，厂商案例显示 Llama 3 405B 在 128 到 1024 片 Ironwood 芯片上达到 57% 的模型算力利用率（MFU）。这些数字均为谷歌一方口径，尚无独立第三方复核。

生态一口气全接上

比后端本身更有分量的是生态接入清单。推理侧：谷歌与 Inferact 合作为 vLLM 打造了基于 TorchTPU 的原生后端，保留连续批处理、量化和多芯片并行能力；与 RadixArk 合作为 SGLang 开发后端，优化多步骤生成和 Agent 任务的延迟与吞吐；llm-d 1.0 则联合红帽、IBM 和超过 500 名贡献者、25 个组织，在 Google Kubernetes Engine 上编排 TPU 与 GPU 混合机队的推理路由。训练侧：TorchTitan 集成、强化学习框架 VeRL 的 TPU 后端（预览版）、Ray 原生支持。部署侧：Hugging Face Transformers v5 里只需设 device_map="tpu" 即可上 TPU；端侧运行时 LiteRT（前身 TensorFlow Lite）也借 LiteRT-Torch 加入了 PyTorch 生态，把兼容执行延伸到 Android、iOS、网页和 IoT。另有一个早期预览的 TPU Agent Suite，负责迁移上云的自动化引导、诊断与推理调优。

Meta PyTorch 工程总监贾娜·范·格鲁嫩表示，Meta 需要在多种硬件架构上运行大模型、广告和推荐系统，TorchTPU 让 Meta 能把 PyTorch 工作负载部署到谷歌 TPU 上。谷歌则透露，TorchTPU 已被部分最大的 Google Cloud 客户用于生产环境的推理和训练。

时机的深意：前有亚马逊，后有 CUDA 护城河

这次开源的时机耐人寻味——就在前一天，亚马逊发布了面向 Trainium 芯片的 TorchNeuron 公测版，同样提供 eager 执行、torch.compile 和标准分布式 API，同样采用 Apache 2.0 协议。两大云厂商背靠背出手，共同瞄准的都是英伟达 CUDA 护城河：超过 400 万开发者、85% 的 AI 应用建立在 CUDA 生态上，每一年的文档和工具积累都在抬高迁移成本。亚马逊甚至更进一步，把代码直接提交到 PyTorch 上游，试图从框架核心剔除英伟达专属假设。而截至本周一，英伟达市值已站上 5.1 万亿美元，对这两次发布尚未公开回应。

冷静看待：开源的是软件，租不到的是芯片

需要指出几个现实限制。简单换一行设备名能让模型跑起来，但拿不到峰值性能——模型的形状、分片和算子往往是按 GPU 架构调优的，榨干 TPU 仍需硬件感知的调优手册。TPU 指令集架构依然私有。TPU 设备目前只支持单个用户态进程，对强化学习训练的算子共置不友好；分布式里的 Skinny Collectives 因 XLA 无法重叠计算与通信，相比融合模块会有 30% 到 40% 的性能损失。

更本质的一点是：Apache 2.0 开放的是软件，芯片本身仍只能以谷歌云租用容量的形式获得。开发者摆脱了英伟达的硬件锁定，却也接上了对 Google Cloud 的新依赖。

但无论如何，PyTorch 与 TPU 之间那道代码迁移的高墙已经被拆掉大半。对关心成本的团队来说，多一个可跑的硬件选项总归不是坏事；对行业来说，模型代码与芯片解耦的这一刀，才刚刚切下去。