Google 将 TorchTPU 开源，让其 TPU 原生支持 PyTorch

Google 于 2026 年 10 月 20 日在 PyTorch Conference 2026 上宣布，其面向张量处理单元（TPU）的原生 PyTorch 技术栈 TorchTPU 现已在 GitHub 上开源，该公司在 Google Cloud 博客上表示。这款软件由 Google 与 Meta 的 PyTorch 团队共同开发，旨在让现有 PyTorch 工作负载只需进行少量代码修改，就能在 TPU 上运行。

Google 表示，此次发布使 TPU 与 GPU、CPU 一样，成为 PyTorch 原生支持的选项，涵盖训练、推理服务和部署，并可借助 vLLM、SGLang、llm-d、TorchTitan、Ray、TorchRec、Hugging Face 和 LiteRT 等工具实现。开发者可以先在即时执行模式下调试模型，再通过 torch.compile 获得生产级性能，从而提高 TPU 利用率并降低内存开销。

Meta 的 PyTorch 工程总监 Jana van Greunen 在公告中表示，TorchTPU 让 Meta 能够与核心机器学习系统团队灵活地共同开发硬件，同时支持大型语言模型（LLM）、广告和推荐系统等重要工作负载。她说：“这使我们能够将 PyTorch 工作负载部署到 Google TPU 上，直接向工程团队提供反馈，并构建所有人都需要的生态系统灵活性。”

模型托管与推理服务集成

Google 与 Hugging Face 合作，使 Hugging Face Hub 上的开放权重模型能够原生运行于 TorchTPU。在 Transformers v5 中，只需传入 device_map="tpu"，Diffusers 库也能在 TPU 上原生运行，并支持即时执行模式、torch.compile 以及跨多个加速器的张量并行。Google 表示，公司还专门构建了 CI/CD 流水线，在 TPU 基础设施上每晚持续测试 Hugging Face Transformers 套件。Hugging Face 产品副总裁 Jeff Boudier 表示，此次合作将该平台对硬件选择的承诺付诸实践。

在推理服务方面，TorchTPU 上的 PyTorch 原生 vLLM 后端由 Google 与 Inferact 合作开发。Google 表示，该后端借助定制的 Pallas 内核，在 TPU 硬件上支持连续批处理、量化和多芯片并行，同时保留上游的工作进程模型、调度器和 API 服务器，因此现有推理技术栈无需改动。Inferact 首席执行官兼联合创始人 Simon Mo 表示，此次集成为 TPU 带来了对 vLLM 的一等支持。

Google 与 RadixArk 共同开发了 SGLang 的 PyTorch 原生后端，面向结构化工作负载和智能体工作负载。公告称，SGLang 使用 RadixAttention 在不同请求间复用 KV 缓存，使重复出现的提示词前缀无需重新计算；TorchTPU 版本则采用解耦式传输后端，该后端也用作 HiCache 的 I/O 后端。RadixArk 首席技术官兼联合创始人 Banghua Zhu 表示，在公司逐步为不同芯片推出 SGLang 支持的过程中，TorchTPU 加快了其工作进展。

Google 还与 Red Hat、IBM 以及由 500 多名贡献者和 25 个组织组成的社区共同发布了 llm-d 1.0，作为云原生计算基金会（Cloud Native Computing Foundation）的 Sandbox 项目。公告将 llm-d 描述为连接 PyTorch 推理服务生态与基于 Kubernetes 的生产部署的工具：其路由器采用感知前缀缓存的路由策略，将每个请求发送给已持有其上下文的副本，并将计算密集型的预填充分离出来，与内存密集型的解码分开处理。1.0 版本新增了多租户流量控制、路由器高可用性、稳定的 v1 API，以及面向智能体的会话感知路由；在 Google Kubernetes Engine 上，通用的 llm-d 端点既可运行于 GPU，也可运行于 TPU。

Zenteiq 创始人 Sashikumaar Ganesan 教授在公告中表示，由 llm-d 驱动的 GKE Inference Gateway，使该公司合成数据生成的总推理吞吐量提高了 4 倍以上，将每日生成的令牌容量从每 1,000 个 TPU v6e 加速器 14.7 亿个提高到 59.7 亿个。

训练与编排

TorchTitan 是 PyTorch 原生的分布式训练参考实现。它现已纳入 Google 所称的生产级 TorchTPU 集成，该集成基于 torch.compile、TorchTitan 的 SimpleFSDP、专用 Pallas 内核和图级优化构建。TorchTitan 提供了 TorchTPU 配方，可用于预训练 Llama 3 70B 和 405B、Qwen 3 8B 和 32B、MuseGlimmer 30B 等稠密模型，也支持 Qwen 3 235B 混合专家模型和 FLUX.1 12B 扩散 Transformer。Google 表示，该集成已在最多 8,000 个 Ironwood TPU（TPU v7x）上完成验证；通过 TorchTitan 使用 TorchTPU 训练时，Qwen 3 8B 的模型浮点运算利用率最高达到 52%，并且在最多 8,192 个 Ironwood TPU 上实现了近线性扩展。

Google 还向 VeRL（一种强化学习框架）贡献了 TPU 后端和端到端参考配方。该 GRPO 配方基于解耦式训练器—采样器架构，并已在数学推理数据集上演示；其目的是让开发者能够在 TPU 上运行 VeRL 工作负载，同时保留现有与硬件无关的模型代码和算法代码。

Ray 为 PyTorch 数据处理、训练、强化学习和服务提供原生 TorchTPU 支持，可在单一工作流中协调 CPU、GPU 和 TPU。Anyscale 首席架构师 Edward Oakes 表示，两家公司致力于将 TPU 打造成 Ray 中的一等加速器，并补充说：“TorchTPU 将 Google TPU 直接集成到这一行业标准技术栈中。”

此前的架构与可用情况

TorchTPU 于 2026 年 4 月 7 日首次公开亮相，当天发布的一篇 Google 开发者博客文章介绍了其“急切优先”设计。该设计基于 PyTorch 的 PrivateUse1 接口，提供 Debug、Strict 和 Fused 三种急切执行模式，并由共享编译缓存提供支持。Google 当时表示，与 Strict Eager 相比，Fused Eager 可带来 50% 至 100% 以上的性能提升；TorchTPU 上的 torch.compile 则通过 Torch Dynamo 捕获 FX 图，并借助 XLA 和 StableHLO 将其降级处理，内核由 Pallas 和 JAX 编写。文章称，该技术栈开箱即支持 Distributed Data Parallel、FSDPv2 和 DTensor，其架构也能够处理不同计算进程间的执行分歧。文章还将推出公开 GitHub 仓库列为 TorchTPU 2026 年路线图中的一项计划。

torch_tpu 仓库现已在 google-pytorch GitHub 组织下公开发布，并采用 Apache 2.0 许可证。其 README 列出了自定义 ATen 内核、XLA 计算的编译缓存、支持数据并行和张量并行的分布式训练、torch.compile 集成、调试和基准测试工具，以及 Llama、Qwen、ResNet、DLRM、minGPT 和 UNet 的示例。该项目支持 Python 3.11 及更高版本；torch-tpu 0.2.0 软件包于 2026 年 10 月 19 日在 PyPI 上发布，提供适用于 CPython 3.11 至 3.14 的 Linux x86-64 轮子。用户指南发布在 torchtpu.readthedocs.io 上。README 还指出，该项目的交互式笔记本必须在 TPU VM 上运行，其中 v6e GCE 实例是最简便的选择。

Google 在同一公告中表示，原名为 TensorFlow Lite 的 LiteRT 将加入 PyTorch 生态系统：LiteRT-Torch 库只需一次 litert-torch.convert() 调用，即可转换 PyTorch 模型；转换后的模型可在 Android、iOS、macOS、Windows、Linux、网页和物联网设备上运行。一套旨在自动完成 PyTorch 模型接入、诊断和推理性能调优的 TPU 智能体工具集，目前已在 GitHub 上开启早期预览。

公告建议开发者查阅 Google Cloud 的 TPU 文档，开始在 TPU 上运行 PyTorch。