Google 与 Meta 为 TPU 带来原生 PyTorch 支持（17 分钟阅读）

# 介绍 TorchTPU - PyTorch 来源: https://pytorch.org/blog/introducing-torchtpu/ ### 精选项目 - PyTorch logo (https://pytorch.org/projects/pytorch/) ## ## TL;DR TorchTPU 是 Google 与 Meta 的开源合作项目，为 Google TPU 引入了原生且设计合理的 PyTorch 后端。它支持开箱即用的无缝模型执行和 TPU 迁移，同时通过 eager 执行优化、torch.compile 集成和原生硬件加速，提供具有行业竞争力的推理和训练性能。 ## 介绍 TorchTPU Google TPU 在 Google 内部和外部都被非常成功地用于大规模训练和部署机器学习模型。TPU 专为 ML 设计，在性能、能效和成本方面具有显著优势。虽然 TPU 历史上一直与 TensorFlow 和 JAX 生态系统相关联，但我们很兴奋地介绍 **TorchTPU**——这是 Google 于 2025 年中期启动的一个新项目，与 Meta 的 PyTorch 团队合作，旨在让 PyTorch 原生运行在 TPU 上。 要在 Python 框架之下对 Google TPU 进行编程，开发者会使用 XLA 编译器栈。2020 年，Google 推出了第一个面向 TPU 的编译版 PyTorch——基于 Lazy Tensor 的 PyTorch/XLA。虽然具有奠基意义，但这一早期集成难以跟上 PyTorch 由社区驱动的快速演进。关键在于，PyTorch/XLA 需要侵入式的源代码修改来显式定义可编译区域，存在严重的编译时延迟，不支持 PyTorch 的原生命令式（"eager"）执行模式，并且在支持 PyTorch 生态系统中新兴的分布式计算原语方面进展缓慢。与 PyTorch/XLA 和 TorchAX 等先前尝试不同——它们将 PyTorch 硬塞进 JAX 的工作模式中——TorchTPU 被构建为 PyTorch 的一个设计合理的 TPU 后端：eager 执行是首要优先级，现有 PyTorch 工作负载只需极少的代码修改即可开箱运行（例如 `.to("cuda")` → `.to("tpu")`）。 更具体地说，在启动 TorchTPU 时，我们特意制定了一系列需要遵循的原则： - **PyTorch 公民身份**： - 运行时中不包含 JAX 依赖； - 扎实的 eager 模式支持； - 基于 Torch Dynamo（而非 Lazy Tensor）的 Torch.compile()，性能与 JAX compile 相当； - 通过在 PyTorch 栈的最低层（例如 Process Group API）提供 TPU 集合通信，广泛支持 PyTorch 分布式惯用模式（DDP、FSDP、DTensor 等）； - 与 PyTorch 生态系统的无缝集成； - 社区参与。 - **以用户为本**： - 将工作负载迁移到 TPU 应该像把 `model.to("cuda")` 改为 `model.to("tpu")` 一样简单。（仅凭这一简单修改无法达到 TPU 的峰值性能，因为模型的形状、分片和算子可能偏向 GPU 架构；但模型应能开箱运行。） - 为运行时/编译器错误提供用户友好且可操作的提示信息； - 与 PyTorch 社区已经在使用的工具和框架集成（例如 torch.profiler、DTensor、Ray）。 - **提升技术前沿**： - 构建一个在大规模工作负载上能与 GPU 上的 PyTorch 竞争的方案； - 为 eager 执行引入 JIT 编译方法，并提供运行时优化以弥合逐算子执行与编译执行之间的性能差距； - 通过引入有界动态性、避免因形状变化导致的 XLA 重新编译，以及探索可提前预编译的动态 TPU 内核，来改进 TPU 编程栈。 TorchTPU 正在积极开发中，已被 Google Cloud 一些最大的客户用于推理和训练工作负载的生产环境，并以开源项目的形式发布在 [此 GitHub 仓库](https://github.com/google-pytorch/torch_tpu)。截至本文撰写时，我们已经： - 添加了对 eager 和编译两种模式的支持，包括一种新的 DeferAndFuse eager 模式； - 添加了约 450 个 Aten 算子的支持（跳过具有结构化或组合委托的算子），包括主要深度学习架构所需的所有 Core ATen 算子的前向和反向传播； - 添加了通过 JAX/Pallas 和 Helion DSL 编写自定义内核的能力； - 在 Hugging Face 上下载量最高的 5,000 个模型上验证了开箱运行，包括： - 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](https://cloud.google.com/tpu) 是专为执行深度学习基础的大规模张量运算而设计的[专用 ASIC](https://en.wikipedia.org/wiki/Application-specific_integrated_circuit)。从 [HW 架构视角](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm)来看，[TPU 芯片](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#chips)（取决于代次）包含： - 一个或多个 **TensorCore**，每个 TensorCore 包含一个或多个矩阵乘法单元（MXU）、一个向量单元 VPU 和一个标量单元（参见[此链接](https://dl.acm.org/doi/pdf/10.1145/3360307)）。根据 TPU 代次，一个 MXU 可包含 256×256 的乘加器，排列为[脉动阵列](https://en.wikipedia.org/wiki/Systolic_array)。 - 一个或多个 **SparseCore**，它们是[专用于稀疏运算](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#sparsecore)（例如 gather、scatter、embedding）和集合通信的专用单元。它们可以与 TensorCore 并行运行，实现计算-计算和计算-通信的重叠。 - 大容量**高带宽内存**（HBM）——在 [TPU 8i](https://cloud.google.com/blog/products/compute/tpu-8t-and-tpu-8i-technical-deep-dive) 中为 288 GB，带宽 8.6 TB/s，可用于存储 ML 数据并与 CPU 共享。 - 双向**芯片间互连**（ICI）——在 [TPU 8t 和 8i](https://blog.google/innovation-and-ai/infrastructure-and-cloud/google-cloud/eighth-generation-tpu-agentic-era/) 中为 19.2 TBps，允许用户配置的 [TPU Slice](https://docs.cloud.google.com/tpu/docs/system-architecture-tpu-vm#slices) 内的 TPU 芯片之间快速通信。Slice 是 [TPU Pod](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu-pod) 的逻辑子组，由预配置数量的 TPU 芯片以 3D 环面（TPU 7、TPU 8t）或 Boardfly（TPU 8i）拓扑连接而成。 - 双向**数据中心网络**（DCN）——在 [TPU 8t](https://blog.google/innovation-and-ai/infrastructure-and-cloud/google-cloud/eighth-generation-tpu-agentic-era/) 中为 400 GBps，允许跨 Slice 的 TPU 芯片之间进行通信。 TPU 通过 [XLA 编译器](https://github.com/openxla/xla)栈编程，有两个不同的入口点：[StableHLO](https://openxla.org/stablehlo/spec)（面向 ML 算子的高层 IR）和 [Mosaic](https://docs.jax.dev/en/latest/pallas/design/design.html#lowering-pallas-to-mosaic-for-tpu)（面向自定义内核的[低层 IR](https://github.com/jax-ml/jax/blob/main/jaxlib/mosaic/dialect/tpu/tpu.td)）。实际的 TPU ISA 是专有的、可能变化的，且不公开提供。实际上，可以将 StableHLO 和（尤其是）Mosaic 类比为 PTX。下图是高层概览。 C++ 程序可以通过实例化 [PjRT Client](https://openxla.org/xla/pjrt/cpp_api_overview) 来访问本地 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 共置（co-location）无法开箱支持，我们正在研究放宽或绕过这一 TPU 限制的方法。 - StableHLO 遵循 [SSA 范式](https://en.wikipedia.org/wiki/Static_single-assignment_form)，因此不允许就地修改，即写操作总是产生一个新值（至少在概念上如此）。这对 TorchTPU 尤其重要，因为 TorchTPU 通过视图分解和反转提供了熟悉的 PyTorch "view" 语义，包括修改操作，后文将详细说明。 - XLA 编译器和 PjRT 运行时假定张量形状在编译时已知。认识到这对 PyTorch 工作负载是一个限制，我们一直在添加对有界动态性的支持，即形状维度的最大边界在编译时已知，但实际维度仅在运行时已知。 - 集合操作被降级为 XLA 的 StableHLO 算子，而不是使用 NCCL 等独立库。 ### 易用性与自动加载机制 我们在设计 TorchTPU 时充分考虑了易用性。当用户在安装了 TorchTPU wheel 包的 TPU 机器上导入 PyTorch 时，由于 PyTorch 的自动加载机制，在 TPU 上运行 PyTorch 代码只需极少的修改，几行代码即可完成： 参见下表进行并排对比。 **修改前（GPU）** **修改后（TPU）** **替代方案（TPU 或 GPU）** ```python import torch import torch.distributed as dist def main(): device = torch.device("cuda") dist.init_process_group(backend="nccl") rank = dist.get_rank() word_size = dist.get_world_size() ... model = model.to(device) ... ``` ```python import torch import torch.distributed as dist def main(): device = torch.device("tpu") dist.init_process_group(backend="tpu_dist") rank = dist.get_rank() word_size = dist.get_world_size() ... model = model.to(device) ... ``` ```python import torch import torch.distributed as dist def main(): device = torch.accelerator.current_accelerator() dist.init_process_group() rank = dist.get_rank() word_size = dist.get_world_size() ... model = model.to(device) ... ``` ### 数值精度 TorchTPU 针对 TPU 硬件适配了标准的 PyTorch 精度 API。与所有高性能机器学习架构一样，TPU 针对其特定代次偏好的原生数据类型进行了优化，以实现峰值性能和能效。对 TPU 而言，原生矩阵乘法格式是 bfloat16（BF16），它在高动态范围和高效吞吐之间取得了平衡。由于 TPU 的矩阵乘法单元（MXU）被设计为使用 BF16 输入和 32 位浮点中间累加执行算术运算，因此这一配置是 TPU 高性能计算的基石。对于非收缩运算（如逐元素加法、激活函数、layernorm 和 softmax），TorchTPU 使用 TPU 向量执行单元（VPU）执行 FP32 张量运算，该单元原生执行完整的 IEEE-754 FP32 算术。 TorchTPU 将 PyTorch 精度模式映射到底层 XLA/StableHLO 配置，以暴露这一硬件能力： - *medium*：乘法使用较低精度（将输入转换为 MXU 原生的 BF16），而累加在 FP32 中进行以保持稳定性。TorchTPU 默认使用此模式以充分利用 TPU 硬件加速，因为带 FP32 累加的单次 BF16 乘法为几乎所有深度学习工作负载提供了足够的精度。 - *high*：使用三次算法，将输入拆分为高位和低位分量以计算三个交叉项，在大约 3 倍的计算开销下恢复大部分 FP32 精度。 - *highest*：使用六次算法，计算所有交叉项以近似达到完整的 IEEE FP32 精度，在 MXU 上带来大约 6 倍的计算开销。 对于绝大多数深度学习工作负载，我们建议使用 *medium*（默认）精度。在 TorchTPU 中显式选择 `high` 或 `highest` 会触发性能警告（`TORCH_WARN_ONCE`）。仅在明确识别到与精度相关的收敛问题时才切换到 *high* 或 *highest*，因为这些模式会带来显著的性能开销。除了全局设置外，TorchTPU 还提供了一个线程本地上下文管理器（`torch.tpu.precision`），用于细粒度的作用域控制（例如，用于数值敏感的层）。 参见下表进行并排对比。 维度 PyTorch on CUDA PyTorch on TPU（TorchTPU） **默认精度模式** highest（纯 FP32） medium（Precision.DEFAULT） **快速数学加速原语** Tensor Cores（TF32） 脉动 MXU（bfloat16） **快速数学位宽** 1 符号位 + 10 尾数位 + 8 指数位（19 位 TF32 格式） 1 符号位 + 7 尾数位 + 8 指数位（16 位 BF16 格式） **累加精度** FP32（highest 和 high 模式），TF32（medium 模式） 所有模式均为 FP32 **High/Highest 精度机制** 专用硬件 FP32/TF32 单元 多次算法分解（3× 或 6× MXU 运行） **high/highest 的性能影响** 基准速度或轻微降低（~1–2×） high 约 3× MXU 计算开销；highest 约 6× 计算开销 **作用域 API** 全局 `torch.set_float32_matmul_precision` 全局 `torch.set_float32_matmul_precision` + 作用域 `torch.tpu.precision` 上下文管理器 ### Eager 执行、StableHLO 降级与 JIT 编译 按照[此 PyTorch 指南](https://docs.pytorch.org/tutorials/advanced/extend_dispatcher.html)，我们通过一个新的 C++ 后端扩展了 PyTorch 分发器，注册为 *PrivateUse1*。此外，我们使用标准的 `at::Tensor` 机制定义了设备特定的张量存储实现（通过在 `c10::StorageImpl` 内使用 `c10::UniqueVoidPtr`）。 TorchTPU eager 执行的一个独特特征是它可以以多种模式将 Aten 算子分发到 TPU： - *立即*执行，实际上是实现逐算子分发； - *延迟*到自动确定的**物化点**（Materialization Point，例如 `.to("cpu")` 语句），并与之前延迟的所有其他算子一起聚合/编译； - 两者之间的任意中间模式。 虽然这种灵活性无疑给设计带来了额外的复杂性，但我们在一开始就有意加入这一特性，因为我们希望尝试不同的策略以改进 eager 性能，超越基本的逐算子分发模式。（添加