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

Introducing TorchTPU

Google 与 Meta 发布开源项目 TorchTPU，为 Google TPU 提供原生 PyTorch 后端，模型迁移只需把 .to("cuda") 改为 .to("tpu")。

Google 与 Meta 官方详解 TorchTPU 的架构设计、精度策略和分布式支持，读者可据此评估把现有 PyTorch 工作流迁到 TPU 的改动成本。

另外 2 家也报道了正文 · AI 翻译中文原文译文尚不完整，完整内容请切换到原文。

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

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

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

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

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

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

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

TPU 通过两个不同的入口点使用 XLA 编译器 栈进行编程：StableHLO（用于 ML 算子的较高级 IR）和 Mosaic（用于自定义内核的较低级 IR）。实际的 TPU ISA 是专有的、可能发生变化且不公开。实际上，可以将 StableHLO 以及（更是如此）Mosaic 视为等同于 PTX。参见下图的高层示意图。

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

TPU 编程模型的显著特点：

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

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

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) ... 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) ... 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（默认）精度。在 TorchTPU 中显式选择 high 或 highest 会触发性能警告（TORCH_WARN_ONCE）。只有在明确识别到与精度相关的收敛问题时才切换到 high 或 highest，因为这些模式会带来显著的性能开销。除全局设置外，TorchTPU 还提供线程本地上下文管理器（torch.tpu.precision）用于细粒度作用域（例如针对数值敏感的层）。

并排对比见下表。

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

TorchTPU 的 eager 执行的一个独特特性是，它可以以多种模式将 Aten 算子分发到 TPU：

虽然这种灵活性确实给设计带来了额外的复杂性，但我们从一开始就刻意加入它，因为我们希望试验不同的策略，以在基本的逐算子分发模式之上提升 eager 性能。（此外，正如本文档后面所讨论的，这种灵活性在处理视图语义和 torch.compile 时非常方便。）

总之，TorchTPU 提供了 3 种 eager 模式，可以通过环境变量或上下文管理器来设置：

现在我们已经描述了 TorchTPU eager 模式的基础知识，让我们勾勒一下当收到来自 PyTorch dispatcher 的回调调用时会发生什么。假设我们使用的是 DeferAndFuse 模式，并以 at::mul.out 为例：

func.func @main(%arg0: tensor<4x16xbfloat16>, %arg1: tensor<4x16xbfloat16>) -> tensor<4x16xbfloat16> { %0 = stablehlo.multiply %arg0, %arg1 : tensor<4x16xbfloat16> return %0 : tensor<4x16xbfloat16> } 一个可选的静态算子参数映射（例如，对于 aten::native_batch_norm 使用 training 和 momentum），用于计算缓存键（下文讨论）； 返回一个新的 at::Tensor，其 c10::StorageImpl（通过间接方式）指向该 Deferred Op 实例。 当收到来自 PyTorch dispatcher 的对应于 mytensor.to("cpu") 或 mytensor.item() 或类似内容的回调时，TorchTPU 将：

表面上看，上述方法似乎效率不高，因为 CPU 只有在遇到 .to("cpu") 语句时才会启动 TPU 执行。对于真实工作负载，我们可以预期，在少数初始迭代/步骤填充编译缓存之后，CPU 执行应仅涉及 PyTorch 追踪和算子分发，因此会始终领先于 TPU 执行，而后者将忙于冗长的数值计算。尽管如此，为了加速对多次重复的算子序列（想想一堆 transformer 块）的物化，我们添加了重复算子启发式（Repeated Ops Heuristic），它会检测此类重复，并在到达 TPU 到 CPU 传输之前就触发物化点。

TorchTPU 采用编译缓存来避免不必要的 XLA 编译。与现代计算机存储系统类似，该缓存包含多个协同工作的层级：

第 1 层缓存始终启用。默认情况下，第 2 层对分布式程序启用，对单设备程序禁用。用户可以根据需要覆盖此行为。第 3 层默认禁用，但可以通过指定共享存储路径显式启用。

缓存中的每个条目都是针对 StableHLO 计算图编译的 TPU 可执行文件。在第 1 层缓存中，每个可执行文件都是一个 C++ 对象，可以直接调度到 TPU 设备上执行。由于对象无法直接跨进程共享，第 2 层和第 3 层缓存将可执行文件保存为序列化的二进制字符串。

出于性能原因，我们不使用 StableHLO 图本身作为缓存键。相反，我们将计算、编译器选项和主要组件的版本指纹化为三个 64 位整数，并使用该元组作为缓存键：

无形状键和维度键合在一起也称为图键。虽然理论上不同图、编译器选项和组件版本之间可能发生缓存键冲突，但我们认为 192 位键空间足够大，实践中不会出现此类冲突。

PyTorch 的操作语义与 StableHLO 语义之间一个特别不同的领域是可变性和共享内存视图的处理：

在存在写入的情况下，在 TPU 上实现正确的视图语义并不简单；然而，TorchTPU 借助一些机制实现了这一点，包括视图分解、视图反转和延迟操作。这一语义桥接使 TorchTPU 能够原生支持原本不可能支持的 PyTorch 模式，例如用于分布式训练的拼接参数数组。

在优化模型性能时，PyTorch 用户可能希望同时执行 3 类操作：

在 GPU 上，这通常通过使用单独的命令队列来实现，以便设备可以将工作分派到不同的硬件单元（SM、DMA 等）。在 PyTorch 中，设备命令队列以 Streams 的形式暴露，它们通常既用于用户代码，也用于框架代码（例如在 FSDPv2 中）。

TorchTPU 也实现了 PyTorch Stream API，但并非在所有情况下都能实现 GPU 上可用的相同功能。考虑到 TPU TensorCore 本质上是一个单线程执行处理器，我们在给定 TensorCore 上一次只能分派一个计算。相反，GPU 包含许多相同的计算核心（SM），可以并行运行。所幸，一个 TPU TensorCore 关联一个或多个 DMA 引擎以及一个或多个 SparseCore 引擎。不出所料，DMA 引擎可用于在主机与设备之间进行数据传输，并与 TensorCore 操作并行执行。更有趣的是，SparseCore 引擎是一个独立单元，针对稀疏操作（例如 gather/scatter）进行了优化，并且还可以支持集合通信操作。原则上它可以支持完整的 StableHLO 操作集，但需要注意的是，它在稀疏数学运算上会优于 TensorCore，而在稠密数学运算上会逊于 TensorCore，而后者正是 TensorCore 内 MXU 的强项。为了让用户完全控制工作运行的位置，可以将 TorchTPU 流手动固定到 SparseCore，并且从性能和操作支持完整性的角度来看，该流与运行在 TensorCore 上的默认流并不相同。流还可以用于协调异步数据传输（使用 DMA 引擎），方式与 GPU 上类似。

我们的首要目标是保持与其他 torch.compile() 后端一致的用户体验，并因此共享相同的编译器前端，即 Dynamo。根据这份 PyTorch 指南，我们注册了一个新的 TorchDynamo 后端，其用法如下：

compiled_model = torch.compile(model, backend="tpu")Dynamo 通过传入 torch.fx.GraphModule 以及一组示例输入张量来调用我们的后端。此时，我们像在 eager 模式下一样对所提供的图进行追踪，但所有算子都被延迟执行，且不会在 TPU 上实际运行。这使我们能够为整个 FX 图生成并编译单个 StableHLO 程序。最后，我们的 Dynamo 后端返回一个 Python 可调用对象，该对象封装了对所生成 TPU 可执行文件的调用。请注意，作为一项临时限制，在 TPU 机器上调用 torch.compile 时，我们会自动设置 dynamic=False，因为支持 SymInt 需要有界动态性（下文将讨论），但该特性目前尚未默认启用。

TPU 与 GPU 的一个主要区别在于，XLA 传统上基于静态形状运行，以生成优化的硬件布局和融合内核。在实践中，这并非严重限制；当张量维度（如批大小或序列长度）在工作负载执行期间发生变化时，开发者通常通过预分配更大的张量并在计算前后手动填充/切片来解决。

为避免手动修改代码和重新编译的开销，我们正在积极开发，以通过 TorchTPU 中的 XLA 有界动态性原生支持动态形状。在有界动态性下，动态缓冲区在主模型模块内按最大上限分配，图只编译一次。计算使用有界缓冲区在 TPU 上执行，而 XLA 会跟踪逻辑维度大小并自动处理填充和切片，不会触发重新编译。

我们区分两类内核：

对于第 1 类，我们最初实现了一种方法：用 Pallas 编写 TPU 内核，然后以符号形状导出到 Mosaic，再通过一个 TorchTPU 步骤在调用时特化内核的形状。该机制被证明不如预期灵活，因此我们正在试验一种新方法，以编写支持动态形状的 Mosaic 内核。

对于第 2 类，我们开发了三种解决方案：

# Define a simple vector addition Pallas kernel def add_vectors_kernel(x_ref, y_ref, o_ref): o_ref[...] = x_ref[...] + y_ref[...] # Wrap the Pallas kernel into a JAX function @torch.tpu.pallas.jax_op("pallas::add_vectors") def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array: out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype) return pl.pallas_call(add_vectors_pallas_kernel, out_shape)(x, y) # Use it like a standard PyTorch function x = torch.tensor([1.0, 2.0, 3.0], device="tpu") y = torch.tensor([4.0, 5.0, 6.0], device="tpu") z = add_vectors(x, y) （生产就绪）用户可利用来自 Tokamax 的内核，这是一个高性能 TPU 内核库，现已可用于与 PyTorch 无缝互操作。Tokamax 提供关键算子的高度优化实现，包括混合专家构建模块（如分组矩阵乘法和 token gather/scatter）、注意力变体（如分页注意力、门控 deltanet、Kimi delta 注意力和多头潜在注意力）以及 TopK 计算。算子组合将随时间演进，以支持在最新硬件加速器上运行的最新 OSS 模型（例如，发布时支持 TPU v8i）。这些内核由稳健的集成基础基础设施提供支持，例如自动调优（如 tile 大小搜索）、低开销的基于 XProf 的基准测试和回归测试，以降低运维复杂性。虽然内核本身用 JAX/Pallas 编写，但该库的基础设施采用 torch.tpu.pallas.jax_op API 来透明地管理互操作，并暴露原生 PyTorch API 以方便开发者。 (Experimental) The user rewrites their GPU kernel in Helion, which enables targeting multiple accelerators with the same PyTorch-based DSL, as shown below # Define the kernel directly in Helion @helion.kernel(backend="pallas") def add_vectors(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: out = torch.empty_like(x) # Device execution loop over automatically partitioned tiles for tile in hl.tile(out.size()): out[tile] = x[tile] + y[tile] return out # Use it like a standard PyTorch function x = torch.tensor([1.0, 2.0, 3.0], device="tpu") y = torch.tensor([4.0, 5.0, 6.0], device="tpu") z = add_vectors(x, y) 分布式支持为确保与现有 PyTorch 分布式 API 和框架兼容，我们决定在 Process Group API 层面实现 TPU 集合通信，通过继承 c10d::Backend；这使我们能够在 eager 和编译模式下拦截集合操作，并无缝支持 PyTorch 分布式框架，如 DTensor、FSDPv2、SimpleFSDP、FairScale 等。与 GPU 上的 PyTorch 类似，我们遵循（并实际强制执行）每设备一进程的方法。从实现角度看，TPU 栈不提供单独的 TPU 间通信库，而是将 TPU 集合操作表示为 StableHLO 操作（stablehlo.all_gather、stablehlo.all_scatter、stablehlo.all_to_all 等），这与上述机制自然契合。

在 Process Group 层面拦截集合操作几乎让我们到达终点，除了与 GPU 相比 TPU 的一些显著特性：

%0 = "stablehlo.all_reduce"(%arg0) <{replica_groups=dense<[[4,0,1,3], [5,2,6,7]]>: tensor<2x4xi64>}> ({... MPMD Support. Historically, the TPU stack was built with SPMD in mind. Thanks to tools such as Shardy, distributed JAX programs are written for a single host, with no communication code included, and then are sharded at massive scale with user-provided directives, while the XLA compiler is fully responsible for inserting collective ops at the right points while maximizing the overlap between computation and communication. Unfortunately, the PyTorch ecosystem doesn’t adopt a clean SPMD approach (a notable exception is SimpleFSDP) and different processes may execute different ML ops leading to a collective op for reasons such as: 用户插入的以进程 rank 为条件的语句； 当张量维度不能被设备数量整除时，DTensor 对共享张量的不均匀拆分。 当包含集合操作的 TPU StableHLO 模块不一致时，可能会发生不可预测的死锁或内存损坏（这是由于生成的 TPU 代码中跨 rank 的同步标志 ID 或缓冲区偏移可能不匹配）。TorchTPU 对此的处理如下：

在 TorchTPU 中初始化分布式工作负载遵循标准 PyTorch 分布式语义，但需注意 TorchTPU 需要两个 TPU 特有的变量，而标准 PyTorch torchrun 不知道如何设置它们：

另一种方法是在 Python 中使用 torch.tpu.distributed.set_tpu_launch_env()，它可以在 Google Cloud 环境中设置这些变量。下面我们展示启动分布式 TorchTPU 工作负载的示例代码。

A. Worker 实现（train.py）

import os import torch import torch.distributed as dist def worker_main(): # 1. Initialize distributed process group on TPU using the tpu_dist backend dist.init_process_group(backend="tpu_dist") # 2. Query rank, world size, and construct target TPU device rank = dist.get_rank() world_size = dist.get_world_size() device = torch.device("tpu") print(f"[Rank {rank}/{world_size}] Initialized on {device}") # 3. Perform a distributed collective operation across TPU workers x = torch.ones(4, 4, device=device) * (rank + 1) dist.all_reduce(x, op=dist.ReduceOp.SUM) print(f"[Rank {rank}] Reduced sum: {x[0, 0].item()}") # 4. Clean up process group resources dist.destroy_process_group() if __name__ == "__main__": worker_main()B. 使用 torchrun 启动工作负载（CLI）

# Single host launch (spawns 1 process per TPU chip on an 8-chip node) eval $(python -m torch.tpu.distributed | sed 's/^/export /') torchrun --nproc_per_node=8 train.py最后，请注意给定节点的默认进程组是该节点的切片。

由于 TPU 独特的连接性，将模型的并行策略（TC、CP、FSDP、DP）与物理 TPU 硬件互连拓扑对齐非常重要，以最大化带宽并最小化网络延迟。回顾一下，TPU 使用分层网络架构：

一个好的规则是使用 TorchTPU 的 topology_aware_mesh() API，并将通信最密集的并行维度映射到最快、最连续的物理 ICI 维度，也就是将 mesh 维度从网络强度最低到网络强度最高进行排序。示例见下表：

以下是一个代码示例：

from torch.distributed.device_mesh import DeviceMesh # Single-Slice 2D Mesh Example (DP + TP) # Logical mesh shape: (dp=4, tp=2) on an 8-chip slice ordered from least # communication-intensive to most communication-intensive: mesh_shape = (4, 2) dim_names = ("dp", "tp") # Computes optimal rank layout aligned with physical ICI rings optimal_rank_tensor = torch.tpu.topology_aware_mesh( mesh_shape=mesh_shape, ) # Instantiate the DeviceMesh mesh_2d = DeviceMesh("tpu", mesh=optimal_rank_tensor, mesh_dim_names=dim_names, preserve_rank_order=True) dp_mesh = mesh_2d["dp"] tp_mesh = mesh_2d["tp"]作为未来工作，我们将探索一种 API，允许用户更直接地自定义其 mesh，从而在特殊场景下实现最优的集合通信性能。

TorchTPU 支持自定义 Pallas 内核，可通过 ICI 上的硬件 RDMA 直接进行芯片间通信，无需主机 CPU 参与。这使你可以编写高效的点对点通信内核（例如 Ring Attention、自定义 All-to-All / Shuffles、流水线并行移位或 Halo 交换），并将它们作为原生 PyTorch 函数调用。示例如下：

# 1. Pallas Kernel Body: Asynchronously sends local buffer to neighbor rank (rank + 1) def pallas_send_to_next_kernel(x_ref, y_ref, send_sem, recv_sem): rank = jax.lax.axis_index("x") world_size = jax.lax.axis_size("x") dest_id = (rank + 1) % world_size # Trigger hardware Remote DMA over TPU interconnect async_copy = pltpu.make_async_remote_copy( src_ref=x_ref, dst_ref=y_ref, send_sem=send_sem, recv_sem=recv_sem, device_id=(dest_id,), ) async_copy.start() async_copy.wait() # 2. JAX Wrapper with ShardMap and DMA Semaphores def build_jax_fn(mesh): def local_pallas_call(x): return pl.pallas_call( pallas_send_to_next_kernel, out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype), grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=0, in_specs=[pl.BlockSpec(x.shape)], out_specs=pl.BlockSpec(x.shape), scratch_shapes=( pltpu.SemaphoreType.DMA, pltpu.SemaphoreType.DMA, ), # send & recv semaphores ), )(x) return jax.shard_map( local_pallas_call, mesh=mesh, in_specs=P("x"), out_specs=P("x"), check_vma=False, ) # 3. Create the TorchTPU callable wrapper def get_send_to_next_op(world_size: int): mesh = torch.tpu.pallas.get_pallas_mesh(axis_names=("x",), mesh_shape=(world_size,)) jax_fn = build_jax_fn(mesh) return torch.tpu.pallas.jax_op( "pallas::send_to_next", build_jax_fn(mesh), mesh=mesh, input_partition_specs=[P("x")], output_partition_specs=[P("x")], )下表总结了如何编写通信型 Pallas 内核。

解决设备映射不一致问题。此前在 TorchTPU 中，使用底层芯片间操作（如 lax.axis_index 和 async_remote_dma / make_async_remote_copy）的 Pallas 内核基于原始硬件设备 ID 而非 PyTorch rank 来路由通信，与 JAX 保持一致。当硬件设备 ID 与 rank 顺序不一致时，环通信和跨设备 DMA 会被路由到非预期的物理设备，而在隔离子组上的执行会因分区数量不匹配而失败。

TorchTPU 通过强制逻辑寻址解决了这些不一致问题：

统一 ProcessGroup 集成。TorchTPU 现在提供获取 rank 对齐 JAX Mesh 的机制，使内核能够继承 PyTorch ProcessGroup 拓扑。这可以在全量世界和隔离子组配置下高效执行，无需手动修补状态或中断计算图。

TorchTPU 以开放方式开发，是 Google 与社区协作的成果，优先考虑透明度和参与度。该开发模式鼓励更广泛的 PyTorch 生态和行业合作伙伴贡献力量。

我们预计 TorchTPU 将服务于多样化的用户群体，包括研究机构和行业领导者。我们还计划将其与 Hugging Face 等平台集成，为广泛的开放源代码模型提供开箱即用的 TPU 支持。

为支持开发者，TorchTPU 提供全面的文档，包括 （参见 http://torchtpu.readthedocs.io/）：

文档使用可直接运行的 notebook 构建。

该项目利用 Bazel 进行稳健的构建和测试管理，测试框架仍在演进，并计划采用 Pytest。CI/CD 基础设施正在集成到官方 PyTorch CI HUD 中，以保持高可靠性并促进社区贡献。

TorchTPU 设计为可扩展，能够与关键 PyTorch 生态组件无缝集成。它支持通过 Pallas 和 Helion 进行自定义内核开发，并通过 vLLM、SGLang 和 TorchTitan 等项目为分布式训练和推理工作负载提供一等支持。这种模块化方法允许用户利用现有的 PyTorch 惯用法，同时针对 TPU 特定的硬件架构进行优化。

TorchTPU 提供与 PyTorch 生态的开箱即用兼容性，使开发者无需修改代码或重新编写即可在 Google Cloud TPU 上运行开放源代码模型。为验证这一承诺，我们构建了一个自动化测试框架，使用 Hugging Face 模型覆盖 NLP、视觉、音频、多模态和生成式扩散领域。该持续集成流水线以 CPU 和 GPU 基线为基准，对 TorchTPU 的覆盖范围、稳定性和数值等价性进行基准测试。我们在从 v5e 开始的 TPU 硬件代际上，以 eager 和编译模式持续评估约 5,000 个 Hugging Face 模型，涵盖数百种架构。

我们采用基于语料分层和均衡架构选择的稳健测试方法。我们的方法依赖于两级选择流程，旨在平衡实际采用与广泛的架构多样性，避免对细微变体进行冗余测试，同时确保真实生产权重得到充分评估。在第一级中，我们根据社区采用度、下载速度和趋势评分优先选择排名靠前的模型——例如 Qwen、Gemma、Mistral、DeepSeek、GLM、Kimi、Llama、Whisper、Stable Diffusion、FLUX 和 Vision Transformer——以确保关键社区热门模型完美运行。第二级是从第一级中选出的聚焦子集，确保我们从超过 500 个不同的模型家族中至少选择 1 个代表性模型。我们使用第二级在全面的 PyTorch 算子和编译器模式集合上验证相对于 GPU 基线的数值准确性。

验证过程从第 1 阶段开始，该阶段侧重于冒烟测试和降低验证，以保证运行时稳定性。在此初始阶段，流水线使用合成权重实例化模型，验证 eager 模式的前向执行，并确认图追踪、FX 到 StableHLO 的图生成以及 XLA 编译均能成功完成且无错误（例如，无崩溃或异常）。一旦确认运行时稳定性，第 2 阶段便使用领域特定的评估指标，对照既定基线评估数值正确性。对于因果 LLM，数学保真度通过困惑度（Perplexity = eL = exp((1/N) Σi=1N Li)）来衡量，若 TPU 与基线困惑度之间的相对差异保持在 0.10 以内，则测试通过。对于视觉、音频和编码器架构中的高维特征空间，我们采用双阈值相对 MSE 指标（Relative MSE = MSE(YTPU, YBaseline) / (Var(YBaseline) + 10−7)），要求相对 MSE ≤ 0.10 或绝对 MSE < 0.05，以防止基线输出接近常数。

下表总结了截至本文撰写时的模型覆盖结果：

对于有界动态性支持（针对 SymInts），我们在 839 个模型（来自 5K 语料库的因果 LLM 列表）上评估了正确性，并确认了 834 个（99%）模型的数值正确。

TorchTPU 致力于成为 PyTorch 的真正原生 TPU 后端。在我们看来，这一期望包括在运行纯 PyTorch、设备无关代码时具有良好的性能。我们通过两个示例来展示这一点。

第一个示例包含在 TorchTPU 仓库中，是一个自包含的、纯 PyTorch 实现的 Muse Glimmer 30B 模型，我们仅使用标准 PyTorch 模块和算子编写，例如 nn.Linear、nn.RMSNorm、F.scaled_dot_product_attention 和 torch.distributed 集合通信。我们有意没有使用任何模型特定的自定义内核（但仍使用了 PyTorch 内部内核）。在此示例中，我们展示了：

请注意，此示例使用 FSDP（通过原生 PyTorch 集合通信实现），world size 为 2，因为单个 TPUv7 芯片向 PyTorch 暴露两个设备（对应两个包含的 TensorCore）。

在第二个示例中，我们启用了激活检查点，它仅在 HBM 中存储前向传播激活的一个子集，并在反向传播计算需要时重新计算/重物化未存储的中间激活。

所考虑示例的一个决定性特征是，同一份代码在除指定加速器类型（例如编写 device="tpu" 而非 device="cuda"）之外没有任何实质性改动的情况下，即可在 TPU 和 GPU 上运行，并且在标准 PyTorch 优化与分片技术（如 torch.compile 和 FSDP）下在 TPU 上表现非常出色。

为进一步验证 TorchTPU 在纯 PyTorch 代码上的性能，我们还从 TorchTitan 项目中选取了一个外部产出的 Muse Glimmer 30B 实现，该实现最初是为 GPU/CUDA 编写的。在对其做最小限度适配到 TorchTPU 后（我们仅将 TorchTPU 尚不支持的 FlexAttention 替换为原生 F.scaled_dot_product_attention，并应用了 torch.compile 和 SimpleFSDP），TorchTPU 移植版实现了：

总之，所有这些示例表明，纯 PyTorch 模型实现，甚至最初为 GPU 编写的实现，只需简单的 PyTorch 改动并配合标准 PyTorch 优化，即可在 TPU 上表现足够出色，而无需借助自定义内核等高级的、设备特定的爬山技术。我们认为这是 TorchTPU 采用的一项重要特性。当然，对于追求峰值性能的生产工作负载，我们绝对建议采用基于 Pallas 内核的最优方案，以及 vLLM-TorchTPU、SGLang-TorchTPU 或 TorchTitan-TorchTPU 等生产级技术栈作为示例。

我们使用 TorchTitan 验证了 TorchTPU 可扩展至 1 个 TPU v7 POD。对于此测试，我们选择了 Qwen3 MoE 235B。作为初步冒烟测试，我们确认在 64 个 TPU v7 配置上，TorchTitan 配合 TorchTPU 的模型性能与 MaxText 上的 JAX 相当（PyTorch/TorchTitan 为 1828 TPS，JAX/MaxText 为 1937 TPS）。

下文讨论的数据基于以下设置计算得出：

在所有拓扑下，我们的数据证实了弱扩展：整体计算量呈线性扩展，如下所示，总每秒 token 数线性增长，而 MFU 基本保持平稳。

我们的数据还显示，在所考虑的各种并行级别下，损失曲线几乎完全相同。

在多切片 HSDP 配置中，前向传播不涉及任何跨切片通信，每个切片处理自己那份批次数据。但在反向传播中，一个切片先计算第 t 层的梯度，然后通过 all-reduce（在 DCN 上，使用 Megascale-XLA）对各切片的所有这些梯度求和。计算随后继续处理第 t-1 层的梯度。得益于 XLA，TorchTPU 能够将第 t+1 层的 all-reduce 与第 t 层的计算完全重叠，如下图所示。实现这种重叠是每个 TensorCore 的 TPS 从一片到八片几乎保持平稳的主要原因，如上表所示。

大语言模型（LLM）推理——尤其是针对大规模混合专家（MoE）架构——是 AI 云基础设施中要求最高、成本最关键的工作负载之一。在 SemiAnalysis 发布的一项独立基准评估中（TPU Inference Externalization Full Steam Ahead — InferenceX），Google 的 TPU v7 搭配 vLLM-TorchTPU 软件栈，交付了具有行业竞争力的每美元性能。该基准提供了对 TorchTPU 软件在大规模混合专家（MoE）工作负载上就绪程度的第三方推理审计，表明面向 TPU 的新 vLLM 软件栈在社区公开可用的 TPU 系统上提供了强大的成本效率和高吞吐量。

在此基础上，Google 最近提交了 Qwen 3.5 在 Ironwood TPU 上的首批 AgentX 性能结果。AgentX 捕捉了智能体工作负载的四个要素——多轮、长上下文、高前缀复用和子智能体突发。TPU 已能够证明其即将推出的 vLLM-TorchTPU 软件栈能够成功处理智能体工作负载。仅用 4 块芯片，TPU v7 就实现了帕累托前沿的全面领先。

这些基准结果反映了在 TPU 系统上实现的强大、真实的在线服务效率，突显了在 Google 定制设计的 ML 芯片上部署大规模模型的巨大经济可行性。

要实现这些有竞争力的性能和 TCO 结果，需要通过自定义 Pallas 内核和有针对性的运行时调优来解决关键的软硬件对齐瓶颈，下表列出了几个示例：

这些结果突显了现代服务基础设施中的 TPU 优化，表明 TCO 和软硬件对齐决定了真实世界的服务效率。最终，带有 TorchTPU 后端的新 vLLM-TPU 使工程团队能够在公开可用的 TPU 基础设施上使用标准 PyTorch 惯用法部署现成的开放权重模型。

TorchTPU 与 TorchTitan 原生集成——TorchTitan 是 PyTorch 面向大规模分布式训练的参考平台——在稠密大语言模型（MuseGlimmer 30B、Llama 3 8B–405B）、稀疏混合专家模型（Qwen 3 235B MoE）以及生成式扩散 Transformer（FLUX.1-dev 12B）上提供统一的高性能训练环境。在 TPU v6e（Trillium）和 TPU v7x（Ironwood）集群上以 BF16 精度进行评估，多切片执行通过 torchrun 在 Google Kubernetes Engine（GKE）上无缝编排。TorchTitan 作为开源项目可在此 GitHub 仓库中找到。

分布式策略与分片架构。高效的多加速器训练要求模型并行策略与 TPU 的物理互连拓扑之间实现无缝对齐。为此，我们采用 2D+ DeviceMesh 映射，将通信密集的维度——例如张量并行和专家并行——与连续的高带宽芯片间互连（ICI）对齐，同时将复制的数据并行流量路由到较慢的跨切片数据中心网络（DCN）。这一策略通过我们与 SimpleFSDP 的集成来实现，后者提供符合 SPMD 规范的操作。为进一步降低延迟，我们将 XLA 运行时配置为把集合通信操作（如 All-Gather 和 Reduce-Scatter）直接卸载到片上 SparseCore 引擎，从而有效地将通信与计算密集的 TensorCore MXU 数学运算重叠。最后，我们通过使用 HLO diffing 工具和 @spmd_safe 注解来确保确定性执行并防止同步死锁，这些工具和注解消除了 torch.compile 下不必要的图中断。

性能与多节点扩展。通过系统性的逐层优化，TorchTPU 上的 TorchTitan 为 FLUX 模型实现了卓越的性能和多节点扩展效率。从开箱即用的 eager 基线步进时间 11.49s（0.26 samples/s）开始，应用 DeferAndFuse eager 模式将步进延迟改善至 9.95s（1.15× 加速）。引入 SimpleFSDP 分片进一步将执行加速至 6.52s（1.76× 加速），而集成自定义 Pallas SplashAttention 内核则将步进时间降至 5.55s（2.07× 加速）。启用完整的 torch.compile(backend=”tpu”) 带来显著的性能飞跃，达到 3.02s（3.80× 加速），而应用目标模型窥孔重写则产生最终优化后的步进时间 2.78s（1.08 samples/s）。在多主机集群中，TorchTitan 展现出接近理想的弱扩展和强扩展：在 FLUX.1-dev 上，使用 SimpleFSDP 从单个 v6e-8 节点（8 芯片）扩展到 v6e-256 切片（64 个节点共 256 芯片），每样本步进时间仅从 0.92s 略微增加到 1.00s（总步进时间 3.01s），在光互连 ICI 网络上展现出接近线性的效率。在更大规模下，TorchTitan 在 TPU v7x（Ironwood）上以 BF16 精度、8K 上下文窗口在 1,024 个芯片上运行 Llama 3 405B，实现 500 tokens/sec/chip 和 50+% MFU。同样，在 8,192 个 TPU v7x 芯片上评估 Qwen3 235B MoE 也展现出高效扩展，在 BF16 精度、4K 序列长度下达到 28% MFU。

关键架构优化。为在多样化的模型架构中维持峰值计算密度并消除硬件瓶颈，我们采用了多项系统和内核优化：

Ongoing development for TorchTitan on TorchTPU targets four strategic priorities: expanding list of models optimized for TPU; integrating FP8 and FP4 low-bit quantization via torchao; and supporting end-to-end Reinforcement Learning (RL) by coupling training loops with vLLM/SGLang sampling backends in VeRL.

Recommendation and ranking architectures—such as Meta’s Deep Learning Recommendation Model (DLRM)—face critical memory and compute bottlenecks due to high-cardinality embedding tables. On standard GPU clusters, managing these tables may require complex software-based sharding via DistributedModelParallel (DMP) and separate optimizer updates, which often introduce latency. To bridge this gap, Google collaborated with the Meta PyTorch team to upstream a native TPU backend directly into the TorchRec repository, allowing recommendation models to bypass host-side compute bottlenecks and tap directly into TPU SparseCore (SC) hardware.

Active development focuses on further saturating bandwidth through TensorCore/SparseCore pipelining—interleaving MLP execution with embedding lookups—and advanced table-stacking techniques to unify memory buffers and reduce kernel overhead. Additionally, ongoing work aims to expose lower-level SparseCore embedding lookup operations (e.g. SparseMapRow, SparseGather etc) via ATen kernels, enabling users to implement custom lookup algorithms and advanced sharding strategies. This integration solidifies the TPU as a high-efficiency engine for the next generation of recommendation architectures.

With TorchTPU now publicly available on GitHub, this initial open-source release represents only the beginning of our development journey. As we engage directly with the broader PyTorch community, our roadmap for the next 6 to 12 months focuses on production hardening, ecosystem expansion, and continually integrating community feedback to bridge the gap between TPU hardware capabilities and PyTorch-native developer workflows. We are prioritizing four strategic pillars to drive this evolution:

TorchTPU is built on the belief that a successful ecosystem thrives through transparent, collaborative engineering. As we open-source our four foundational repositories (TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, and TorchTitan) at the October 2026 PyTorch Conference, we actively seek feedback on model architecture prioritization and workload sequencing. We invite developers, researchers, and partners to collaborate with us via GitHub or our active TPU Builders program.

Claudio Basile, Michael Voznesensky, Gunhyun Park, Ben Wilson, Parker Holloway, Zhanyong Wan, Vlad Belous, Manish Khanna, Paul Ganssle, Kevin Gleason, Andrew Dame, Brendan Folie, Bhavya Bahl, Haifeng Jin, Pedro Goncalves Mokarzel, Arian Arfaian, Luke Boyer, Chi Zhang, Yiling Yuan, Adam Mainz, Aditya Patel, Huangze Tang, Jay Gan, Lizzie Liang, Will Froom, Tushar Agarwal, Karlo Basioli, Theotime Combes, Emilio Cota, Minos Papadopoulos, Tudor Pipernea, Yarong Mu, Qi Zhou, Yaoshiang Ho, Aleksey Vlasenko, Alexander Brauckmann, Ziyin Huang, Jack Oh, Jerry Hsu, Tanvi Bajpai, Jialei Chen, Wenjun Gu, Chintan Shah, Chris Chan, Sayce Falk, Chris Jones, Rob Torres, Kat Ko, Lee Howes, Freya Azad, Andi Gavrilescu, Joe Pamer, Bill Jia

Plus many, many others…

The post Introducing TorchTPU appeared first on PyTorch.

来源：PyTorch：Blog · pytorch.org