How to use Google microbenchmarks for evaluating TPU performance

如何使用 Google 微基准测试评估 TPU 性能 - Google 开发者博客

Google Tag Manager (noscript)

End Google Tag Manager (noscript)

HTML

如何使用 Google 微基准测试评估 TPU 性能

2026 年 7 月 30 日

Junjie Qian

软件工程师

Chi Shuen Lee

Yu-Hsuan (Amy) Lin

Haixiong (Sean) Wang

高级软件工程师

分享

• Facebook

• Twitter

• LinkedIn

• 邮件

要确定 TPU 的实际性能，需要超越产品规格，进行细粒度的真实测量。虽然理论吞吐量指标表明了显著的计算能力，但关键的衡量标准仍然是这些加速器在特定架构环境和工作负载限制下的行为表现。为此，我们编写了一套微基准测试套件，用于对示例工作负载进行 TPU 评估。

微基准测试是评估 TPU 等加速器的关键工具，因为它们能提供对基础硬件能力的细粒度量化指标。它们将性能评估划分为关键功能领域，使您能够准确判断设备是否达到理论性能规格，并识别特定的性能差距或架构特有的瓶颈。

在本博客中，我们将更详细地探讨这套微基准测试套件：它测量的内容、与 TPU 架构的关系以及如何调整它。我们还将展示一个示例工作负载的案例研究。

关键类别和指标

微基准测试套件在五个核心领域提供标准化性能评估：

• 网络：表征通过芯片间互连（ICI）进行集体通信操作的互连性能，这对于跨多个芯片扩展模型至关重要。关键测试包括 all-gather、all-reduce、reduce-scatter 和 all-to-all。指标包括吞吐量（GB/s）和延迟（秒）。

• 计算：量化原始计算吞吐量。主要操作是通用矩阵乘法（GEMM/Matmul），以 TFLOPs 和模型 FLOPs 利用率（MFU）衡量，这表明工作负载利用 TPU 矩阵单元（MXU）架构的效率。

• 高带宽内存（HBM）：测量芯片内内存的带宽。这评估了在旨在饱和内存结构的大向量操作期间的基本内存访问速度。核心指标是按 GB/s 衡量的内存带宽。

• 主机传输：通过测量主机 CPU 和加速器 HBM 之间的 PCIe 互连的数据传输速率来表征 I/O 效率。这包括主机到设备（H2D）和设备到主机（D2H）传输，以 GB/s 衡量。

• 稀疏分页注意力（RPAv3）：专注于针对基于变压器的模型的基准测试，特别测量注意力（tokamax splash）和批量矩阵乘法（BMM）等原语的吞吐量。这些有助于预测端到端推理性能，并识别首次令牌延迟（TTFT）中的瓶颈。

微基准测试 TPU 模块

我们的套件中的网络、计算、HBM 和主机传输微基准测试通过测试独立组件，提供对理论性能的细致评估。这种详细的组件分析对于跨多个维度调试和优化系统性能至关重要。例如，它可以帮助确定设备是否达到其规格，或者特定的架构瓶颈（如 VMEM 局部性或互连效率）是否限制了工作负载的规模和速度。

微基准测试套件

这种实证评估的核心工具是我们的加速器微基准测试 GitHub 仓库。该套件通过几个关键目录提供了一种结构化的硬件验证方法：

• src/ : 包含微基准测试的基本实现和源代码
• configs/ : 管理每个测试的调优参数，如数据缓冲区大小、迭代次数和支持的数据类型（例如 bf16、fp8、fp32）

# 使用示例配置运行计算微基准测试

kubectl apply -f Ironwood/guides/collectives/tpu7x-2x2x1-ici-all-gather-microbenchmark.yaml

Shell

已复制

通过利用这些配置方案，您可以将之前讨论的高层指标（如 ICI 带宽或 MXU 利用率）与实际运行时数据关联起来。例如，configs 文件夹中的集体 YAML 文件允许直接测量架构部分描述的 ICI 互连效率，帮助确保物理互连的健康状况符合理论预期。

性能调优影响

微基准测试将性能优化从试错过程转变为实证工程学科。通过建立“光速（SOL）”基线（硬件模块的理论极限），它们为训练和推理效率提供了明确的目标。

Roofline 模型：诊断的北极星

微基准测试的核心价值在于其定义 Roofline 模型边界的能力。通过测量 MXU 峰值吞吐量和 HBM 带宽，开发者可以立即确定工作负载的瓶颈类型：

• 计算受限 : 模型达到 MXU 平台。调整内存访问不会带来任何好处；优化必须聚焦于内核效率或减少 FLOPs。
• 内存受限 : 模型受限于 Roofline 的斜率“屋檐”。通过优化数据局部性（VMEM）或减少 HBM 通信可以提升性能。
• 网络受限 : ICI 或 DCN 的停滞阻止芯片发挥计算潜力。这表明需要更好的分片策略或通信重叠技术。

硬件感知的模型架构

微基准测试揭示了 TPU 硬件的具体“特性”。例如，在 Ironwood（TPU7x）架构中，256x256 系统阵列对操作数形状施加了物理约束。基准测试验证了使用 head_dim 为 128（在较早的 Llama 变体中常见）的模型 MXU 利用率较低。这一洞察促使研究人员与硬件效率最大化对齐，共同设计符合 256 字节边界的前沿模型。

软件杠杆与优化策略

一旦确定瓶颈，微基准测试可以帮助您选择合适的软件杠杆：

• 内核选择：您可以将标准 JAX 操作与专用的 Pallas 内核或 Tokamax Splash Attention 进行对比。微基准测试提供了确定哪种实现能最小化反向传播开销的真实依据。

• 分片与网格调优：通过测量不同拓扑结构（如 2x2x2 与 4x2x1）下的集体通信性能，您可以优化全分片数据并行（FSDP）参数，确保通信时间完全隐藏在计算之后，即任务执行速度足够快，使 TPU 不需要停止等待响应。

• 重新计算：微基准测试有助于精确计算计算周期与内存压力之间的权衡，指导决策在何时将计算数据存储在内存中，或选择“重新计算”（重新计算）以避免 HBM 溢出。

预测性优化

最终，微基准数据可用于对大规模部署进行预测性建模。无需在完整的 TPU 切片上运行昂贵的测试，即可利用结果构建分析模型，准确预测扩展前的训练吞吐量、首次令牌生成时间（TTFT）、每输出令牌耗时（TPOT）和内存利用率（MFU）。

案例研究：Ironwood（TPU 7x）性能分析

在其中一个案例中，我们分析了在 4x4x4 TPU 7x 配置上部署的 110B 混合专家（MoE）训练工作负载的性能调优，利用微基准测试建立硬件基准并驱动对复杂稀疏工作负载的优化。以下是我们的发现：

• 屋顶线分析：对训练工作负载的诊断分析揭示了性能瓶颈的明显划分。前向传播和密集核心操作主要受计算限制，达到 1.85 PFLOPS。相反，MoE 架构中的路由和注意力原语保持内存限制，由于 HBM 饱和，仅达到 SOL 的 30-60%，在专家调度期间，稀疏激活模式使操作变得高度 I/O 限制，将传入请求发送到子网络。

• 与微基准测试的相关性：微基准测试结果直接指导了干预措施，将训练步骤时间减少了 21.2%。ICI 集体基准测试识别出通过 SparseCore 集体卸载解决的通信阻塞，而 HBM 带宽测试识别出内存限制的注意力原语，促使采用 Tokamax Splash Attention。

为 TPU 性能优化注入规范性

总之，微基准测试套件提供了必要的结构化验证，以弥合理论规格与实际运行时性能之间的差距，使您能够充分利用 TPU 架构在大规模部署中的潜力。

欲了解更多信息，请查看以下资源：

• Ironwood：首个面向推理时代的 Google TPU
• 在 TPU v7x（Ironwood）上优化前沿模型训练
• AI 加速器性能与基准测试
• 使用 Ironwood TPU 进行训练的开发者指南

发布于：

• AI
• 云服务
• 操作指南
• 学习
• 探索

上一篇

下一篇

相关文章

列表

AI

云服务

教程

案例研究

扩展智能体强化学习：使用 Tunix 进行高吞吐量智能体训练

2026 年 7 月 21 日

操作指南

公告

在 TPU 上运行 Ray，第 1 部分：基础

2026 年 7 月 20 日

在 TPU 上运行 Ray，第 2 部分：Ray AI 库

2026 年 7 月 24 日

导航点