文章目录Google 与 Meta 发布 TorchTPU：PyTorch 原生 TPU 后端
Gemini 使用指南封面
Gemini 使用指南封面
Gemini 中文指南 专题文章

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

Google 与 Meta 发布 TorchTPU，为 Google TPU 提供原生 PyTorch 后端。本文完整还原其架构、编译、精度、分布式通信、Pallas 内核、模型覆盖和训练推理基准，帮助开发者评估 PyTorch 工作流迁移到 TPU 的方式与边界。

本文根据 PyTorch 官方博客 Introducing TorchTPU 的中文译文整理，尽量保留原文的章节、代码和基准数据。文中的性能结果来自项目团队或文中引用的第三方评测，不能直接视为本站实测。

TL;DR

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

推出 TorchTPU

Google TPU 已在 Google 内外被用于大规模训练和部署大型机器学习模型。它们专为机器学习设计，在性能、能耗和成本方面具有优势。尽管 TPU 历史上与 TensorFlow 和 JAX 生态密切相关，Google 于 2025 年年中启动了 TorchTPU，并与 Meta 的 PyTorch 团队合作，让 PyTorch 原生运行在 TPU 上。

在 Python 框架下编程 Google TPU，需要使用 XLA 编译器栈。2020 年，Google 推出了首个面向 TPU 的 PyTorch 编译版本 PyTorch/XLA，基于 Lazy Tensor。这个早期集成具有奠基意义，但难以跟上 PyTorch 由社区驱动的快速演进：用户需要侵入式修改源代码以显式定义可编译区域，编译延迟较高，不支持 PyTorch 原生的命令式（eager）执行模式，在支持 PyTorch 生态新兴的分布式计算原语方面也相对滞后。

与 PyTorch/XLA 和 TorchAX 等将 PyTorch 强行适配到 JAX 运作方式的尝试不同，TorchTPU 被设计为 PyTorch 的原生 TPU 后端。eager 执行是首要优先级，现有 PyTorch 工作负载只需极少的代码改动即可运行，例如把 .to("cuda") 改成 .to("tpu")。

项目启动时制定了以下原则。

  • PyTorch Citizenship（遵循 PyTorch 生态）： 运行时不依赖 JAX；提供稳健的 eager 模式；基于 TorchDynamo（而不是 Lazy Tensor）实现 torch.compile()，目标是接近 JAX compile 的性能；在 PyTorch 栈底层（例如 Process Group API）提供 TPU 集合通信，支持 DDP、FSDP、DTensor 等分布式惯用法；无缝集成 PyTorch 生态，并接受社区参与。
  • Focus on the User（以用户为中心）： 将工作负载移植到 TPU 应当像把 model.to("cuda") 改为 model.to("tpu") 一样简单。这个改动本身不保证 TPU 峰值性能，因为模型形状、分片和算子可能偏向 GPU 架构，但模型应能够开箱运行；运行时和编译器错误应给出友好且可操作的提示；并与 torch.profiler、DTensor、Ray 等 PyTorch 社区正在使用的工具和框架集成。
  • Improve the State of the Art（推进技术边界）： 构建能够在大规模工作负载上与 GPU 上 PyTorch 竞争的方案；为 eager 执行引入 JIT 编译方法，以运行时优化弥合逐算子执行与编译执行之间的性能差距；通过有界动态性、减少形状变化导致的 XLA 重新编译，以及探索可提前预编译的动态 TPU 内核，改进 TPU 编程栈。

TorchTPU 仍在积极开发，已经被部分大型 Google Cloud 客户用于生产环境中的推理和训练工作负载。项目源码位于 google-pytorch/torch_tpu。截至本文原文发布时，项目已经完成以下工作：

  • 支持 eager 和编译模式，包括新的 DeferAndFuse eager 模式；
  • 支持约 450 个 ATen 算子（跳过具有结构化或复合委托的算子），覆盖主要深度学习架构所需的 Core ATen 前向和反向传播；
  • 支持通过 JAX/Pallas 和 Helion DSL 编写自定义内核；
  • 验证与 Hugging Face Transformers v5、Diffusers 和 Kernels 的内置兼容性，在下载量最高的约 5,000 个模型上开箱运行；
  • 验证基于 DTensor 的分片；
  • 验证多主机推理（vLLM、SGLang）和微调/训练工作负载（TorchTitan），并采用 Ray 编排；
  • 验证与 vLLM、SGLang、TorchTitan、Ray 和 TorchRec 的集成。

下面按原文顺序介绍 TorchTPU 的架构和初步结果。

架构深入探讨

TPU 入门

Google TPU 是专门用于执行深度学习大规模张量运算的 ASIC。从硬件架构看，一个 TPU 芯片根据代际不同包含：

  • 一个或多个 TensorCore，每个 TensorCore 包含一个或多个矩阵乘法单元（MXU）、一个向量单元（VPU）和一个标量单元。根据代际，一个 MXU 可以包含 256×256 个乘加器，以脉动阵列排列。
  • 一个或多个 SparseCore，用于 gather、scatter、embedding 和集合通信等稀疏操作。它们可以与 TensorCore 并行运行，实现计算与计算、计算与通信的重叠。
  • 大容量高带宽内存（HBM）。例如 TPU 8i 为 288 GB，带宽为 8.6 TB/s，可存储机器学习数据并与 CPU 共享。
  • 双向芯片间互连（ICI）。TPU 8t 和 8i 的 ICI 为 19.2 TB/s，允许同一 TPU Slice 内的芯片高速通信。Slice 是 TPU Pod 的逻辑子组，由预配置数量的 TPU 芯片组成，TPU 7、TPU v8t 使用 3D torus，TPU 8i 使用 Boardfly 配置连接。
  • 双向数据中心网络（DCN）。TPU 8t 的 DCN 为 400 GBps，用于跨 Slice 的 TPU 芯片通信。

TPU 通过 XLA 编译器栈的两个入口进行编程：StableHLO 是面向机器学习算子的高级 IR，Mosaic 是用于自定义内核的低级 IR。实际 TPU ISA 是专有的，可能变化且不公开；在实践中，可以把 StableHLO（更确切地说是 Mosaic）理解为类似 GPU 的 PTX。

C++ 程序可以实例化 PjRT Client 访问本地 TPU。PjRT API 负责将 MLIR（StableHLO 或 Mosaic）编译成 TPU 可执行文件、启动可执行文件以及在 TPU 之间复制数据。TPU 执行是异步的，CPU 可以等待关联的 PjRTBuffer 获取结果。

TPU 编程模型有几个重要特点：

  • TensorCore 可以看作带一个或多个大规模 MXU 的单线程处理器。要达到峰值性能，模型形状需要足够宽，以充分利用 MXU；这与 GPU 不同，GPU 可以通过并行启动很多较小内核来提高利用率。
  • XLA 可以把稀疏运算和集合运算卸载到 SparseCore，以实现更高效或并行的执行。这个过程可以自动发生，也可以通过编译后的 StableHLO 模块中的 MLIR 注解控制。
  • 一个 TPU 设备只能由一个用户空间进程访问。为与 PyTorch 生态集成，TorchTPU 的分布式工作负载采用每设备一进程模型，与 JAX 的每主机一进程模型不同。由于一个 TPU 只能被一个用户空间进程访问，RL 共置目前不能开箱即用，项目正在研究如何放宽或绕过这一限制。
  • StableHLO 遵循 SSA 范式，不允许原地修改，写操作总是产生新值（至少在语义上如此）。TorchTPU 需要通过视图分解和反转提供熟悉的 PyTorch 视图语义，包括修改视图的能力。
  • XLA 编译器和 PjRT 运行时通常假设编译时已知张量形状。TorchTPU 正在加入有界动态性：编译时知道形状维度的最大边界，但实际维度可以到运行时才确定。
  • 集合操作被降级为 StableHLO 操作交给 XLA，而不是使用 NCCL 等独立库。

可用性与自动加载机制

在安装 TorchTPU wheel 的 TPU 机器上导入 PyTorch 后，得益于 PyTorch 的自动加载机制，只需几行代码即可运行 TPU 工作负载。下面是 GPU、显式 TPU 和自动选择设备的并排对比。

之前：GPU

import torch
import torch.distributed as dist

def main():
    device = torch.device("cuda")
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

之后：TPU

import torch
import torch.distributed as dist

def main():
    device = torch.device("tpu")
    dist.init_process_group(backend="tpu_dist")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

替代方案：自动选择 TPU 或 GPU

import torch
import torch.distributed as dist

def main():
    device = torch.accelerator.current_accelerator()
    dist.init_process_group()
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

数值精度

TorchTPU 适配标准 PyTorch 精度 API。与所有高性能机器学习架构一样，TPU 针对特定代际的原生数据类型进行了优化。TPU 的原生矩阵乘法格式是 bfloat16（BF16）：MXU 使用 BF16 输入执行乘法，并以 FP32 进行中间累加。对于逐元素加法、激活、LayerNorm 和 Softmax 等非收缩运算，TorchTPU 使用原生执行 IEEE-754 FP32 的 TPU 向量单元（VPU）。

TorchTPU 将 PyTorch 精度模式映射到 XLA/StableHLO：

  • medium： 矩阵乘法将输入转换为 MXU 原生 BF16，累加以 FP32 执行。TorchTPU 默认采用此模式，适合绝大多数深度学习工作负载。
  • high： 使用三遍算法，将输入拆成高、低分量并计算三个交叉项，以约三倍计算成本恢复大部分 FP32 精度。
  • highest： 使用六遍算法计算所有交叉项，在 MXU 上产生约六倍计算开销，接近完整 IEEE FP32 精度。

官方建议大多数任务使用默认的 medium。显式选择 high 或 highest 会触发 TORCH_WARN_ONCE 性能警告；只有确认存在精度导致的收敛问题时，才应评估更高精度。除全局设置外，TorchTPU 还提供线程本地的 torch.tpu.precision 上下文管理器，用于只对数值敏感层提高精度。

维度CUDA 上的 PyTorchTPU 上的 PyTorch（TorchTPU）
默认精度模式highest（纯 FP32）medium（Precision.DEFAULT）
快速数学加速器Tensor Cores（TF32）脉动 MXU（BF16）
快速数学格式1 位符号 + 10 位尾数 + 8 位指数（19 位 TF32）1 位符号 + 7 位尾数 + 8 位指数（16 位 BF16）
累加精度highest/high 为 FP32，medium 为 TF32所有模式均为 FP32
更高精度机制专用 FP32/TF32 硬件单元三遍或六遍 MXU 分解算法
high/highest 性能影响基线速度或适度减慢（约 1–2 倍）high 约 3 倍 MXU 成本，highest 约 6 倍
作用域 API全局 torch.set_float32_matmul_precision全局设置 + torch.tpu.precision 上下文管理器

Eager 执行、StableHLO 降级和 JIT 编译

TorchTPU 按照 PyTorch 后端扩展指南扩展 PyTorch dispatcher，并注册为 PrivateUse1。它使用标准 at::Tensor 机制定义设备特定的存储实现（通过 c10::StorageImpl 中的 c10::UniqueVoidPtr）。TorchTPU 的 eager 执行可以用多种模式将 ATen 算子分发到 TPU：立即逐算子分发、延迟到自动确定的物化点后聚合编译，以及介于两者之间的模式。

TorchTPU 提供三种 eager 模式，可通过环境变量或上下文管理器设置：

  • Strict（默认）： 一次分发一个 ATen 算子，算子在 TPU 上与 CPU 异步执行，行为类似 GPU 的 eager 模式。
  • Debug（TORCH_TPU_LAUNCH_BLOCKING=1）： 一次分发一个 ATen 算子，但 TPU 与 CPU 同步执行，速度很慢，主要用于调试，并提供 TorchTPU 边界检查，类似 GPU 的 CUDA_LAUNCH_BLOCKING。
  • DeferAndFuse（TORCH_TPU_DEFER_AND_FUSE=1）： 延迟多个 ATen 算子并融合为更大的块。这是 TorchTPU 的 eager/编译混合模式，官方称它比 Strict Eager 快 50% 或更多，同时降低峰值 HBM 内存，因为融合减少了临时结果的存储。

以 at::mul.out 为例，DeferAndFuse 模式会创建一个包含 StableHLO 构建器的 Deferred Op。构建器在调用时返回实现 aten::mul 语义的 StableHLO 函数（隐式广播和 dtype promotion 由其他步骤处理）：

func.func @main(%arg0: tensor<4x16xbfloat16>,
                %arg1: tensor<4x16xbfloat16>)
    -> tensor<4x16xbfloat16> {
  %0 = stablehlo.multiply %arg0, %arg1 : tensor<4x16xbfloat16>
  return %0 : tensor<4x16xbfloat16>
}

当收到 mytensor.to("cpu")、mytensor.item() 等回调时，TorchTPU 会：

  1. 分析累积的、产生 mytensor 的延迟算子区域 R；
  2. 进行活跃性分析，确定需要物化的算子及产生 mytensor 的算子。集合通信等有副作用的算子始终保留且保持顺序；
  3. 为需要物化的算子图 G 创建编译缓存键；
  4. 在编译缓存中查找 G。如果已有缓存，就把可执行文件分发到 TPU；如果没有，就异步编译后分发；
  5. 等待 mytensor 的 TPU 值就绪，并将结果从 TPU HBM 复制到目标 at::Tensor 的 CPU 存储。

表面上看，CPU 似乎要等到 .to("cpu") 才启动 TPU 执行。对重复的真实工作负载，初始迭代填充编译缓存后，CPU 只需追踪和分发算子，可以始终领先于忙于数值计算的 TPU。为了加速重复算子序列（例如一组 Transformer block）的物化，TorchTPU 还提供 Repeated Ops Heuristic，在到达 TPU 到 CPU 传输之前检测重复序列并提前触发物化点。

编译缓存

TorchTPU 使用编译缓存避免不必要的 XLA 编译。缓存分为三层：

  1. 进程本地内存缓存： 同一进程共享编译结果，速度最快；
  2. 机器本地共享内存： 支持同一主机的跨进程共享，适合多设备训练和推理；
  3. 远程持久缓存： 由远程文件系统（例如 Google Cloud Storage）支持，可跨主机、程序和作业重启复用，适合 Colab 和数据中心环境。

第一层始终启用。默认情况下，第二层对分布式程序启用，对单设备程序禁用，用户可以覆盖。第三层默认关闭，但可以显式指定共享存储路径启用。

第一层中的每个可执行文件是可直接调度到 TPU 的 C++ 对象；第二层和第三层将可执行文件保存为序列化二进制字符串。为提高性能，缓存键不是 StableHLO 图本身，而是由三个 64 位整数构成的元组：

  • 无形状键： 记录计算图结构，不记录张量输入输出大小；
  • 维度键： 记录张量大小；
  • 编译选项键： 记录传给 XLA 的选项，以及 TorchTPU、libtpu 和 PjRT 运行时版本。

无形状键和维度键合称图键。理论上不同图、编译器选项和组件版本之间可能发生冲突，但项目认为 192 位键空间在实践中足够大。

视图语义

PyTorch 与 StableHLO 在可变性和共享内存视图上存在显著差异。

PyTorch 将每个 Tensor 看作指向线性字节数组的指针，并在 Tensor 层保存 stride 元数据。多个 Tensor 可以共享同一底层字节数组，因此可以对任意 Tensor 原地写入，并从其他 Tensor 看到修改。视图操作通常只需在 dispatch 阶段更新 stride 元数据；as_strided 还可以定义任意复杂的访问模式。

StableHLO 是纯函数式语言，张量值创建后不可变，物理布局由编译器决定，用户不能直接操作 stride，也不支持原地写入。选择性更新需要用 dynamic_update_slice 把新值合并到原张量的子切片。

在存在写入时实现正确的视图语义并不简单。TorchTPU 通过视图分解、视图反转和延迟操作桥接两种语义，因此可以支持原本难以支持的 PyTorch 模式，例如分布式训练中使用拼接参数数组。

流与 SparseCore

优化模型性能时，PyTorch 用户可能希望同时执行：

  • 计算操作；
  • 集合通信操作；
  • 主机与设备之间的数据传输。

GPU 通常用不同命令队列让设备将工作分派到不同硬件单元。PyTorch 用 Stream 暴露这些队列，用户代码和框架代码（例如 FSDPv2）都可能使用它们。

TorchTPU 也实现了 PyTorch Stream API，但不可能在所有情况下实现 GPU 上相同的并行方式。一个 TensorCore 本质上是单线程执行处理器，在同一 TensorCore 上一次只能分派一个计算；GPU 则有许多可以并行运行的 SM。TPU 的优势在于一个 TensorCore 关联一个或多个 DMA 引擎以及一个或多个 SparseCore 引擎：DMA 可负责主机与设备传输并与 TensorCore 并行，SparseCore 可处理 gather/scatter 等稀疏操作和集合通信。

用户可以把 TorchTPU 流手动固定到 SparseCore。它与运行在 TensorCore 上的默认流在性能特征和操作支持完整性上并不相同。流也可以像 GPU 一样协调异步 DMA 数据传输。

编译模式

TorchTPU 的目标是保持与其他 torch.compile() 后端一致的用户体验，因此复用 TorchDynamo 编译器前端。按照 PyTorch 自定义后端指南，使用方式如下：

compiled_model = torch.compile(model, backend="tpu")

Dynamo 向后端传入 torch.fx.GraphModule 和一组示例输入张量。TorchTPU 像 eager 模式一样追踪这些算子，但将所有算子延迟执行，不在 TPU 上实际运行，从而为整个 FX 图生成并编译单个 StableHLO 程序。后端最终返回一个封装 TPU 可执行文件调用的 Python 可调用对象。

现阶段，在 TPU 机器上调用 torch.compile 会自动设置 dynamic=False。支持 SymInt 需要有界动态性，而该能力当时尚未默认启用。

有界动态性

TPU 与 GPU 的主要区别之一是，XLA 传统上基于静态形状生成优化的硬件布局和融合内核。实践中，开发者通常通过预分配更大的张量并手动填充和切片来处理批大小或序列长度变化，但这会带来代码修改和重新编译开销。

TorchTPU 正在通过 XLA 有界动态性原生支持动态形状：动态缓冲区在主模型模块内按最大上限分配，计算图只编译一次；实际计算使用有界缓冲区，XLA 跟踪逻辑维度并自动处理填充和切片，而不触发重新编译。

内核

TorchTPU 将内核分为两类：

  1. PyTorch 框架提供的内核，例如 torch.nn.functional.scaled_dot_product_attention；
  2. 用户提供的、用于加速模型特定部分的自定义内核。

对第一类，早期方案是用 Pallas 编写 TPU 内核，以符号形状导出到 Mosaic，再通过 TorchTPU 在调用时特化形状。该机制灵活性不足，项目正在试验支持动态形状的 Mosaic 内核。

对第二类，项目提供三种路径。

Pallas（生产就绪）

用户可以用 Pallas 编写内核，再通过 TorchTPU API 与 PyTorch 集成：

import jax
from jax.experimental import pallas as pl
import torch

# 定义简单的向量加法 Pallas 内核
def add_vectors_kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] + y_ref[...]

# 包装成 JAX 函数，再注册为 PyTorch TPU 算子
@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_kernel, out_shape=out_shape)(x, y)

# 像普通 PyTorch 函数一样使用
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（生产就绪）

来自 Tokamax 的内核库已经可以与 PyTorch 无缝互操作，提供混合专家模块（分组矩阵乘法、token gather/scatter）、分页注意力、门控 DeltaNet、Kimi Delta Attention、多头潜在注意力和 TopK 等高度优化的 TPU 实现。内核组合会随硬件和开源模型演进；其基础设施提供自动调优、低开销 XProf 基准测试与回归测试。内核本身用 JAX/Pallas 编写，底层通过 torch.tpu.pallas.jax_op 管理互操作，并暴露原生 PyTorch API。

Helion（实验性）

用户可以在 Helion 中重写 GPU 内核，用同一个 PyTorch 风格 DSL 面向多个加速器：

import torch
import helion
import helion.language as hl

@helion.kernel(backend="pallas")
def add_vectors(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    # 对自动分块的 tile 执行设备循环
    for tile in hl.tile(out.size()):
        out[tile] = x[tile] + y[tile]
    return out

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 和框架，TorchTPU 在 Process Group API 层实现 TPU 集合通信，继承 c10d::Backend。这样可以在 eager 和编译模式下拦截集合操作，并支持 DTensor、FSDPv2、SimpleFSDP、FairScale 等框架。

与 GPU 上的 PyTorch 类似，TorchTPU 强制采用每设备一进程。TPU 栈没有单独的 TPU 间通信库，而是将集合操作表示为 StableHLO 操作，例如 stablehlo.all_gather、stablehlo.all_scatter 和 stablehlo.all_to_all。

TPU 分布式运行有几个与 GPU 不同的地方：

  • 初始化序列不同： GPU 上由 torchrun 启动的进程可以看到主机上的所有设备；TPU 进程只能看到一个设备，并不知道主机上有多少设备。初始化前必须知道是否要运行分布式程序、world size 和 TPU 类型，因为 TPU 类型会影响拓扑。
  • 默认 Process Group 信息： torch.distributed.all_reduce 只需要当前子组，而 stablehlo.all_reduce 还需要所有子组的信息。例如子组 (4, 0, 1, 3) 和 (5, 2, 6, 7) 执行 all-reduce 时，生成的 StableHLO 需要同时知道两个 replica group：
%0 = "stablehlo.all_reduce"(%arg0)
  <{replica_groups = dense<[[4, 0, 1, 3], [5, 2, 6, 7]]> : tensor<2x4xi64>}> ({
  ...
})
  • MPMD 支持： TPU 栈历史上以 SPMD 为中心，JAX 程序通常由单主机程序和用户分片指令组成，XLA 负责插入集合操作。PyTorch 并不完全遵循这种方式：不同进程可能因为 get_rank() 条件分支，或 DTensor 在维度不能整除设备数时进行不均匀拆分，而执行不同的算子。

含集合操作的 TPU StableHLO 模块如果在不同 rank 上不一致，可能产生死锁或内存损坏，因为生成的 TPU 代码中同步标志 ID 或缓冲区偏移可能不匹配。TorchTPU 提供三种处理方法：

  • Skinny Collectives（默认）： 每个集合操作编译为独立的“skinny” StableHLO 模块，保证 rank 间缓冲区和同步标志 ID 一致。代价是 XLA 不能将通信与计算重叠，端到端可能带来 30%–40% 的性能损失。项目正在把 Skinny Collectives 放到绑定 SparseCore 的独立流上，以挽回大部分损失。
  • @spmd_safe（可选）： 如果用户确定一段代码是 SPMD（例如不存在按 get_rank() 条件执行的分支），可以用 @spmd_safe 标记。torch.compile 会消除集合通信周围的图断点，生成包含集合通信的大 StableHLO 模块，让 XLA 移动通信位置并最大化计算重叠。SimpleFSDP 和手动放置集合操作时已经观察到明显收益，但错误使用可能导致脆弱行为。
  • Compilation Handshake（可选）： 当 torch.compile 触发包含通信的模块编译时，所有副本先进行共识轮次，验证 FX 图一致，再允许关闭集合操作周围的图断点。环境变量 TORCH_TPU_SPMD_HANDSHAKE_STAGE 可设为 OFF（默认）、COMPILE_STAGE 或 DISPATCH_STAGE；后者只适用于预期存在 rank 特定重新编译的情况，开销更大。相关 RFC 见 torch_tpu issue #2886。

分布式初始化

TorchTPU 使用标准 PyTorch 分布式语义，但需要标准 torchrun 不会自动设置的两个 TPU 变量：

  • TORCH_TPU_TOPOLOGY：TPU Slice 的物理 mesh 几何结构，例如 2×2×1 Slice 写作 2,2,1；
  • TORCH_TPU_SLICEBUILDER_ADDRESSES：构建 ICI 网络 mesh 的开放端口或地址列表，每个设备一个，例如 localhost:8471,localhost:8472,...。

也可以在 Python 中使用 torch.tpu.distributed.set_tpu_launch_env()，让 Google Cloud 环境设置这些变量。下面是一个最小分布式工作负载示例。

A. Worker 实现（train.py）

import torch
import torch.distributed as dist

def worker_main():
    # 1. 使用 tpu_dist 初始化 TPU 进程组
    dist.init_process_group(backend="tpu_dist")

    # 2. 查询 rank、world size，并构造 TPU 设备
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = torch.device("tpu")
    print(f"[Rank {rank}/{world_size}] Initialized on {device}")

    # 3. 在 TPU worker 间执行集合通信
    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. 清理进程组
    dist.destroy_process_group()

if __name__ == "__main__":
    worker_main()

B. 使用 torchrun 启动

# 单主机启动：在 8 芯片节点上每个 TPU 芯片启动一个进程
eval "$(python -m torch.tpu.distributed | sed 's/^/export /')"
torchrun --nproc_per_node=8 train.py

给定节点的默认进程组对应于该节点的 TPU Slice。

最优设备 Mesh 映射

TPU 的独特连接性意味着模型并行策略（TC、CP、FSDP、DP）必须与物理 TPU 互连拓扑对齐，才能最大化带宽、最小化延迟。TPU 使用分层网络：

  • ICI： 以 Torus mesh 直接连接芯片，带宽从数十到数百 GB/s；
  • DCN： 连接不同 TPU Slice 的主机级网络，带宽更低、延迟更高。

一般应使用 TorchTPU 的 topology_aware_mesh()，将通信最密集的并行维度映射到最快且最连续的 ICI 维度，即按网络通信强度从低到高排列 mesh 维度。

并行类型通信模式频率 / 数据量最优物理放置
张量并行（TP）all_reduce、reduce_scatter每个 Transformer 层前向和反向最快的 ICI 环，同一 tray/host 的连续 1-hop 芯片
上下文/序列并行（CP/SP）all_to_all、Ring Attention每个 attention 层都很高Slice 内本地 ICI 轴
专家并行（EP）all_to_all每个 MoE router 步骤本地 ICI 轴
FSDP / 完全分片 DPall_gather、reduce_scatter每个模块/层一次本地 ICI 轴
复制数据并行（DP）all_reduce（仅梯度）每个训练步骤一次ICI 或跨 Slice 的 DCN

代码示例：

import torch
from torch.distributed.device_mesh import DeviceMesh

# 单 Slice 的二维 Mesh：8 芯片上使用 (dp=4, tp=2)
# 从通信较少到通信较多排序
mesh_shape = (4, 2)
dim_names = ("dp", "tp")

# 计算与物理 ICI 环对齐的最优 rank 布局
optimal_rank_tensor = torch.tpu.topology_aware_mesh(
    mesh_shape=mesh_shape,
)

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"]

项目后续将探索更直接的 mesh 定制 API，以便特殊场景获得最优集合通信性能。

分布式 Pallas 内核

TorchTPU 支持使用 Pallas 内核直接通过 ICI 硬件 RDMA 在芯片间通信，无需主机 CPU 参与。这样可以编写 Ring Attention、自定义 All-to-All/Shuffle、流水线移位和 Halo 交换等点对点通信内核，并作为原生 PyTorch 函数调用。

下面的示例把本地 buffer 异步发送到下一个 rank。

import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl
from jax.experimental.pallas import tpu as pltpu
from jax.sharding import PartitionSpec as P
import torch

# 1. Pallas 内核：异步发送到 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

    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. 用 ShardMap 和 DMA 信号量包装

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,
                ),
            ),
        )(x)

    return jax.shard_map(
        local_pallas_call,
        mesh=mesh,
        in_specs=P("x"),
        out_specs=P("x"),
        check_vma=False,
    )

# 3. 创建 TorchTPU 可调用包装

def get_send_to_next_op(world_size: int):
    mesh = torch.tpu.pallas.get_pallas_mesh(
        axis_names=("x",),
        mesh_shape=(world_size,),
    )
    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 内核的关键步骤如下：

阶段关键步骤与操作
1. Pallas 内核主体查询逻辑 rank 与目标 ID，调用 pltpu.make_async_remote_copy(...)，用 desc.start()/desc.wait() 控制异步调用生命周期
2. 分片映射与信号量在 scratch_shapes 中分配 DMA 信号量，用 JAX Mesh 和 PartitionSpec 配合 jax.shard_map 包装
3. TorchTPU 注册使用 torch.tpu.pallas.jax_op 包装，提供 mesh 和 input_partition_specs，在 PyTorch 张量上直接调用，兼容 eager 与 torch.compile

此前，使用 lax.axis_index 和 async_remote_dma/make_async_remote_copy 的 Pallas 内核按原始硬件设备 ID 路由通信，而不是按 PyTorch rank 路由。硬件 ID 与 rank 顺序不一致时，环通信和跨设备 DMA 可能被发送到错误物理设备；隔离子组也可能因分区数量不匹配而失败。

TorchTPU 现在强制逻辑寻址：

  • Rank 对齐映射： Pallas 内核直接映射到 PyTorch rank，而不是原始设备 ID；get_pallas_mesh 将 Pallas 分区和轴映射到 PyTorch Process Group 的 rank 顺序。
  • 逻辑设备寻址： 使用 DeviceIdType.LOGICAL 和规范设备分配 da(0, p) = p，保证逻辑 rank 映射到正确的物理目标。

TorchTPU 还提供获取 rank 对齐 JAX Mesh 的机制，让内核继承 PyTorch ProcessGroup 拓扑，在全 world 和隔离子组配置下执行，无需手动修补状态或打断计算图。

开发模型

TorchTPU 以开放方式开发，是 Google 与社区协作的成果，强调透明度和参与度。项目面向研究机构和行业用户，并通过与 Hugging Face 合作，在 Transformers v5、Diffusers 和 Kernels 中提供集成，让大量开源模型通过 to("tpu") 原生运行，并使用硬件加速内核。

项目文档位于 torchtpu.readthedocs.io，包括：

  • 教程： 虚拟机配置和模型移植的引导式学习路径；
  • 配方： 性能优化和内核实现等面向问题的方案；
  • 参考： 技术规范和 API 映射；
  • 说明： 架构与硬件基础的深入解释。

文档使用可直接运行的 notebook 构建。项目使用 Bazel 管理构建和测试，测试框架仍在演进并计划采用 Pytest。CI/CD 基础设施正在接入官方 PyTorch CI HUD，以保持可靠性并促进社区贡献。

TorchTPU 可通过 Pallas 和 Helion 进行自定义内核开发，并通过 vLLM、SGLang、TorchTitan 为分布式训练和推理提供一等支持。模块化设计让用户保留 PyTorch 惯用法，同时针对 TPU 专用硬件架构优化。

案例研究

案例研究 1：Hugging Face 模型覆盖

TorchTPU 通过 PyTorch 生态的开箱兼容性，让开发者无需重写模型代码即可在 Google Cloud TPU 上运行开源模型。项目构建了覆盖 NLP、视觉、音频、多模态和扩散模型的自动化测试框架，以 CPU 和 GPU 基线验证覆盖范围、稳定性和数值等价性。

从 v5e 开始的各代 TPU 都参与了 eager 与编译模式评估。测试语料约 5,000 个模型，包含数百个架构。模型选择采用两层筛选：第一层根据社区采用度、下载速度和趋势评分选择 Qwen、Gemma、Mistral、DeepSeek、GLM、Kimi、Llama、Whisper、Stable Diffusion、FLUX 和 Vision Transformer 等热门模型；第二层从第一层中精选，确保 500 多个模型系列至少有一个代表。

第一阶段是冒烟测试和 lowering 验证：用合成权重实例化模型，在 eager 模式验证前向执行，并验证图追踪、FX 到 StableHLO 图生成和 XLA 编译不崩溃。第二阶段使用领域指标验证数值正确性：

  • 因果 LLM 使用困惑度（Perplexity）评估，TPU 与基线的相对差异不超过 0.10 即通过；
  • 视觉、音频和编码器使用双阈值相对 MSE：Relative MSE = MSE(Y_TPU, Y_Baseline) / (Var(Y_Baseline) + 10^-7)，要求相对 MSE ≤ 0.10 或绝对 MSE < 0.05，以避免基线接近常数时指标失真。

截至原文撰写时的结果：

指标类别范围 / 定义结果
已验证模型总数端到端测试的不同开源模型4,917 个模型
独特架构系列Transformer、Diffuser 等架构类别500+ 个架构系列
第一层（流行度）通过率下载量最高的 LLM、视觉、音频和多模态模型，无崩溃98%，4,821/4,917
第二层（架构）数值一致性不同架构的代表模型，相对 GPU 基线97%，638/657

对于有界动态性（SymInt），项目在 839 个来自 5K 语料库的因果 LLM 上评估，834 个（99%）达到数值正确。

案例研究 2：开箱即用的性能

TorchTPU 的目标是成为 PyTorch 的原生 TPU 后端，包括在纯 PyTorch、设备无关代码上提供良好性能。

第一个示例是 TorchTPU 仓库中的 Muse Glimmer 30B，完全使用标准 PyTorch 模块和算子编写，包括 nn.Linear、nn.RMSNorm、F.scaled_dot_product_attention 和 torch.distributed 集合通信，没有使用模型专用自定义内核：

  • 单个 TPUv7 芯片推理（prefill）： 序列长度 5,120 时约 25.7k TPS、59% MFU；序列长度 10,240 时约 22.4k TPS、51% MFU；
  • 4 个 TPUv7 芯片训练： 使用 FSDP，序列长度 5,120 时每 TPU 约 5.6k TPS、38% MFU；序列长度 10,240 时每 TPU 约 4.6k TPS、32% MFU。

该示例使用 FSDP，world size 为 2，因为单个 TPUv7 芯片向 PyTorch 暴露两个设备，对应两个 TensorCore。训练示例启用了激活检查点，只在 HBM 中保存部分前向激活，在反向需要时重新计算其余中间激活。

同一份代码只需指定加速器（例如 device="tpu" 而不是 device="cuda"）即可运行在 TPU 和 GPU 上，并在 torch.compile、FSDP 等标准 PyTorch 优化和分片技术下获得性能。

第二个示例来自 TorchTitan 的 Muse Glimmer 30B GPU/CUDA 实现。移植到 TorchTPU 时只把尚未支持的 FlexAttention 换成 F.scaled_dot_product_attention，并应用 torch.compile 和 SimpleFSDP；在 4 个 TPUv7 芯片上训练达到：

  • 序列长度 5,120：每 TPU 约 5.3k TPS、36% MFU；
  • 序列长度 10,240：每 TPU 约 4.2k TPS、29% MFU。

这些结果显示，纯 PyTorch 模型，甚至最初为 GPU 编写的实现，只需少量改动和标准 PyTorch 优化即可在 TPU 上运行。对于追求峰值性能的生产工作负载，官方仍建议使用 Pallas 内核以及 vLLM-TorchTPU、SGLang-TorchTPU 或 TorchTitan-TorchTPU 等生产级技术栈。

案例研究 3：将训练扩展到 Pod 规模

项目使用 TorchTitan 验证 TorchTPU 可扩展到一个 TPU v7 Pod，模型选择为 Qwen3 MoE 235B。作为初步冒烟测试，在 64 个 TPU v7 上，TorchTitan + TorchTPU 得到 1,828 TPS，JAX + MaxText 得到 1,937 TPS。

测试设置如下：

  • 模型：qwen3_moe_235b_scan；
  • 本地批大小（每 TensorCore）：16；序列长度：4,096；
  • 单 Slice 使用 FSDP（SimpleFSDP），多 Slice 使用 HSDP；
  • BF16 精度，FP32 混合精度归约（梯度归约、RoPE 和优化器状态使用 F32）；
  • Pallas 内核包括 Tokamax Splash Attention、融合交叉熵损失和 MoE 分组矩阵乘法；
  • 94 个解码器层在堆叠权重上作为一个扫描循环，配合激活卸载、集合通信流水线和完整激活检查点；
  • 使用 TorchTitan Graph Trainer 联合执行前向和反向；
  • 使用辅助损失缓解专家不均衡；
  • 训练数据为 Hugging Face C4 en，并使用自定义数据集分片。
拓扑TPU v7 芯片数TensorCore 数总批大小每 TensorCore TPSMFU步进时间
1 个 Slice，FSDP1,0242,04832,7682,23628.8%29.3 s
2 个 Slice，HSDP2,0484,09665,5362,21428.5%29.6 s
4 个 Slice，HSDP4,0968,192131,0722,09526.9%31.3 s
8 个 Slice，HSDP8,19216,384262,1442,10427.1%31.2 s

这些数据体现了弱扩展：整体计算量线性增加，每秒 token 数线性增长，MFU 基本保持平稳，损失曲线在不同并行级别下几乎完全相同。

多 Slice HSDP 的前向传播不涉及跨 Slice 通信，每个 Slice 处理自己的批次。反向传播时，一个 Slice 先计算第 t 层梯度，再通过 DCN 上的 all-reduce（使用 Megascale-XLA）跨 Slice 求和，然后继续计算第 t-1 层。借助 XLA，第 t+1 层的 all-reduce 可以与第 t 层计算重叠，这是每 TensorCore TPS 从一个 Slice 扩展到八个 Slice 仍然相近的主要原因。

案例研究 4：在 vLLM 上使用 Qwen3.5 进行 LLM 推理（SemiAnalysis InferenceX）

LLM 推理，尤其是大规模 MoE 架构，是 AI 云基础设施中要求最高、成本最关键的工作负载之一。在 SemiAnalysis 的独立基准 TPU Inference Externalization Full Steam Ahead — InferenceX 中，Google TPU v7 搭配 vLLM-TorchTPU 软件栈交付了具有竞争力的每美元性能。

Google 还提交了 Qwen 3.5 在 Ironwood TPU 上的首批 AgentX 结果。AgentX 覆盖多轮、长上下文、高前缀复用和子智能体突发四类智能体工作负载；在测试中，4 块芯片的 TPU v7 实现了完整的帕累托前沿领先。主要数字为：

  • InferenceX 8K/1K：每百万 token 总成本 2.01 美分，低于 9 月首次发布时的 3.6 美分；
  • AgentX：每百万 token 总成本 0.83 美分，并具备领先的交互响应性。

这些数字依赖特定硬件、服务配置、模型、输入输出长度和计价口径，不能当作所有工作负载的固定价格。取得结果需要 Pallas 内核和运行时调优，原文列举了几个软硬件对齐案例：

优化瓶颈实现
将专家并行 MoE 层作为单个 SPMD 程序DP8 + EP 下，vLLM MoE 层的 token 收集、路由、置换、两次分组矩阵乘法、合并和 reduce-scatter 中，prefill 密集步骤的 reduce-scatter 约耗时 1,190 μs，占设备时间 22%，期间没有其他操作并行8 个 TensorCore 运行同一个 SPMD 程序；使用 SparseCore 异步 all-gather、Pallas top-k、分组矩阵乘法和 pltpu.make_async_remote_copy，把结果推送到拥有 token 的 TensorCore，并与后续专家矩阵乘法重叠
Sequence-on-Lane 注意力布局（无填充 FP8 KV 缓存）TPU 向量单元以 32 位字工作，标准 KV 布局在 8 路张量并行下每个字有一半填充，每 token 每核心缓存为 1,024 字节但数据只有 512 字节将页面 token 沿 128 个 lane 排列、头维度沿子 lane 排列，让一个字包含同一 token 的四个值，KV 内存和每次解码读取的字节数减半
Gated DeltaNet 优化60 层中有 45 层使用 Gated DeltaNet，每头维护 128×128 状态；逐 token 解码需读写每个序列状态，prefill 递归使矩阵单元空闲Pallas 内核融合卷积和递归，批处理解码序列并以 64 token 分块处理 prefill；状态和 KV 缓存共享分页池，直接流式读取状态并原地更新

新版 vLLM-TPU 后端让工程团队能够使用标准 PyTorch 惯用法，在公开可用 TPU 基础设施上部署开放权重模型，同时获得针对 TPU 的服务效率。

案例研究 5：使用 TorchTitan 进行 LLM 微调

TorchTPU 与 TorchTitan（PyTorch 面向大规模分布式训练的参考平台）原生集成，为 MuseGlimmer 30B、Llama 3 8B–405B、Qwen 3 235B MoE 和 FLUX.1-dev 12B 提供统一训练环境。评估覆盖 TPU v6e（Trillium）和 TPU v7x（Ironwood），多 Slice 通过 torchrun 在 Google Kubernetes Engine（GKE）上编排。TorchTitan 源码位于 google-pytorch/torchtitan。

高效多加速器训练需要让模型并行策略与 TPU 物理拓扑对齐。项目采用 2D+ DeviceMesh，把通信密集的张量并行和专家并行映射到连续、高带宽的 ICI，把复制数据并行映射到跨 Slice 的 DCN。SimpleFSDP 提供符合 SPMD 的操作；All-Gather 和 Reduce-Scatter 可以卸载到片上 SparseCore，与 TensorCore MXU 计算重叠。HLO diffing 和 @spmd_safe 用于保证确定性并避免同步死锁。

FLUX 模型的逐层优化结果如下：

阶段步时相对 eager 基线
开箱即用 eager11.49 s（0.26 samples/s）1.00×
DeferAndFuse eager9.95 s1.15×
SimpleFSDP6.52 s1.76×
Pallas SplashAttention5.55 s2.07×
torch.compile(backend="tpu")3.02 s3.80×
目标模型 peephole 重写后2.78 s（1.08 samples/s）4.13×

在多主机集群上，FLUX.1-dev 使用 SimpleFSDP 从单个 v6e-8 节点（8 芯片）扩展到 v6e-256 Slice（64 节点、256 芯片），每样本步时只从 0.92 s 增至 1.00 s，总步时为 3.01 s。TPU v7x 上，Llama 3 405B 在 1,024 芯片、BF16、8K 上下文达到 500 tokens/sec/chip 和 50% 以上 MFU；Qwen3 235B MoE 在 8,192 个 TPU v7x 芯片、BF16、4K 序列长度下达到约 28% MFU。

关键架构优化包括：

  • SimpleFSDP： FSDP 的编译器友好版本。FSDP2 通过模块钩子注入 all-gather 和 reduce-scatter，容易使 torch.compile 退化为逐算子调度；SimpleFSDP 把通信表示为追踪图中的普通操作，让 XLA 可见并进行融合、重叠和 SparseCore 卸载。
  • Scan 高阶算子： torch._higher_order_ops.scan 将重复 Transformer block 作为堆叠权重上的一个循环，保持编译时间和 gather 参数内存近似恒定，并让每层集合通信处于循环内以便重叠。
  • Pallas SplashAttention： 支持完整 GQA、bias 和 dropout，替换不适合 TPU 向量 tile 的默认注意力实现，在扩散注意力 block 上带来 17.5% 加速。
  • 融合 Pallas Cross-Entropy Loss： 将 Tokamax 在线归约内核移植到原地计算损失，避免大词表 [B×S×V] logits 物化导致 OOM。
  • Graph Trainer： 将完整训练步骤（前向、loss、反向、优化器更新）捕获为单图，编译成一个 XLA 程序，让编译器进行全局融合、通信调度和内存规划，并决定激活保存、重算、卸载及优化器 buffer donation。
  • Activation Offloading： 前向结束后将反向所需激活移到主机内存，使用前再复制回 HBM；在 scan 循环中逐层移动激活，并通过异步传输与计算重叠。这对 Llama3-405B 和 Qwen3-235B 等大模型达到目标序列长度至关重要。

TorchTitan on TorchTPU 的后续重点包括扩大 TPU 优化模型列表、通过 torchao 集成 FP8/FP4 低比特量化，以及通过 VeRL 将训练循环与 vLLM/SGLang 采样后端结合，提供端到端强化学习（RL）。

案例研究 6：使用 TorchRec 和 SparseCore 加速推荐模型

推荐和排序架构（例如 Meta 的 Deep Learning Recommendation Model，DLRM）由于高基数 embedding 表而面临内存和计算瓶颈。在 GPU 集群上，这类表通常需要 DistributedModelParallel（DMP）软件分片和独立优化器更新，会引入延迟。

Google 与 Meta PyTorch 团队将原生 TPU 后端上游到 TorchRec，使推荐模型可以绕过主机计算瓶颈，直接使用 TPU SparseCore 硬件。

  • 架构集成： 多线程 C++ PyBind 后端在主机 CPU 上并行执行分区映射和索引填充，生成直接映射 SparseCore 的结构化 batch。TPU 内部实现 In-Place Optimizer Fusion：反向传播直接进入 SparseDenseMatmulGrad，不在 HBM 中物化 .grad 张量，SGD、AdaGrad 或 Adam 更新都可以原地执行。SparseDenseMatmulOp（前向查找）和 RaggedAllToAll（ICI 跨芯片集合通信）等 ATen 算子直接运行硬件加速内核。
  • 性能结果： 在 32 芯片 TPU v7x（Ghostfish）集群的行业标准 DLRM v2 上，达到 Google 高度优化 JAX 基线的 100% 性能，step latency 同为 18.3 ms。这个结果使用的是基础 SparseCore 内核，尚未应用更高级优化；收敛和数学结果也与 JAX 实现一致。

后续工作将通过 TensorCore/SparseCore 流水线交错 MLP 与 embedding 查找，使用表堆叠统一内存 buffer 并降低内核开销，同时将 SparseMapRow、SparseGather 等更底层 SparseCore 查找操作以 ATen 内核暴露给用户。

Future Work

TorchTPU 已在 GitHub 公开，当前开源版本只是开发起点。未来 6–12 个月的路线图聚焦生产加固、生态扩展和社区反馈，重点包括四个方向。

  1. 基础能力加固与提升用户效率： 完善 True Eager 体验，为 ATen 算子构建预编译动态内核库；让 XLA 有界动态性达到生产可用，使 torch.compile 能处理 SymInt 而无需额外重编译；简化 GKE 打包，扩大算子覆盖并推进 FP8/FP4 量化，面向后续 TPUv8 架构，并通过嵌套编译区域降低编译时间。
  2. Pod 规模分布式训练与高级训练： 在 Qwen3 Pod 规模验证基础上扩展更多模型架构和分片技术；通过 Ray Core 编排的 vLLM/SGLang 采样后端实现端到端 RL；验证基于共识的 compilation handshake，在不依赖脆弱 @spmd_safe 注解的情况下安全优化集合通信与计算的重叠。
  3. 新兴架构、推理与服务： 支持推理模型、多模态和生成式扩散模型（如 Wan 2.2、Flux Klein）以及新型 MoE；优化 vLLM 和 SGLang 低延迟生产服务，在未来 TPU 代际加入原生 FP8/FP4 量化，确保新架构首日得到 TPU 支持。
  4. 高效通信与诊断： 优化 TPU Slice 内外集合操作，加入零拷贝 DMA 以支持硬件辅助内存卸载；将 Trace Viewer 迁移到 WebGPU Trace Viewer 2.0，并部署持续 profiling（XProfz flight recorder），定位生产环境中的间歇性瓶颈。

TorchTPU 团队表示，项目相信透明、协作的工程方式才能形成健康生态。随着 TorchTPU、vLLM-TorchTPU、SGLang-TorchTPU 和 TorchTitan 四个基础仓库在 2026 年 10 月 PyTorch Conference 开源，团队希望开发者、研究人员和合作伙伴通过 GitHub 或 TPU Builders 计划反馈模型架构优先级和工作负载顺序。

TorchTPU 团队

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，以及许多其他贡献者。

与 Gemini 专题的关系

TorchTPU 是 Google TPU 的 PyTorch 后端，并不是 Gemini API、Gemini 模型权重或 Gemini 服务端实现。它可能让使用 Gemma 等开放权重模型、并具备 TPU 集群的团队更方便地迁移 PyTorch 训练和推理工作流；Gemini 模型服务、API 权限和 Google Cloud TPU 资源仍是不同层面的产品与访问问题。

Gemini API 的模型、区域和配额应以 Gemini API 官方文档和 Vertex AI 文档为准。对站内 Gemini 开发入口，可参考 Gemini AI Studio 与开发者入口指南。

来源

性能数字、硬件代际、软件版本和价格会随基准配置变化。部署前请以 PyTorch、TorchTPU、Google Cloud TPU、vLLM 和相关模型仓库的最新文档为准，并遵守所在地法律法规及服务条款。

相关主题#TorchTPU#Google TPU#PyTorch#Google AI#Gemini 开发
返回 Gemini 专题