谷歌在 PyTorch Conference 2026 上宣布,面向谷歌 TPU 的原生 PyTorch 技术栈 TorchTPU 已在 GitHub 开源。据谷歌介绍,该项目由谷歌与 Meta 的 PyTorch 团队共同开发,让 TPU 与 GPU、CPU 一样成为 PyTorch 的原生选项,开发者可沿用熟悉的工具,在 TPU 上训练、服务和部署模型。
据谷歌介绍,TorchTPU 可让现有 PyTorch 工作负载经过少量代码改动后运行于 TPU。开发者可以先用 eager 模式调试模型,再通过 torch.compile 优化 TPU 利用率、降低内存开销。项目同时支持标准 PyTorch API、诊断工具和原生分布式原语。
Meta PyTorch 工程总监 Jana van Greunen 表示,Meta 不能受限于单一硬件架构,TorchTPU 让其能够将 PyTorch 工作负载部署到谷歌 TPU 上,并向工程团队直接提供反馈。
GitHub README 列出的底层能力包括 TPU 版 PyTorch ATen 内核、用于优化 XLA 计算执行的编译缓存,以及数据并行和张量并行训练支持。仓库还提供 Llama、Qwen、ResNet、DLRM、minGPT 和 UNet 模型示例。
谷歌与 Hugging Face 合作,为核心模型库接入 TorchTPU。据谷歌介绍,在 Transformers v5 中,传入 device_map="tpu" 即可将模型部署到 TPU;Diffusers 也支持在 TPU 上原生运行,并支持 eager 模式、torch.compile 和张量并行。谷歌还为 Transformers 测试套件搭建了专门的 TPU 每夜测试流水线。
推理方面,谷歌与 Inferact 合作开发的 vLLM 原生 PyTorch 后端支持连续批处理、量化和多芯片并行,同时保留上游的工作进程模型、调度器和 API 服务器。
谷歌还宣布,与 Red Hat、IBM,以及由超过 500 名贡献者、25 家组织组成的社区共同发布 llm-d 1.0。该项目面向分布式推理,支持根据前缀缓存路由请求,并将预填充与解码阶段分离。
据谷歌发布方自测,经由 TorchTitan 的 TorchTPU 训练路径在 Qwen 3 8B 上实现最高 52% 的模型浮点运算利用率(MFU),并在最多 8,192 颗 Ironwood TPU 上展示了近线性扩展与稳定性。
安装方面,项目 README 标明目前支持 Python 3.11 及以上版本,列出的 PyPI 最新公开版本为 0.2.0。开发者也可使用 Bazel 从源码构建,项目采用 Apache 2.0 许可证。