Introducing TorchTPU

介绍 TorchTPU - PyTorch

特色项目

TL;DR

TorchTPU 是 Google 与 Meta 合作推出的开源项目，为 Google TPUs 提供了原生且原理性的 PyTorch 后端。通过即时执行优化、torch.compile 集成和原生硬件加速，它实现了开箱即用的模型执行和 TPU 迁移，同时提供行业领先的推理和训练性能。

介绍 TorchTPU

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

为了在 Python 框架之下编程 Google 的 TPU，开发者使用 XLA 编译器堆栈。2020 年，Google 推出了基于 Lazy Tensors 的首个 PyTorch TPU 编译版本 PyTorch/XLA。虽然这是基础性的工作，但最初的集成难以跟上 PyTorch 快速的社区驱动发展。关键问题在于：PyTorch/XLA 需要对源代码进行侵入性修改以显式定义可编译区域，存在严重的编译时延迟，不支持 PyTorch 原生的即时（"eager"）执行模式，并在支持 PyTorch 生态系统中的新兴分布式计算原语方面存在滞后。与之前的 PyTorch/XLA 和 TorchAX 项目（将 PyTorch 强行适配到 JAX 的工作方式）不同，TorchTPU 被设计为 PyTorch 的原理性 TPU 后端：即时执行是首要优先级，现有 PyTorch 工作负载只需简单的代码修改（例如 .to("cuda") → .to("tpu")）即可开箱即用。

更具体地说，我们在启动 TorchTPU 时明确制定了以下原则：

• PyTorch 公民身份：运行时无 JAX 依赖；稳固的即时模式；基于 Torch Dynamo 的 Torch.compile（非 Lazy Tensor）且性能与 JAX 编译相当；通过在 PyTorch 堆栈最低层提供 TPU 集合操作（如 Process Group API），广泛支持 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 eager mode）；
• 新增对约 450 个 Aten 操作（跳过具有结构或复合代理的操作）的支持，包括所有主要深度学习架构所需核心 ATen 操作的前向和反向传播；
• 新增通过 JAX/Pallas 和 Helion DSL 编写自定义内核的能力；
• 验证了对下载量最高的 5,000 个 Hugging Face 模型的开箱即用执行，包括：LLMs：Qwen、DeepSeek、Gemma、LLama、Mistral、GPT、OPT、BLOOM、Falcon、StableML；Audio/Speech：Whisper、Wav2Vec2、HuBERT、MMS；Vision：DinoV2、SigLIP、CLIP、Swin Transformer、ConvNeXt、ViT、ResNet、DETR；Text Encoders & Search：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 芯片包含（视代际版本而定）：

• 一个或多个 TensorCores，每个包含一个或多个矩阵乘法单元（MXU）、向量单元（VPU）和标量单元（参见此链接）。根据 TPU 代际，MXU 可包含 256×256 的乘加器，以 systolic array（ systolic 数组）形式排列；
• 一个或多个 SparseCores，这是专为稀疏操作（如 gather、scatter、embeddings）和 collectives 设计的单元。它们可以与 TensorCore 并行运行，实现计算-计算和计算-通信重叠；
• 大容量高带宽内存（HBM）——在 TPU 8i 中为 288 GB，带宽达 8.6 TB/s，可用于存储 ML 数据并与 CPU 共享；
• 双向芯片间互连（ICI）——在 TPU 8t 和 8i 中为 19.2 TBps，允许用户配置的 TPU Slices（TPU Pod 的逻辑子组，由预配置数量的 TPU 芯片以 3D torus（TPU 7、TPUv 8t）或 Boardfly（TPU 8i）配置连接）内的 TPU 芯片快速通信；
• 双向数据中心网络（DCN）——在 TPU 8t 中为 400 GBps，允许跨 Slices 的 TPU 芯片通信。

TPU 通过两个入口点使用 XLA 编译器堆栈进行编程：StableHLO（机器学习操作的高级 IR）和 Mosaic（自定义内核的低级 IR）。实际的 TPU 指令集架构（ISA）是专有的，可能更改且不公开。实际上，可以将 StableHLO（更不用说 Mosaic）视为等同于 PTX。下图是高层次的示意图。

C++ 程序可通过实例化 PjRT Client 访问本地 TPU，其 API 提供将 MLIR（StableHLO 或 Mosaic）编译为 TPU 可执行文件、启动该可执行文件以及在 TPU 与 CPU 之间复制数据的方法。请注意，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操作，而不是使用NCCL等独立库。

易用性和自动加载机制

我们设计TorchTPU时充分考虑了易用性。当用户在已安装TorchTPU wheel的TPU机器上导入PyTorch时，通过PyTorch的自动加载机制，仅需进行微小更改即可在TPU上运行PyTorch代码，并且只需几行代码即可完成：

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

GPU之前

TPU之后

替代方案（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()
    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 高性能计算的核心。对于非收缩操作（如逐元素加法、激活函数、层归一化和 softmax），TorchTPU 会通过 TPU 向量执行单元（VPU）执行 FP32 张量运算，该单元原生支持完整的 IEEE-754 FP32 算术运算。

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

• medium：使用较低精度进行乘法运算（将输入转换为 MXU 原生 BF16），同时使用 FP32 进行累积以保持稳定性。TorchTPU 默认使用此模式以充分利用 TPU 硬件加速，因为单次传递的 BF16 乘法配合 FP32 累积对几乎所有深度学习工作负载都具有足够的精度。

• 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

更高/最高精度机制

专用硬件 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 上下文管理器

立即执行、StableHLO 降级和 JIT 编译

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

TorchTPU 的即时执行模式的一个独特特性是它能够以多种方式将 Aten 操作分发到 TPU：

• 立即执行，实现逐个操作的分发；
• 延迟到自动确定的 Materialization Point（例如 .to("cpu") 语句），并与其他所有延迟到该点的操作一起聚合/编译；
• 介于两者之间的任何方式。

虽然这种灵活性确实增加了设计的复杂性，但我们从一开始就刻意添加了它，因为我们希望尝试不同的策略以提高即时执行模式的性能。此外，这种灵活性在处理视图语义和 torch.compile 时也非常有用，如本文后面将讨论的那样。

总结来说，TorchTPU 提供了三种即时执行模式，可以通过环境变量或上下文管理器进行设置：

• Strict（默认）：我们逐个分发 Aten 操作，它们在 TPU 上异步于 CPU 执行。这与在 GPU 上使用 PyTorch 的即时执行模式效果相同。
• Debug（`export TORCH_TPU_LAUNCH_BLOCKING=1`）：我们逐个分发 Aten 操作，但 TPU 执行与 CPU 同步。这种方式速度较慢，主要用于调试，等同于在 GPU 上使用 CUDA_LAUNCH_BLOCKING，并且额外启用了 TorchTPU 的边界检查。
• DeferAndFuse（`export TORCH_TPU_DEFER_AND_FUSE=1`）：我们将多个 Aten 操作延迟并融合为更大的块。这是 TorchTPU 引入的一种新的即时/编译混合模式，相比 Strict 模式可以提升 50% 或更高的性能，同时减少峰值 HBM 内存使用（操作融合减少了 HBM 中的临时存储）。

现在我们已经介绍了 TorchTPU 即时模式的基础知识，接下来我们简要说明当从 PyTorch 分发器接收到回调调用时会发生什么。假设我们使用的是 DeferAndFuse 模式，并以 at::mul.out 为例：

• 创建一个 Deferred Op 实例，其中包含：
• 一个 StableHLO 构建器对象，当被调用时会返回一个实现 aten::mul 语义的 StableHLO 函数（如下所示，注意任何隐式的广播和 dtype 提升会在其他地方处理）：

    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 的训练和动量参数），用于计算缓存键（如下文所述）；
• 返回一个新的 at::Tensor，其 c10::StorageImpl 指向该 Deferred Op 实例（通过间接引用）。

当从 PyTorch 分发器接收到与 mytensor.to("cpu") 或 mytensor.item() 等类似的回调时，TorchTPU 会：

• 分析当前累积的、导致 mytensor 计算的 延迟操作区域 R；
• 应用活跃性分析，确定在 R 中需要实例化的操作（包括生成 mytensor 的操作）；重要的是，具有副作用的操作（如集合操作）始终会被保留（包括其调用顺序），而不在 R 之外活跃的操作可以跨不同工作负载副本确定性地检测到，因为 Python 的内存管理使用引用计数。

• 为要实例化的操作图 G 创建一个编译缓存键，以便未来能够快速且唯一地识别该图；

• 使用计算出的缓存键检查编译缓存，判断图 G 是否已被编译。如果已存在，直接在 TPU 上调度图 G 的可执行文件；如果不存在，则使用线程池异步编译图 G，随后在 TPU 上调度其执行。

• 等待 TPU 上 mytensor 的值就绪，然后将其从 TPU HBM 复制到与目标 at::Tensor 关联的 CPU 存储中。

从表面来看，上述方法似乎效率不高，因为 CPU 只有在遇到 .to("cpu") 语句时才会启动 TPU 执行。但对于实际工作负载，我们预计在初始几次迭代/步骤（用于填充编译缓存）之后，CPU 执行将仅涉及 PyTorch 跟踪和操作调度，因此会始终领先于 TPU 执行（后者将忙于执行耗时的数值计算）。尽管如此，为了加速重复执行的操作序列（例如变压器块堆栈）的实例化，我们引入了“重复操作启发式方法”，该方法可检测此类重复并在达到 TPU 到 CPU 数据传输之前触发实例化点。

编译缓存

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

• 第1层：进程内内存缓存，用于在同一进程内共享编译结果。这是三个层级中最快的。

• 第2层：机器本地共享内存，支持同一主机上的跨进程共享。这对于多设备训练/推理非常有用。

• 第3层：由远程文件系统（如 Google Cloud Storage）支持的持久化缓存，允许跨主机、程序和作业重启重复使用。这在 Colab 和数据中心等执行环境中运行 TPU 程序时非常有用。

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

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

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

• 一个无形状键，用于捕获计算图的结构，但不包含张量输入/输出的尺寸，

• 一个维度键，用于捕获张量尺寸，

• 一个编译选项键，用于捕获传递给 XLA 编译器的选项以及 TorchTPU、libtpu 和 PjRt 运行时的版本。

无形状键和维度键共同也被称为图键。理论上不同图、编译器选项和组件版本之间可能出现缓存键冲突，但考虑到 192 位键空间的规模，实际使用中此类冲突不会发生。

视图语义

PyTorch 的操作语义与 StableHLO 语义在可变性和共享内存视图的处理上存在显著差异：

• PyTorch 假设每个 Tensor 是指向线性字节数组的指针，Tensor 层级的步长元数据用于描述逻辑索引元组到该数组线性偏移量的转换。多个 Tensor 可以共享相同的底层字节数组，从而允许从任意 Tensor 执行原地写入，并使其他 Tensor 可见该变更。借助附加的步长元数据，视图操作可以在 CPU 的调度阶段通过简单更新元数据以低成本执行。尽管大多数视图操作属于少量常见操作，但也可以通过 as_strided 操作符定义任意复杂的模式。

• 相比之下，StableHLO 是一种纯函数式语言，张量值在构造后不可变。物理布局对用户抽象，反映编译器的分块决策，因此不支持直接步长操作。此外，StableHLO 不支持原地写入，而是通过 dynamic_update_slice 操作实现选择性更新，将新值合并到现有张量的子切片中。

在 TPU 上实现存在写入操作时的正确视图语义并不直接；然而，TorchTPU 通过视图分解、视图反转和延迟操作等机制实现了这一目标。这种语义桥梁使 TorchTPU 能原生支持否则无法实现的 PyTorch 模式，例如分布式训练中使用的拼接参数数组。

流与 SparseCore

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

• 计算操作；

• 集合操作；

• 主机与设备之间的数据传输。

在 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 内部的 MXU 所擅长的领域。为了赋予用户对工作执行位置的完全控制，TorchTPU 流可以手动固定到 SparseCore，这种流在性能和操作支持完整性方面将与默认运行在 TensorCore 上的流有所不同。流还可以像 GPU 上一样用于协调异步数据传输（通过 DMA 引擎）。

编译模式

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

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

Dynamo 通过传递 torch.fx.GraphModule 加上示例输入张量列表来调用我们的后端。此时我们以急切模式追踪提供的图，但所有操作都会被延迟处理且不会在 TPU 上实际执行。这使我们能够为整个 FX 图生成并编译一个 StableHLO 程序。最后，我们的 Dynamo 后端返回一个封装生成 TPU 可执行文件调用的 Python 可调用对象。需要注意的是，由于有界动态性（下文将讨论）需要支持 SymInts，因此作为临时限制，当在 TPU 机器上调用 torch.compile 时，我们会自动设置 dynamic=False，但该功能尚未默认启用。

有界动态性

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

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

内核

我们区分两种类型的内核：

• PyTorch 框架提供的内核，例如 torch.nn.functional.scaled_dot_product_attention。

• 用户提供的自定义内核，用于加速模型的特定部分；这些内核通常与具体模型相关。

对于 1)，我们最初实现了一种方法，即在 Pallas 中编写 TPU 内核，然后通过符号形状导出到 Mosaic，并通过 TorchTPU 步骤在调用时专门化内核的形状。该机制的灵活性低于预期，因此我们正在尝试一种新方法来编写 Mosaic 内核以支持动态形状。

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

• （生产就绪）用户使用 Pallas 编写 TPU 内核，并通过 TorchTPU API 与 PyTorch 集成，如下所示：

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

# 将 Pallas 内核包装成 JAX 函数
@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)

# 像标准 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，这是一个高性能 TPU 内核库，现在可与 PyTorch 实现无缝互操作。Tokamax 提供了关键操作的高性能实现，包括专家混合构建模块（如分组矩阵乘法和 token 聚合/分散）、注意力变体（如分页注意力、门控 deltanet、Kimi delta 注意力和多头潜在注意力）以及 TopK 计算。操作符组合将随着时间推移不断演进，以支持最新 OSS 模型在最新硬件加速器（如发布时的 TPU v8i）上的运行。这些内核由强大的基础架构支持，包括自动调优（如瓷砖大小搜索）、低开销基于 XProf 的基准测试和回归测试，以降低操作复杂性。虽然内核本身使用 JAX/Pallas 编写，但库的基础设施采用 torch.tpu.pallas.jax_op API 透明管理互操作性，并为开发者提供原生 PyTorch API 以提高便利性。

• （实验性）用户使用 Helion 重写其 GPU 内核，这使得可以通过相同的 PyTorch 基于 DSL 针对多种加速器，如下所示：

# 直接在 Helion 中定义内核
@helion.kernel(backend="pallas")
def add_vectors(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    # 自动分区的瓷砖设备执行循环
    for tile in hl.tile(out.size()):
        out[tile] = x[tile] + y[tile]
    return out

# 像标准 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)

分布式支持

为确保与现有 PyTorch 分布式 API 和框架的兼容性，我们决定在进程组 API 层级实现 TPU 集合操作，通过继承自 c10d::Backend；这使我们能够在即时模式和编译模式下拦截集合操作，并无缝支持 PyTorch 分布式框架（如 DTensor、FSDPv2、SimpleFSDP、FairScale 等）。与 GPU 上的 PyTorch 类似，我们也遵循（甚至强制要求）每个设备对应一个进程的方案。从实现角度来看，TPU 堆栈并未提供专门用于 TPU 间通信的独立库，而是将 TPU 集合操作表达为 StableHLO 操作（如 stablehlo.all_gather、stablehlo.all_scatter、stablehlo.all_to_all 等），这与上述机制自然契合。

在进程组层级拦截集合操作已使我们接近完成，仅需处理几个与 GPU 相比的 TPU 特殊之处：

• 不同的初始化顺序。考虑一个连接了多个设备的主机。在 GPU 场景中，通过 torchrun 启动的进程可以看见所有连接的设备，即使它只会使用其中一个（例如 device=f"cuda:{local_rank}"），并在调用 torch.distributed.init_process_group() 之前将模型移动到该设备。但在 TPU 场景中，进程只能看到一个设备，且无法得知主机连接了多少设备。为了正确初始化 TPU 运行时，进程需要在调用 torch.distributed.init_process_group 之前明确以下信息：(a) 有分布式运行的意图，(b) TPU 数量（世界规模），(c) TPU 类型（因为它会影响拓扑结构）。

• 默认进程组知识。考虑两个子组（4, 0, 1, 3）和（5, 2, 6, 7）在组内执行 all_reduce 操作。注意 torch.distributed.all_reduce 的输入仅包含需要执行集合操作的子组。相比之下，stablehlo.all_reduce 需要所有子组的信息（即 PyTorch 的默认进程组），即使当前进程不属于这些子组，如以下 MLIR 代码片段所示：

  %0 = "stablehlo.all_reduce"(%arg0) <{replica_groups=dense<[[4,0,1,3], [5,2,6,7]]>: tensor<2x4xi64>}> ({...

• MPMD 支持。历史上，TPU 堆栈是基于 SPMD（单程序多数据）设计的。借助 Shardy 等工具，分布式 JAX 程序可以为单个主机编写，无需包含通信代码，然后通过用户提供的指令在大规模上进行分片，而 XLA 编译器会负责在适当的位置插入集体操作，同时最大化计算与通信的重叠。不幸的是，PyTorch 生态系统并未采用干净的 SPMD 方法（一个显著例外是 SimpleFSDP），不同进程可能执行不同的 ML 操作，导致集体操作的原因包括：用户插入的基于进程等级的条件语句；DTensor 在张量维度无法被设备数量整除时对共享张量的不均匀分割。当 TPU 的 StableHLO 模块包含集体操作且模块不一致时（这可能是由于生成的 TPU 代码中不同等级的同步标志 ID 或缓冲区偏移不匹配），可能导致不可预测的死锁或内存损坏。TorchTPU 的处理方式如下：（默认）瘦集体操作。默认情况下，TorchTPU 会将每个集体操作编译到其自己的“瘦”StableHLO 模块中，这确保了不同等级的缓冲区和同步标志 ID 的一致性。这种方法有效，但可能带来显著的性能损失（例如端到端性能下降 30-40%），因为无法让 XLA 编译器重叠计算和通信。当前正在进行的优化是将瘦集体操作运行在单独的流上，并固定到 SparseCore，使其能与 TensorCore 上的数学运算重叠。我们预计该优化可以恢复瘦集体操作的大部分性能损失。（可选）@spmd_safe 装饰器。如果用户知道某段代码是 SPMD（例如没有基于 torch.distributed.get_rank() 的条件语句），可以使用 TorchTPU 的 @spmd_safe 装饰器标记该代码段。在 torch.compile 下，这会消除 TorchTPU 否则会在集体操作周围插入的图断点（见上文瘦集体操作），从而生成一个包含嵌入式集体操作的大 StableHLO 模块，XLA 编译器可以通过移动集体操作来最大化计算与通信的重叠，从而高效优化；这已被证明在 SimpleFSDP 或手动放置的集体操作中能显著提升性能。该方法的一个挑战是 MPMD 代码可能难以察觉，用户可能在代码并非严格 SPMD 的情况下使用 @spmd_safe 装饰器，导致系统变得脆弱。（可选）编译握手（RFC）。当给定的 torch.compile 语句触发包含通信操作的模块编译时，该机制会启动。它会在所有副本之间发起一轮共识，验证编译后的 FX 图是否一致，从而允许副本在给定的 FX 图周围禁用图断点，类似于上文所述的 @spmd_safe 机制。握手阶段由环境变量 TORCH_TPU_SPMD_HANDSHAKE_STAGE 控制，可设置为：OFF – 从不执行握手（当前默认值）；COMPILE_STAGE – 在开始 FX 图编译前执行握手；DISPATCH_STAGE – 在开始编译的 FX 图执行前执行握手。此模式仅应在用户预期特定等级需要重新编译时使用，因为它比前一个选项有更高的开销。

分布式初始化

在 TorchTPU 上初始化分布式工作负载遵循标准 PyTorch 分布式语义，但需要注意 TorchTPU 需要两个特定于 TPU 的变量，而标准 PyTorch torchrun 无法设置这些变量：

• TORCH_TPU_TOPOLOGY : TPU 切片的物理网格几何形状（例如，2x2x1 切片配置的“2,2,1”）。

• TORCH_TPU_SLICEBUILDER_ADDRESSES : 用于构建 TPU 芯片间互连（ICI）网络网格的开放端口/地址列表（每个设备一个，例如“localhost:8471,localhost:8472,…”）。

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

A. 工作节点实现（train.py）

import os
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}] 在 {device} 上初始化")

  # 3. 在 TPU 工作节点之间执行分布式集合操作
  x = torch.ones(4, 4, device=device) * (rank + 1)
  dist.all_reduce(x, op=dist.ReduceOp.SUM)

  print(f"[Rank {rank}] 归约总和: {x[0, 0].item()}")

  # 4. 清理进程组资源
  dist.destroy_process_group()

if __name__ == "__main__":
  worker_main()

B. 使用 torchrun（CLI）启动工作负载

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

最后请注意，给定节点的默认进程组是该节点的切片。

最优设备网格映射

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

• ICI : 直接的光/铜互连，将芯片连接在环形网格中。提供数十至数百 GB/s 的带宽。

• DCN : 用于连接不同 TPU 切片的标准主机级以太网。与 ICI 相比带宽显著更低且延迟更高。

一个良好的规则是使用 TorchTPU 的 topology_aware_mesh() API，并将通信量最大的并行维度映射到最快的、最连续的物理 ICI 维度，即按网络强度从低到高排列网格维度。请参见下表中的示例：

| 并行类型 | 通信模式 | 频率/体积 | 最优物理位置 | |---------|---------|---------|-------------| | 张量并行（TP） | all_reduce, reduce_scatter | 每个 Transformer 层（前向+反向） | 最快的 ICI 环（同一托盘/主机，连续的1跳芯片） | | 上下文/序列并行（CP/SP） | all_to_all, Ring Attention | 高（每个注意力层） | 本地 ICI 轴（切片内） | | 专家并行（EP） | all_to_all | 每个 MoE 路由步骤 | 本地 ICI 轴 | | FSDP / 全分片 DP | all_gather, reduce_scatter | 每个模块/层1次 | | | 数据并行（DP） | all_reduce（仅梯度） | 每次训练步骤1次 | ICI 或 DCN（多切片） |

以下是代码示例：

from torch.distributed.device_mesh import DeviceMesh

# 单切片二维网格示例（DP + TP）

# 逻辑网格形状：在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,
)

# 实例化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，允许用户更直接地自定义网格，以在特定场景中实现最优的集体通信性能。

分布式Pallas内核

TorchTPU支持自定义Pallas内核，可通过硬件RDMA在ICI上直接进行芯片间通信，无需主机CPU参与。这使您能够编写高效的点对点通信内核（如环形注意力、自定义All-to-All/洗牌、流水线并行移位或Halo交换），并将其作为原生PyTorch函数调用。以下示例展示了具体实现方式：

# 1. Pallas内核主体：异步将本地缓冲区发送到相邻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
  # 触发TPU互联上的硬件远程DMA
  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信号量的JAX包装器
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,))
  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内核的步骤：

阶段 | 关键步骤与操作 --- | ---

1. Pallas内核主体 | 查询逻辑rank和目标ID<br>调用<br>pltpu.make_async_remote_copy(...)<br>控制异步调用生命周期：<br>desc.start()<br>/<br>desc.wait()
2. ShardMap与信号量 | 在<br>scratch_shapes中分配DMA信号量<br>使用<br>jax.shard_map<br>包装JAX Mesh与PartitionSpec
3. TorchTPU注册 | torch.tpu.pallas.jax_op<br>提供mesh和<br>input_partition_specs

直接在 PyTorch 张量上调用（Eager &

torch.compile

)

解决设备映射不一致性问题。在 TorchTPU 之前，Pallas 内核使用低级芯片间操作（如 lax.axis_index 和 async_remote_dma / make_async_remote_copy）时，通信是基于原始硬件设备 ID 而非 PyTorch 等级进行路由的，与 JAX 保持一致。当硬件设备 ID 与等级顺序不匹配时，环形通信和跨设备 DMA 会路由到意外的物理设备，且在孤立子组上执行时会因分区数量不匹配而失败。

TorchTPU 通过强制使用逻辑寻址解决了这些不一致性：

• 按等级对齐的映射：TorchTPU 现在将 Pallas 内核直接映射到 PyTorch 等级而非原始物理设备 ID，通过 get_pallas_mesh API 确保一致路由，该 API 将 Pallas 中的分区和轴直接映射到 PyTorch Process Group 的等级顺序。

• 逻辑设备寻址：TorchTPU 现在强制使用 DeviceIdType.LOGICAL 和规范设备分配（da(0, p) = p），保证每个逻辑等级始终映射到正确的物理目标。

统一的 ProcessGroup 集成。TorchTPU 现在提供机制获取与等级对齐的 JAX Mesh，使内核能够继承 PyTorch ProcessGroup 拓扑。这使得在完整世界和孤立子组配置中都能高效执行，无需手动状态修补或图中断。

开发模式

TorchTPU 作为 Google 与社区的协作项目开放开发，优先考虑透明度和参与度。开发模式鼓励来自更广泛 PyTorch 生态系统和行业合作伙伴的贡献。

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

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

• 教程：包含虚拟机配置和模型移植的指导学习路径。

• 实用方案：针对性能优化和内核实现等任务的问题导向解决方案。

• 参考：技术规范和 API 映射。

• 解释：对架构和硬件基础的深入解析。

文档使用即开即用的笔记本构建。

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

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

案例研究

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

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

我们采用基于语料库分层和架构平衡选择的稳健测试方法。该方法通过两层筛选机制实现现实应用普及度与架构多样性的平衡，避免对次要变体的冗余测试，同时确保真实生产权重得到充分验证。第一层筛选优先选择基于社区采用率、下载速度和趋势评分的头部模型（如 Qwen、Gemma、Mistral、DeepSeek、GLM、Kimi、Llama、Whisper、Stable Diffusion、FLUX 和 Vision Transformers），以确保关键社区模型完美运行。第二层筛选从第一层中选取子集，确保从 500 多个不同模型家族中至少选择 1 个代表性模型。我们使用第二层筛选在全面的 PyTorch 操作符和编译器模式集合上，验证 TPU 与 GPU 基线的数值准确性。

验证过程从第一阶段开始，重点进行烟雾测试和降级验证以确保运行时稳定性。在此初始阶段，流水线使用合成权重实例化模型，验证 eager 模式前向执行，并确认图追踪、FX-to-StableHLO 图生成和 XLA 编译成功无误（例如无崩溃或异常）。确认运行时稳定性后，第二阶段通过领域特定评估指标，将数值正确性与既定基线进行对比。对于因果 LLM，通过困惑度（Perplexity = e L = exp((1/ N ) Σ i =1 N L i )）衡量数学保真度，当 TPU 与基线困惑度的相对差异保持在 0.10 以内时测试通过。对于视觉、音频和编码器架构中的高维特征空间，我们采用双重阈值相对 MSE 指标（Relative MSE = MSE(Y TPU , Y Baseline ) / (Var(Y Baseline ) + 10 −7 )），要求相对 MSE ≤ 0.10 或绝对 MSE < 0.05，以防止近似常数基线输出。

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

| 指标类别 | 范围/定义 | 结果 | |----------------------|------------------------------------------------|--------------------------| | 总验证模型语料库 | 端到端测试的独立开源模型数量 | 4,917 个模型 | | 独特架构家族 | 跨 Transformers、Diffusers 等的独立架构类别数量 | 500+ 架构家族 | | 第一层（流行度）- 通过率（即无崩溃） | 社区下载量高的 LLM、视觉、音频和多模态模型 | 98% - 4,821/4,917 | | 第二层（架构）- 数值一致性 | 在树内独立架构代表模型 | 97% - 638/657（相对于 GPU） | | 有界动态支持（针对 SymInts） | 在 839 个模型（来自 5K 语料库的因果 LLM 列表）上评估正确性 | 确认 834 个模型（99%）数值正确 |

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

TorchTPU 致力于成为 PyTorch 的真正原生 TPU 后端。在我们看来，这一目标意味着在运行纯 PyTorch、设备无关代码时应具有良好的性能。我们通过两个示例来证明这一点。

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

• 在单个 TPUv7 芯片上，对于长度为 5,120 的序列，推理（预填充）性能达到约 25.7k TPS，模型利用率为 59%；对于长度为 10,240 的序列，推理性能达到约 22.4k TPS，模型利用率为 51%。请注意，此示例使用了 FSDP（通过原生 PyTorch 集合操作实现），由于单个 TPUv7 芯片向 PyTorch 暴露了两个设备（对应两个包含的 TensorCores），因此世界规模设置为 2。

• 在 4 个 TPUv7 芯片（即完整单主机）上使用 FSDP 进行训练，对于长度为 5,120 的序列，每个 TPU 的性能约为 5.6k TPS，模型利用率为 38%；对于长度为 10,240 的序列，每个 TPU 的性能约为 4.6k TPS，模型利用率为 32%。在第二个示例中，我们启用了激活检查点，该功能仅在 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 移植版本实现了：

• 在 4 个 TPUv7 芯片上训练：对于长度为 5,120 的序列，每个 TPU 的性能约为 5.3k TPS，模型利用率为 36%；对于长度为 10,240 的序列，每个 TPU 的性能约为 4.2k TPS，模型利用率为 29%。

总之，所有这些示例表明，即使是一个原本为 GPU 编写的纯 PyTorch 模型实现，仅通过简单的 PyTorch 修改和标准 PyTorch 优化（如 torch.compile 和 FSDP），而无需诉诸于高级、设备特定的优化技术（如自定义内核），仍然可以在 TPU 上表现良好。我们认为这是 TorchTPU 采用的重要特性。当然，对于追求峰值性能的生产工作负载，我们绝对建议采用基于 Pallas 内核和生产级堆栈（如 vLLM-TorchTPU、SGLang-TorchTPU 或 TorchTitan-TorchTPU）的最优方法。

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

我们使用 TorchTitan 验证了 TorchTPU 在 1 个 TPU v7 Pod 上的扩展能力。为此测试，我们选择了 Qwen3 MoE 235B 模型。作为初步的烟雾测试，我们确认了 TorchTitan 中使用 TorchTPU 实现的模型性能与 JAX 在 MaxText 上的性能相当（在 64 个 TPU v7 的设置下，PyTorch/TorchTitan 的性能为 1828 TPS，JAX/MaxText 的性能为 1937 TPS）。

以下为翻译后的 Markdown 内容：

以下讨论的数据是基于以下配置计算得出的：

• 模型：qwen3_moe_235b_scan
• 本地批量大小（即每个 TensorCore）：16，序列长度：4096
• 并行性：单切片运行使用 FSDP（通过 SimpleFSDP），多切片运行使用 HSDP
• 精度：BF16 混合精度（FP32），包括梯度归约（f32）、RoPE（f32）、优化器状态（f32）
• 内核（Pallas）：Tokamax Splash Attention（显著提升性能）、融合交叉熵损失内核、MoE 层的分组矩阵乘法
• 扫描优化：94 个解码器层以堆叠权重的形式运行在一个扫描循环中
• 解码器层输入的激活卸载（作为扫描优化的一部分）
• 集体流水线技术，用于重叠相邻层的计算（扫描的一部分）
• 完全激活检查点
• 使用 TorchTitan 的 Graph Trainer 进行联合前向和反向传播
• 专家不平衡：使用 TorchTitan 的辅助损失防止 MoE 不平衡
• 训练数据集：Hugging Face C4 en 数据集，采用自定义数据分片

拓扑结构

| TPU v7 芯片数量 | TensorCore 数量 | 总批量大小 | 每 TensorCore TPS | MFU | 步骤时间 | |----------------|----------------|------------|------------------|-----|---------| | 1 slice with FSDP | 1024 | 2048 | 32768 | 2,236 | 28.8% | 29.3s | | 2 slices with HSDP | 4096 | 65536 | 2,214 | 28.5% | 29.6s | | 4 slices with HSDP | 8192 | 131072 | 2,095 | 26.9% | 31.3s | | 8 slices with HSDP | 16384 | 262144 | 2,104 | 27.1% | 31.2s |

在所有拓扑结构中，我们的数据验证了弱扩展性：整体计算量呈线性增长，如图所示，每秒处理的总 token 数呈线性增长，MFU 基本保持稳定。

我们的数据还显示，不同并行级别下的损失曲线几乎完全一致。

在多切片 HSDP 配置中，前向传播过程不涉及切片间的通信，每个切片独立处理其分配的批量部分。但在反向传播过程中，切片首先计算第 t 层的梯度，然后通过 all-reduce（基于 DCN 和 Megascale-XLA）将所有切片的梯度相加。计算随后继续处理第 t-1 层的梯度。得益于 XLA，TorchTPU 能够完全重叠第 t+1 层的 all-reduce 操作与第 t 层的计算，如图所示。实现这种重叠是导致 TensorCore 每秒处理量（TPS）从 1 到 8 切片基本保持稳定的主要原因，如上表所示。

案例研究 4：基于 vLLM（SemiAnalysis InferenceX）的 Qwen3.5 大语言模型推理

大语言模型（LLM）推理——特别是针对大规模专家混合（MoE）架构的推理——代表了 AI 云基础设施中最苛刻且成本敏感的工作负载之一。在 SemiAnalysis 发布的独立基准测试报告（《TPU 推理外部化全速前进——InferenceX》）中，Google 的 TPU v7 配合 vLLM-TorchTPU 栈实现了行业领先的性价比。该基准测试对 TorchTPU 在大规模 MoE 工作负载上的软件就绪状态进行了第三方推理审计，证明了新的 vLLM 栈在 TPU 上实现了强大的成本效率和高吞吐量，适用于社区公开可用的 TPU 系统。

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

• 在InferenceX 8K/1K上，每百万tokens成本最低为2.01美分（较9月首次发布时的3.6美分下降）；
• 在AgentX上，每百万tokens成本最低为0.83美分，并实现SOTA交互性（用户响应速度）。在帕累托前沿全面领先。

这些基准测试结果反映了TPU系统在实际场景中强大的服务效率，凸显了在Google自研ML芯片上部署大规模模型的巨大经济可行性。

实现这些具有竞争力的性能和总体拥有成本（TCO）结果，需要通过自定义Pallas内核和针对性运行时调优来解决关键的软硬件对齐瓶颈，以下表格中列举了部分示例：

优化项 | 瓶颈 | 实现方式 --- | --- | --- 专家并行MoE层作为单一SPMD程序 | 在DP8 + EP下，vLLM的MoE层将每个rank的tokens收集到所有rank，进行路由、置换、运行两次分组矩阵乘法和合并操作，然后将结果reduce-scatter回传。在预填充密集型步骤中，concurrency为256时，该reduce-scatter操作每层耗时1,190微秒，占设备总时间的22%，且无其他操作并行执行 | 一个SPMD程序在所有八个TensorCore上运行该层。每个TensorCore将其行量化为FP8，并在SparseCore上异步all-gather。同时，通过Pallas top-k选择属于自己的tokens的专家，并在第二次小规模all-gather中共享专家索引和行缩放因子。Pallas内核随后在其路由到的行上运行每个本地专家的两个矩阵乘法，通过索引获取数据。它以FP8格式（带每行缩放因子）通过远程DMA（pltpu.make_async_remote_copy）将结果行推送到拥有该token的TensorCore。这些推送操作与后续专家的矩阵乘法重叠。每个TensorCore随后根据其路由得分对到达的行进行加权求和。每层操作取代了路由行的argsorts、SparseCore的置换和合并、两个分组矩阵乘法以及合并reduce-scatter操作。

基于通道的序列注意力布局（无填充FP8 KV缓存） | TPU的向量单元以32位字工作，KV缓存布局与之匹配，因此每个字包含四个FP8值。标准KV布局用多个头的K和V填充一个字。在8路张量并行下，每个核心仅持有模型两个KV头中的一个，导致每个字的一半为填充：每个核心每个token的缓存占用1024字节，仅存储512字节数据 | 基于通道的布局将页面的tokens沿128个通道排列，头维度沿子通道排列，使每个字存储相同token的四个值，无任何浪费。这将KV内存减少一半，每解码步骤读取的KV字节数也减少一半。基于通道的布局是vllm-torchtpu现在支持的通用布局。

Gated DeltaNet优化

60层中有45层使用Gated DeltaNet，该结构为每个注意力头保持固定大小的128 × 128状态，而非使用可扩展的KV缓存。解码过程需要在每一步读写所有序列的状态；在并发数为64时，移除所有内核的算术运算仅使速度提升6%。逐词计算时，预填充是顺序递归过程，导致矩阵单元处于空闲状态。

一个Pallas内核将卷积和递归操作融合，批量处理解码序列，并以64词块为单位处理预填充，这些操作大部分被重写为矩阵乘法。状态与KV缓存共享相同的分页池。内核直接从池页中流式读取数据，并通过修改字节寻址方式而非转换数据格式，在原地读取递归状态。结果是共享池不会影响吞吐量：在我们的测试中，所有配置下共享池的吞吐量均等于或高于独立缓存。

这些结果突显了现代服务基础设施中的TPU优化，证明总拥有成本（TCO）和软硬件对齐程度决定了实际服务效率。最终，配备TorchTPU后端的新版vLLM-TPU使工程团队能够使用标准PyTorch语法，在公开可用的TPU基础设施上部署现成的开源权重模型。

案例研究5：使用TorchTitan进行大语言模型微调

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

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

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

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

• 简易FSDP：一种编译器友好的FSDP版本。FSDP2通过模块钩子注入all-gathers和reduce-scatters，这会将torch.compile拆分为每个操作的分发，而简易FSDP则在追踪图中将它们表示为普通操作。这使集合操作对XLA可见，TPU编译器可以围绕它们进行融合、与MXU计算重叠，并将其卸载到SparseCore。

• 扫描高阶操作（torch._higher_order_ops.scan）：将重复的Transformer块作为堆叠权重的单个循环运行。这使编译时间和收集参数内存随深度增长保持恒定。每个层的集合操作位于循环内部，因此XLA可以将其与计算重叠：延迟隐藏调度器将每个层的all-gathers隐藏在该层的计算之后，集合流水线器可选择性预取下一层的all-gather，在反向传播中允许一层的reduce-scatter在下一层计算时完成。

• Pallas SplashAttention内核：用自定义的Pallas内核替换缺乏矢量瓦片对齐的默认注意力实现，支持完整的分组查询注意力（GQA）、偏置和dropout，在扩散注意力模块上实现+17.5%的加速。

• 融合Pallas交叉熵损失：通过将Tokamax的在线约简内核移植到原地计算损失，避免在大词汇量下生成[B×S×V]logits时出现内存不足（OOM）故障。

• 图训练器：将完整的训练步骤（前向、损失、反向和优化器更新）捕获为单个图，编译成一个XLA程序。这为TPU编译器提供了全局视图，用于融合、集合调度和内存规划，并让TPU特定的传递决定每个激活是否保存、重新计算或卸载到主机，同时捐赠优化器缓冲区使状态就地更新而非双缓冲。

• 激活卸载：反向传播所需的激活数据可以在前向计算后移动到主机内存，并在反向传播使用前复制回来，这为更大模型和更长序列释放了高带宽内存（HBM）。图训练器会将卸载和重新加载操作插入到追踪的图中。在扫描操作下，这些操作位于循环体内部，因此激活数据逐层移动。XLA将这些复制操作转换为异步传输，在反向传播时，当前层计算的同时将下一层的激活数据复制回来。这对于最大规模的模型（如Llama3-405B和Qwen3-235B）在目标序列长度下适配至关重要。

TorchTitan在TorchTPU上的持续开发聚焦于四个战略重点：扩展针对TPU优化的模型列表；通过torchao集成FP8和FP4低比特量化；以及通过将训练循环与vLLM/SGLang采样后端结合，在VeRL中支持端到端强化学习（RL）。

案例研究6：使用TorchRec和SparseCore实现的硬件加速推荐模型

推荐和排序架构（如Meta的深度学习推荐模型DLRM）由于高基数嵌入表面临关键的内存和计算瓶颈。在标准GPU集群上，管理这些表可能需要通过分布式模型并行（DMP）进行复杂的软件分片，并进行独立的优化器更新，这通常会引入延迟。为弥补这一差距，Google与Meta PyTorch团队合作，将原生TPU后端直接上游到TorchRec仓库，使推荐模型能够绕过主机侧计算瓶颈，直接接入TPU SparseCore（SC）硬件。

• 架构集成。该实现创建了一个无缝的数据流水线，其中高基数查找通过专用硬件加速。多线程C++ PyBind后端在主机CPU上运行，以并行化分区映射和索引填充，输出直接映射到TPU SparseCore单元的结构化批次。在TPU内部，我们实现了原地优化器融合：与传统在高带宽内存（HBM）中显式生成梯度张量（.grad）（这一过程通常会复杂化内存管理）不同，我们直接将反向传播映射到SparseDenseMatmulGrad。这确保参数更新（SGD、AdaGrad或Adam）在原地进行，完全消除了梯度内存开销。最后，低级ATen操作符（如SparseDenseMatmulOp（用于前向查找）和RaggedAllToAll（用于光学芯片间集体通信））提供了直接在硅芯片上执行的硬件加速内核。

• 性能与结果。在分布式32芯片TPU v7x（Ghostfish）集群上对行业标准DLRM v2模型进行评估时，该集成实现了与Google高度优化的JAX基线100%的性能一致性，达到相同的18.3毫秒步延迟。值得注意的是，这种性能匹配是在应用高级优化之前，使用基线SparseCore内核建立的。收敛性和数学保真度同样稳健，与JAX实现相比损失轨迹完全一致。

当前开发重点在于通过 TensorCore/SparseCore 流水线进一步提升带宽利用率——将 MLP 执行与嵌入查找操作交错进行，并采用先进的表格堆叠技术统一内存缓冲区、降低内核开销。此外，正在进行的工作旨在通过 ATen 内核暴露更底层的 SparseCore 嵌入查找操作（例如 SparseMapRow、SparseGather 等），使用户能够实现自定义查找算法和高级分片策略。这种集成巩固了 TPU 作为下一代推荐架构高效引擎的地位。

未来工作方向

随着 TorchTPU 现已在 GitHub 上公开发布，此次初始开源版本仅代表我们开发历程的开端。在与更广泛的 PyTorch 社区直接互动的过程中，我们未来 6 至 12 个月的路线图将聚焦于生产环境强化、生态扩展以及持续整合社区反馈，以弥合 TPU 硬件能力与 PyTorch 原生开发工作流之间的差距。我们正在优先推进四个战略支柱以推动这一演进：

• 基础强化与用户效率：我们正在通过构建 ATen 操作符的预编译动态内核库来完善“True Eager”体验。同时使 XLA 有界动态性完全具备生产就绪能力，让 torch.compile 堆栈能够处理 SymInts 而无需额外重新编译，并优化 GKE 打包流程以降低跨 PyTorch 版本部署的门槛。此外，我们持续扩展操作符覆盖范围并推进量化支持——特别针对下一代 TPUv8 架构的 FP8 和 FP4 格式，以及实现嵌套编译器区域以减少编译时间。

• Pod 级分布式训练与高级训练：在成功验证 Qwen3 的 POD 级训练基础上，我们正在扩展至更大规模的模型架构和分片技术。通过将可扩展训练与基于 Ray Core 协调的高吞吐量 vLLM/SGLang 采样后端紧密耦合，我们正在实现端到端强化学习（RL）训练循环。为增强分布式稳定性，我们正在验证基于共识的编译握手机制，以在不依赖脆弱 @spmd_safe 注解的情况下安全优化集体操作。

• 新兴架构、推理与服务：我们承诺主动支持开源生态系统中出现的前沿模型架构——包括推理模型、多模态和生成式扩散模型（如 Wan 2.2、Flux Klein）以及新型专家混合（MoE）变体。我们的工作重点是优化 vLLM 和 SGLang 的服务栈以实现低延迟生产，整合未来 TPU 世代的原生 FP8/FP4 量化支持，并确保新架构范式在发布时即具备 TPU 支持。

• 高效通信与诊断：我们正在优化 TPU 切片内部及跨切片的集体操作，并添加零拷贝 DMA 以实现硬件辅助内存卸载。为支持大规模调试，我们正在转向基于 WebGPU 的 Trace Viewer 2.0，并部署“飞行记录仪”持续性能分析（XProfz）以识别生产环境中的间歇性瓶颈。

TorchTPU 的构建基于一个信念：成功的生态系统通过透明且协作的工程实践蓬勃发展。随着我们在 2026 年 10 月的 PyTorch 大会上开源四个核心仓库（TorchTPU、vLLM-TorchTPU、SGLang-TorchTPU 和 TorchTitan），我们积极寻求关于模型架构优先级和工作负载排序的反馈。我们邀请开发者、研究人员和合作伙伴通过 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, Bill Jia

还有许多其他成员……

/post-content

/inner-wrap