Google 联合 Meta 开源 TorchTPU 项目，为 TPU 硬件打造原生 PyTorch 后端

Google 与 Meta 正式联合开源 TorchTPU，给 TPU 硬件搭建起原生 PyTorch 运行后端，开发者可以直接使用 PyTorch 原生语法在 TPU 上运行模型，不用再修改大量代码适配编译框架，同时兼顾即时 eager 执行与编译加速两种模式，大幅降低 PyTorch 模型迁移到 TPU 集群的开发门槛。过往 TPU 主要适配 TensorFlow 与 JAX 生态，早期 PyTorch 对接 TPU 只能依靠 PyTorch/XLA 方案，需要改动源码划定编译区域，编译等待时间长，还不支持 PyTorch 原生即时执行模式，跟不上 PyTorch 社区迭代速度。本次 TorchTPU 采用混合即时执行 + 编译架构，深度对接 PyTorch 2.0 的 torch.compile 编译体系，直接调用 TPU 底层 PJRT 运行时，不需要额外封装适配层，原生支持算子融合、分布式训练、自定义内核开发，不管是大模型预训练还是云端推理，都能跑出和 GPU 持平甚至更优的运行效率。使用者只需要简单切换设备标识，就能把原本运行在 GPU 上的 PyTorch 模型一键迁移到 TPU 硬件，整个迁移过程基本不用改动模型业务代码，还能借助 TPU 的矩阵单元获得更高算力与更低使用成本。整套项目完全开源开放，配套完整测试用例与开发文档，已经在谷歌云 TPU 集群完成线上业务落地，能够稳定支撑大语言模型、计算机视觉模型的规模化训练与部署。

一、TorchTPU 与旧方案核心能力对比

1. 新旧 TPU 适配方案对照表 表格 34
2. 核心技术亮点清单

• 亮点 1：原生 PyTorch 后端，深度对接 ATen 算子库，绝大部分 PyTorch 算子开箱即用，无额外适配成本。
• 亮点 2：混合执行架构，提供严格即时执行、延迟算子融合两种模式，算子融合最高可提升 50% 运行速度，同时减少 HBM 显存占用。
• 亮点 3：完整接入 PyTorch 分布式组件，原生支持 FSDP2、torch.distributed 多卡训练，适配 TPU Pod 多机集群。
• 亮点 4：打通 Helion 高层内核开发工具，开发者可以用 PyTorch 风格代码编写自定义算子，自动编译为 TPU 可执行代码，不用手写底层 Pallas 代码。
• 亮点 5：项目采用 Bazel 构建工具，配套完善 CI 自动化测试，持续跟进 PyTorch 新版本迭代，长期保持版本同步。

二、TorchTPU 整体技术架构分层

1. 上层对接层 直接对接 PyTorch 主框架，实现 TorchTPU 设备注册，把"cuda"设备平滑扩展为"tpu"，模型代码无需重写，自动分发算子到 TPU 硬件。
2. 中间混合执行层 分为两大运行模式：严格即时模式用于调试，单算子同步执行；DeferAndFuse 延迟融合模式自动合并连续算子，减少内存临时变量，兼顾调试灵活性与运行性能。
3. 底层运行时层 调用 Google 官方 PJRT 运行时，直接调用 libtpu 硬件驱动，跳过多余封装层，充分释放 TPU 矩阵计算单元、稀疏计算单元的硬件算力。
4. 内核扩展层 兼容 Pallas 底层语言 + Helion 高层 DSL 两条开发路径，既支持资深工程师手写高性能内核，也支持普通开发者用 Python 快速开发自定义算子。

三、TorchTPU 快速部署操作步骤

第一步：环境准备，在谷歌云 TPU 虚拟机中安装对应版本 PyTorch，拉取 TorchTPU 开源代码仓库，使用 Bazel 完成项目编译安装。

第二步：设备切换，仅修改一行代码，将device="cuda"改为device="tpu"，其余模型、数据加载代码保持原样。

第三步：选择运行模式，调试阶段开启严格即时执行；正式训练时开启延迟算子融合，同时开启 torch.compile 编译优化。

第四步：开启分布式训练，直接调用 PyTorch 原生分布式接口，自动适配 TPU 多芯片互联架构，无需额外配置集群通信规则。

第五步：自定义算子开发，使用 Helion DSL 编写算子代码，TorchTPU 后端自动将代码编译为 TPU 硬件可执行指令，完成算子加速。

四、落地应用场景

• 大模型预训练：LLM 大语言模型、多模态模型在 TPU 集群上开展大规模预训练，依靠算子融合降低显存压力，提升多卡并行效率。
• 云端模型推理：vLLM、SGLang 等推理框架对接 TorchTPU，在 TPU 硬件上实现高吞吐、低延迟的线上推理服务。
• 科研模型开发：科研人员可以直接用 PyTorch 即时调试模型，不用忍受长时间编译等待，大幅缩短模型迭代周期。
• 端到端模型迁移：原有 GPU 上的 PyTorch 项目，一键迁移至 TPU 平台，低成本完成算力扩容，降低整体算力支出。

【相关 FAQ】

问题 1：TorchTPU 作为 TPU 原生 PyTorch 后端，相比之前的 PyTorch/XLA 最大优势是什么？

回答 1：TorchTPU 是原生 PyTorch 后端，支持 PyTorch 原生 eager 即时执行模式，深度兼容 torch.compile 编译体系，不用手动修改代码划定编译区域；而 PyTorch/XLA 只支持惰性图执行，代码改动量大，还无法跟上 PyTorch 社区更新节奏，TorchTPU 大幅降低了 PyTorch 模型向 TPU 迁移的成本。

问题 2：TorchTPU 的混合执行架构里，DeferAndFuse 模式能带来哪些性能提升？

回答 2：DeferAndFuse 模式会自动把多个连续算子合并成一个计算块，减少中间张量占用 HBM 显存，同时降低算子调度开销，相比纯即时执行模式，整体运行速度最高可以提升 50%，兼顾调试便利性与线上运行性能。

问题 3：使用 TorchTPU 把 PyTorch 模型迁移到 TPU，需要大规模改写原有业务代码吗？

回答 3：不需要大规模改写业务代码，TorchTPU 完成了完整设备适配，只需要把代码里的 GPU 设备标识替换为 TPU，绝大部分模型网络、数据处理代码可以直接复用，仅在自定义算子场景下做少量适配，迁移工作量极低。

问题 4：TorchTPU 支持在 TPU 多机集群上做分布式训练吗？

回答 4：TorchTPU 原生兼容 PyTorch 全套分布式工具，支持 FSDP2、torch.distributed 等主流多卡训练组件，依托 PJRT 运行时可以轻松组建 TPU Pod 大规模集群，稳定支撑千亿参数大模型并行训练任务。

【数据来源】

[2026-10-20]（Introducing TorchTPU）PyTorch 官方博客

[2026-10-20]（TorchTPU – A Hybrid Eager/Compiled Architecture for Native PyTorch Execution on TPU）arXiv 论文

[2026-10-21]（Google 与 Meta 联合推出 TorchTPU，打通 PyTorch 与 TPU 硬件）新智元

[2026-10-21]（TorchTPU 开源，PyTorch 正式原生支持 TPU 算力平台）机器之心

免责声明：本文仅对 TorchTPU 开源项目公开技术文档、论文内容与官方发布信息进行客观整理，不构成硬件选型与工程开发决策建议；实际运行性能会受 TPU 硬件型号、模型结构、集群环境影响；项目功能持续迭代，具体能力以 GitHub 开源仓库最新版本为准。