谷歌联手Meta奇袭英伟达，开源TorchTPU，给PyTorch代码搬家

智东西（公众号：zhidxcom）
编译 杨京丽
编辑 李水青

智东西10月21日消息，今日，谷歌宣布开源TorchTPU，这是谷歌和Meta PyTorch团队共同开发的原生TPU软件栈，开发者可以用较少代码改动，将现有PyTorch模型部署到谷歌TPU上。

据TechCrunch报道，目前多数AI应用运行在英伟达GPU上，谷歌此次开源，或许试图打破英伟达在AI产业的主导地位。

PyTorch是用于开发和训练AI模型的开源深度学习框架，由Facebook AI Research（现Meta AI）团队于2016年启动开发，并于2017年正式开源发布。如今，PyTorch已成为全球AI研究和应用开发中使用最广泛的框架之一，覆盖大语言模型、计算机视觉和生成式AI等领域。

围绕TorchTPU，谷歌还将TPU接入vLLM、SGLang、llm-d、TorchTitan、VeRL、Ray、Hugging Face和LiteRT等工具，覆盖大模型训练、推理、强化学习、集群调度以及端侧部署等场景。

谷歌称，基于TorchTPU的训练方案已在最多8000颗Ironwood TPU上完成验证，在Qwen3-8B训练测试中实现最高52%的模型浮点运算利用率；在推理侧，印度AI模型公司Zenteiq.ai使用相关方案后，整体推理吞吐量提升超过4倍。

对于开发者而言，已有的PyTorch模型可以更容易迁移到TPU，常用的训练、推理和任务调度工具也能继续使用。开发者还可以根据模型规模、运行速度和成本，在TPU、GPU和CPU之间灵活选择。

开源地址：

https://github.com/google-pytorch/torch_tpu

一、TorchTPU开源，PyTorch模型可直接迁移至TPU

TorchTPU支持开发者使用标准PyTorch API在TPU上运行模型。开发者可以先通过eager mode调试模型，再使用torch.compile进行编译优化，从而提升TPU利用率并降低内存开销。

谷歌称，TorchTPU还支持分布式计算和多种诊断工具，可以减少开发者针对不同硬件反复修改代码的工作量。

Meta PyTorch工程总监贾娜·范·格鲁嫩（Jana van Greunen）说：“Meta需要在多种硬件架构上运行大模型、广告和推荐系统，TorchTPU让Meta能够将PyTorch工作负载部署到谷歌TPU上，并直接向相关工程团队反馈使用情况。”

谷歌还与Hugging Face合作，让Hugging Face Hub上的热门开源模型原生运行在TorchTPU上。

在大模型开发库Transformers v5中，开发者只需设置device_map=”tpu”，即可将模型部署到TPU。在扩散模型工具库Diffusers中，开发者可以使用torch.compile和张量并行，在多颗TPU上运行大型生成模型。

二、多款开源工具接入TPU，降低大模型推理成本

随着推理模型和Agent模型规模扩大，推理成本与延迟逐渐成为生产环境中的主要瓶颈。谷歌已将TorchTPU接入多个开源推理框架，并通过Kubernetes集群管理系统的原生路由能力，帮助开发者在TPU和GPU之间调度工作负载。

此次谷歌重点围绕三类工具展开合作：vLLM是用于高效部署大语言模型的推理引擎，SGLang主要面向结构化生成和Agent工作负载，llm-d则负责在大规模服务器集群中编排和调度推理服务。相关更新覆盖模型运行、请求处理和集群调度等环节。

谷歌已与Inferact合作，为vLLM推出基于TorchTPU的原生后端，支持连续批处理、量化和多芯片并行，同时保留vLLM原有的模型执行、请求调度和API服务架构，开发者无需大幅调整现有系统。

谷歌还与RadixArk合作，为SGLang开发TorchTPU后端。SGLang可以复用不同请求之间的缓存，减少重复提示词的计算，并针对多步骤生成和Agent任务优化延迟与吞吐量。

此外，谷歌宣布与红帽、IBM及超过500名贡献者、25家机构共同发布llm-d 1.0。该项目可以根据缓存情况分配请求，并将推理过程中的不同阶段拆分处理。1.0版本新增多租户流量控制、高可用路由和面向Agent的会话路由。

在谷歌 Kubernetes Engine（GKE）上，开发者可以通过统一的llm-d端点，在GPU和TPU之间选择更适合具体工作负载的硬件。

据AI模型公司Zenteiq.ai创始人萨希库马尔·加内桑（Sashikumaar Ganesan）称，该公司使用基于llm-d的GKE Inference Gateway生成合成数据后，整体推理吞吐量提升超过4倍，在1000颗TPU v6e加速器上的日生成Token数从14.7亿提升至59.7亿。

三、支持8000颗TPU系统训练大模型，效率最高52%

在训练侧，谷歌将TorchTPU接入PyTorch社区的训练、强化学习和编排框架。

TorchTitan是PyTorch原生的可扩展分布式训练参考实现，目前已完成与TorchTPU的生产级集成。

TorchTitan为多类模型提供了预训练配方，包括Llama 3 70B和405B、Qwen3-8B和Qwen3-32B、MuseGlimmer 30B等稠密模型，Qwen3-235B等混合专家模型，以及FLUX.1系列12B等扩散Transformer模型。

谷歌称，TorchTitan已在8000颗Ironwood TPU（TPU v7x）上完成验证。在Qwen3-8B训练测试中，TorchTPU训练路径实现最高52%的模型浮点运算利用率，并在最多8192颗Ironwood TPU上保持接近线性的扩展能力。

针对强化学习场景，谷歌向VeRL贡献了TPU后端和端到端参考配方。该方案基于训练器与采样器分离的架构，支持在数学推理数据集上运行GRPO（群体相对策略优化），同时保留原有的硬件无关模型和算法代码。

谷歌也为Ray加入TorchTPU支持。开发者可以在同一套Ray工作流中调度CPU、GPU和TPU资源，用于数据处理、训练、强化学习和推理，并将不同阶段的任务分配到更适合的加速器上。

四、PyTorch模型可运行在手机、电脑和物联网设备上

谷歌此次还宣布，LiteRT正式加入PyTorch生态。LiteRT原名TensorFlow Lite，是谷歌面向端侧AI推理推出的高性能框架。

通过LiteRT-Torch，开发者可以使用一次调用将PyTorch模型转换为LiteRT格式。该方案基于Core ATen算子构建，模型代码仍可保持在PyTorch体系内。

对于大语言模型，LiteRT-Torch可以将Hugging Face Transformers模型导出到LiteRT-LM，支持在Android、iOS、macOS、Windows、Linux、Web和物联网设备上运行。

在硬件支持方面，LiteRT可通过XNNPack调用CPU，通过开源ML Drift调用GPU，并通过统一API适配主要芯片平台的NPU。

谷歌还在GitHub上预览发布TPU agent suite。该工具通过插件接入多种代码编排工具和模型组合，可协助开发者完成PyTorch模型在TPU上的部署、问题诊断、用户问答和推理性能调优。

结语：谷歌为TPU争取更多开发者

谷歌此次开源TorchTPU，让开发者更容易把已有PyTorch模型迁移到TPU上，并继续使用熟悉的训练、推理和部署工具，由此可以降低开发者对单一硬件平台的依赖。

目前，英伟达GPU仍拥有成熟的软件生态和广泛的开发者基础，TorchTPU很难在短期内改变这一格局。但如果谷歌能够持续完善TPU工具链，并在性能和价格上展现优势，TPU有望在部分大模型训练和推理场景中获得更多应用。

来源：谷歌云、TechCrunch