Google 与 Meta 发布 TorchTPU：PyTorch 原生 TPU 后端

精选项目

TorchTPU 是 Google 与 Meta 之间的开源合作项目，为 Google TPU 引入了原生且有原则的 PyTorch 后端。它支持无缝的开箱即用模型执行和 TPU 迁移，同时通过即时执行优化、torch.compile 集成和原生硬件加速，提供具有行业竞争力的推理和训练性能。

隆重推出 TorchTPU

Google TPU 在 Google 内外都被非常成功地用于大规模训练和部署大型机器学习模型。它们专为机器学习而设计，因此在性能、能效和成本方面具有显著优势。虽然 TPU 历史上与 TensorFlow 和 JAX 生态系统相关联，但我们很高兴推出 TorchTPU——这是 Google 于 2025 年中期启动、与 Meta 的 PyTorch 团队合作开展的新工作，旨在在 TPU 上原生运行 PyTorch。

为了在 Python 框架之下对 Google TPU 进行编程，开发者使用 XLA 编译器栈。2020 年，Google 推出了首个面向 TPU 的 PyTorch 编译版本 PyTorch/XLA，它基于惰性张量。尽管具有奠基意义，这一早期集成难以跟上 PyTorch 由社区驱动的快速演进。关键在于，PyTorch/XLA 需要侵入式修改源代码以显式定义可编译区域，存在严重的编译时延迟问题，缺乏对 PyTorch 原生命令式（“即时”）执行模式的支持，并且在支持 PyTorch 生态内新兴的分布式计算原语方面滞后。与此前的 PyTorch/XLA 和 TorchAX 等工作将 PyTorch 勉强塞入 JAX 的运作方式不同，TorchTPU 是作为 PyTorch 的有原则的 TPU 后端而构建的：即时执行是首要优先级，现有 PyTorch 工作负载可通过极少的代码更改开箱即用（例如，.to("cuda") → .to("tpu")）。

更具体地说，在启动 TorchTPU 时，我们有意制定了一套遵循的原则：

PyTorch 公民身份：运行时无 JAX 依赖；

坚实的即时模式；

基于 Torch Dynamo（而非惰性张量）的 Torch.compile()，性能与 JAX compile 相近；

通过在 PyTorch 栈的最底层（例如进程组 API）提供 TPU 集合通信，广泛支持 PyTorch 分布式惯用法（DDP、FSDP、DTensor 等）；

无缝集成到 PyTorch 生态系统中；

社区参与。

以用户为中心：将工作负载迁移到 TPU 应像将 model.to("cuda") 改为 model.to("tpu") 一样简单。（仅凭这一简单更改无法达到 TPU 的峰值性能，因为模型的形状、分片和算子可能偏向 GPU 架构；但模型应能开箱即用。）

为运行时/编译器错误提供用户友好且可操作的提示信息；

与 PyTorch 社区已在使用的工具和框架集成（例如 torch.profiler、DTensor、Ray）。

提升技术水平：构建一个能够在大规模工作负载上与 GPU 上的 PyTorch 竞争的解决方案；

为即时执行引入 JIT 编译方法，通过运行时优化弥合逐算子执行与编译执行之间的性能差距；

通过引入有界动态性、避免因形状变化导致的 XLA 重新编译，以及探索可提前预编译的动态 TPU 内核，改进 TPU 编程栈。

TorchTPU 正处于积极开发中，已用于部分最大的 Google Cloud 客户的生产推理和训练工作负载，并可在本 GitHub 仓库中作为开源项目找到。在撰写本文时，我们已：

添加对即时模式和编译模式的支持，包括新的 DeferAndFuse 即时模式；

添加对约 450 个 Aten 算子的支持（跳过具有结构化或复合委托的算子），包括主要深度学习架构所需的所有 Core ATen 算子的前向和反向传播；

添加通过 JAX/Pallas 和 Helion DSL 编写自定义内核的能力；

在下载量最高的 5,000 个 Hugging Face 模型上验证了开箱即用执行，包括：LLM：Qwen、DeepSeek、Gemma、LLama、Mistral、GPT、OPT、BLOOM、Falcon、StableML

音频/语音：Whisper、Wav2Vec2、HuBERT、MMS

视觉：DinoV2、SigLIP、CLIP、Swin Transformer、ConvNeXt、ViT、ResNet、DETR

文本编码器与搜索：BERT、ALBERT 与 Electra、RoBERTa、DeBERTa、T5、MT5、Flan-T5

验证了基于 DTensor 的分片；

验证了多主机推理（vLLM 和 SGLang）以及微调/训练工作负载（TorchTitan），并配合基于 Ray 的编排；

验证了与 vLLM、SGLang、TorchTitan、Ray 和 TorchRec 的集成。

在本文档的剩余部分，我们将讨论 TorchTPU 架构的关键方面以及一些初步结果。

架构深入探讨

TPU 入门

Google TPU 是专用 ASIC，专门设计用于执行构成深度学习基础的大规模张量运算。从硬件架构角度看，TPU 芯片包含（取决于代际）：

一个或多个 TensorCore，每个包含一个或多个矩阵乘法单元（MXU）、一个向量单元（VPU）和一个标量单元（见此链接）。根据 TPU 代际，一个 MXU 可以包含 256×256 个乘法累加器，它们排列为脉动阵列。

一个或多个 SparseCore，它们是用于稀疏操作（例如 gather、scatter、embeddings）和集合通信的专用单元。它们可以与 TensorCore 并行运行，从而实现计算-计算和计算-通信重叠。

大容量高带宽内存（HBM）——在 TPU 8i 中为 288 GB、8.6 TB/s，可用于存储机器学习数据并与 CPU 共享。

双向芯片间互连（ICI）——在 TPU 8t 和 8i 中为 19.2 TBps，允许用户配置的 TPU Slice 内的 TPU 芯片之间快速通信；TPU Slice 是 TPU Pod 的逻辑子组，由预配置数量的 TPU 芯片组成，这些芯片以 3D 环面（对于 TPU 7、TPUv 8t）或 Boardfly（TPU 8i）配置连接。

双向数据中心网络（DCN）——在 TPU 8t 中为 400 GBps，允许跨 Slice 的 TPU 芯片之间通信。

TPU 通过两个不同的入口使用 XLA 编译器栈进行编程：StableHLO（用于机器学习操作的高层 IR）和 Mosaic（用于自定义内核的低层 IR）。实际的 TPU ISA 是专有的、可能变更且不公开可用。实际上，可以将 StableHLO 以及（更是如此）Mosaic 视为等同于 PTX。参见下图了解高层草图。

C++ 程序可以通过实例化一个 PjRT Client 来访问本地 TPU，其 API 提供方法将 MLIR（StableHLO 或 Mosaic）编译为 TPU 可执行文件、启动该可执行文件，以及在 TPU 与主机之间复制数据。注意，TPU 执行是异步的，CPU 可以通过等待关联的 PjRTBuffer 来等待 TPU 结果就绪。

TPU 编程模型的显著特性：

TensorCore 设备可以被视为一个单线程处理器，具有一个或多个大型矩阵乘法单元（MXU）。因此，要达到峰值性能，关键在于模型的形状要足够宽，以充分发挥 MXU 的潜力。这与 GPU 不同，在 GPU 上并行启动多个较小的内核有助于最大化硬件利用率。

XLA 编译器可以指示 TensorCore 将稀疏操作和集合通信操作卸载到 SparseCore，以实现更高效和/或并行的执行。这要么自动发生，要么通过编译后的 StableHLO 模块中用户提供的 MLIR 注解发生。

一个 TPU 设备只能由一个用户空间进程访问。遵循我们在 PyTorch 生态系统中无缝集成的目标，对于分布式 TorchTPU 工作负载，我们采用每设备一个进程的模型，这与 JAX 的每主机一个进程模型相反。然而，由于一个 TPU 只能由一个用户空间进程访问，RL 共置无法开箱即用地支持，我们正在研究放宽或绕过这一 TPU 限制的方法。

StableHLO 遵循 SSA 范式，因此不允许原地修改，即写操作总是产生一个新值（至少在概念上如此）。这对 TorchTPU 尤为重要，TorchTPU 确实通过视图分解和反转提供了熟悉的 PyTorch“视图”语义（包括修改），如本文档后面所述。

XLA 编译器和 PjRT 运行时假定张量形状在编译时已知。认识到这对 PyTorch 工作负载是一个限制，我们一直在增加对有界动态性的支持，即形状各维度的最大边界在编译时已知，但实际维度仅在运行时已知。

集合通信操作会作为 StableHLO 操作下沉到 XLA，而不是使用诸如 NCCL 之类的单独库。

可用性与自动加载机制

我们在设计 TorchTPU 时考虑了可用性。当用户在已安装 TorchTPU wheel 的 TPU 机器上导入 PyTorch 时，得益于 PyTorch 的自动加载机制，在 TPU 上运行 PyTorch 代码只需进行很小的更改，几行代码即可完成：

请参见下表进行并排比较。

数值精度

TorchTPU 为 TPU 硬件适配标准 PyTorch 精度 API。与所有高性能机器学习架构一样，TPU 经过优化，使用其特定代际首选的原生数据类型来实现峰值性能和能效。对于 TPU，原生矩阵乘法格式是 bfloat16（BF16），它在高动态范围与高效吞吐之间取得平衡。由于 TPU 矩阵乘法单元（MXU）被设计为使用 BF16 输入执行算术运算，并以 32 位浮点进行中间累加，这种配置是 TPU 高性能计算的基石。对于非收缩操作（例如逐元素加法、激活、layernorm 和 softmax），TorchTPU 使用 TPU 向量执行单元（VPU）执行 FP32 张量操作，VPU 原生执行完整的 IEEE-754 FP32 算术。

TorchTPU 将 PyTorch 精度模式映射到底层 XLA/StableHLO 配置，以暴露这一硬件能力：

medium：乘法使用较低精度（将输入转换为 MXU 的原生 BF16），而累加在 FP32 中执行以保持稳定性。TorchTPU 默认使用此模式以充分利用 TPU 硬件加速，因为单遍 BF16 乘法配合 FP32 累加可为几乎所有深度学习工作负载提供足够精度。

high：使用 3 遍算法，将输入拆分为高分量和低分量来计算三个交叉项，以大约 3 倍的计算成本恢复大部分 FP32 精度。

highest：使用 6 遍算法计算所有交叉项，以近似完整的 IEEE FP32 精度，在 MXU 上产生大约 6 倍的计算开销。

对于绝大多数深度学习工作负载，我们建议使用 medium（默认）精度。在 TorchTPU 中显式选择 high 或 highest 会触发性能警告（TORCH_WARN_ONCE）。只有在明确识别出与精度相关的收敛问题时，才切换到 high 或 highest，因为这些模式会触发显著的性能开销。除全局设置外，TorchTPU 还提供线程本地上下文管理器（torch.tpu.precision），用于细粒度作用域（例如用于数值敏感的层）。

请参阅下表进行并排比较。

Eager 执行、StableHLO 降级和 JIT 编译

遵循这一 PyTorch 指导，我们使用新的 C++ 后端扩展了 PyTorch 调度器，并将其注册为 PrivateUse1。此外，我们使用标准 at::Tensor 机制来定义张量存储的设备特定实现（在 c10::StorageImpl 内使用 c10::UniqueVoidPtr）。