在 Cloud TPU 节点上分析模型的性能

通过对模型进行性能剖析，您可以优化 Cloud TPU 上的训练性能。 要分析模型，您可以使用 TensorBoard 和 Cloud TPU TensorBoard 插件。安装用 说明，请参阅 TensorBoard 安装说明。

如需详细了解如何将 TensorBoard 与其中一个支持的框架搭配使用，请参阅以下文档：

前提条件

TensorBoard 作为 TensorFlow 的一部分安装。TensorFlow 默认安装在 Cloud TPU 节点中。您也可以手动安装 TensorFlow。 无论采用哪种方式，都需要一些额外的依赖项。通过运行以下命令进行安装：

(vm)$ pip3 install --user -r /usr/share/models/official/requirements.txt

安装 Cloud TPU TensorBoard 插件

通过 SSH 连接到您的 TPU 节点：

$ gcloud compute ssh your-vm --zone=your-zone

运行以下命令：

  pip3 install --upgrade "cloud-tpu-profiler>=2.3.0"
  pip3 install --user --upgrade -U "tensorboard>=2.3"
  pip3 install --user --upgrade -U "tensorflow>=2.3"

捕获性能剖析文件

您可以使用 TensorBoard 界面或通过编程方式捕获配置文件。

使用 TensorBoard 捕获配置文件

当您启动 TensorBoard 时，它会启动 Web 服务器。将浏览器指向 TensorBoard 网址后，浏览器会显示一个网页。您可以在该网页上手动捕获配置文件并查看配置文件数据。

启动 TensorFlow 性能分析器服务器

tf.profiler.experimental.server.start(6000)

这将在您的 TPU 虚拟机上启动 TensorFlow 性能分析器服务器。

启动训练脚本

运行训练脚本并等待，直到您看到表明您的模型正在训练的输出。其样式取决于您的代码和模型。查找 Epoch 1/100 之类的输出。或者，您也可以前往 Cloud TPU 页面 在 Google Cloud 控制台中，选择您的 TPU 并查看 CPU 利用率图表。虽然 但这并不显示 TPU 利用率， 训练模型。

启动 TensorBoard 服务器

打开新的终端窗口，并通过端口转发通过 SSH 连接到您的 TPU 虚拟机。这个 可让您的本地浏览器与运行 使用 TPU 虚拟机

 gcloud compute tpus execution-groups ssh your-vm --zone=us-central1-a --ssh-flag="-4 -L 9001:localhost:9001"

在您刚刚打开的终端窗口中运行 TensorBoard，并指定目录 其中，TensorBoard 可以使用 --logdir 标志写入性能剖析数据。例如：

TPU_LOAD_LIBRARY=0 tensorboard --logdir your-model-dir --port 9001

TensorBoard 会启动一个网络服务器并显示其网址：

Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.3.0 at http://localhost:9001/ (Press CTRL+C to quit)

打开网络浏览器并转到 TensorBoard 输出中显示的网址。点击 TensorBoard 页面右上角的“重新加载”按钮，以确保 TensorBoard 已完全加载性能剖析数据。默认情况下，会显现选中 Scalars 标签的 TensorBoard 页面。

图片

捕获 TPU 节点上的性能剖析文件

  1. 从屏幕顶部的下拉菜单中选择“剖析”
  2. 选择“捕获性能剖析文件”按钮
  3. 选择“TPU 名称”单选按钮
  4. 输入您的 TPU 名称
  5. 选择“拍摄”按钮。

以编程方式捕获配置文件

以编程方式捕获配置文件的方式取决于您使用的机器学习框架。

如果您使用的是 TensorFlow，则使用以下命令启动和停止性能分析器： tf.profiler.experimental.start()和tf.profiler.experimental.stop() 。如需了解详情，请参阅 TensorFlow 性能指南。

如果您使用的是 JAX，请使用 jax.profiler.start_trace() 和 jax.profiler.stop_trace() 分别启动和停止性能分析器。如需了解详情，请参阅剖析 JAX 程序的性能。

捕获配置文件常见问题

有时，当您尝试捕获跟踪记录时，可能会看到 以下：

No trace event is collected after xx attempt(s). Perhaps, you want to try again
(with more attempts?).Tip: increase number of attempts with --num_tracing_attempts.

Failed to capture profile: empty trace result

如果 TPU 未在积极执行计算、训练步骤耗时过长或存在其他原因，就可能会出现这种情况。如果您看到这条消息， 请尝试以下操作：

  • 尝试在运行几个周期后捕获配置文件。
  • 请尝试在 TensorBoard 的 Capture Profile 对话框中延长性能分析时长。训练步骤可能耗时过长。
  • 确保虚拟机和 TPU 使用相同的 TF 版本。

使用 TensorBoard 查看性能剖析文件数据

捕获一些模型数据后，系统会显示分析标签页。您可能需要点击 TensorBoard 页面右上角的“重新加载”按钮。数据可用后，点击 性能剖析 (Profile) 标签页会显示一系列有助于性能分析的工具：

Trace Viewer

Trace viewer 是性能剖析 (Profile) 下的 Cloud TPU 性能分析工具。此工具使用 Chrome 跟踪记录事件分析查看器，因此仅适用于 Chrome 浏览器。

Trace Viewer 包含一个时间轴，其中显示的内容包括：

  • 由 TensorFlow 模型执行的操作的时长。
  • 系统的哪个部分（TPU 还是主机）执行了操作。通常，主机执行馈入操作，它会预处理训练数据并将其传输到 TPU，而 TPU 执行实际的模型训练。

您可以通过 Trace Viewer 找出模型中的性能问题，然后采取措施来加以解决。例如，从总括层面来讲，您可以确定占用大部分时间的是馈入还是模型训练。展开细目，您可以确定哪些 TensorFlow 操作执行时间最长。

请注意，Trace Viewer 针对每个 Cloud TPU 仅限访问 100 万个事件。如果您需要访问更多事件，请使用 Streaming Trace Viewer。

Trace Viewer 界面

如需打开 Trace Viewer，请转到 TensorBoard，点击屏幕顶部的性能剖析 (Profile) 标签页，然后从工具下拉列表中选择 trace_viewer。Trace Viewer 会显示最近的运行：

图片

此屏幕包含以下主要元素（标有上述数字）：

  1. 运行下拉列表。包含已捕获其跟踪记录信息的所有运行。默认视图为最近的运行，您也可以打开下拉列表选择其他运行。
  2. 工具下拉列表。选择不同的性能剖析工具。
  3. 主机下拉列表。选择包含 Cloud TPU 集的主机。
  4. 时间轴窗格。显示 Cloud TPU 和主机随时间执行操作的情况。
  5. 详细信息窗格。显示时间轴窗格中所选操作的更多信息。

以下是时间轴窗格的详细信息：

图片

时间轴窗格包含以下元素：

  1. 顶栏。包含各种辅助控件。
  2. 时间轴。显示相对于跟踪记录开始时的时间。
  3. 版块和跟踪记录标签。每个版块包含多个跟踪记录，左侧有一个三角形，点击此三角形可展开和收起该版块。系统中的每个处理元素都有一个版块。
  4. 工具选择器。包含与 Trace Viewer 交互的各种工具。
  5. 事件。这些事件显示操作的执行时间或元事件（例如训练步）的时长。
  6. 垂直标签栏。它对 Cloud TPU 没有用处。此栏是 Chrome 提供的通用跟踪记录查看器工具的一部分，用于各种性能分析任务。

版块和跟踪记录

Trace Viewer 包含以下版块：

  • 每个 TPU 节点有一个版块，上面标有 TPU 芯片的编号及芯片内 TPU 节点的编号（例如“芯片 2: TPU 核心 1”）。每个 TPU 节点版块包含以下跟踪记录：
    • 步。显示 TPU 上运行的训练步的时长。
    • TensorFlow 操作。显示在 TPU 上执行的 TensorFlow 操作。
    • XLA 操作。显示在 TPU 上运行的 XLA 操作。（每个操作都会转换为一个或多个 XLA 操作。XLA 编译器会将这些 XLA 操作转换为在 TPU 上运行的代码。）
  • 针对在主机的 CPU 上运行的线程的版块，其标签为“Host Threads”。对于每个 CPU 线程，此版块都包含一个跟踪记录。注意：您可以忽略版块标签旁边显示的信息。

时间轴工具选择器

您可以使用 TensorBoard 中的时间轴工具选择器与时间轴视图进行交互。您可以点击时间轴工具或使用以下键盘快捷键激活并突出显示该工具。如需移动时间轴工具选择器，请点击顶部的虚线区域，然后将选择器拖动到所需位置。

以下是各个时间轴工具的使用方法：

 选择工具
点击事件可将其选中，拖动可选择多个事件。有关所选事件（名称、开始时间和时长）的更多信息将显示在“详细信息”窗格中。

平移工具
拖动可在水平和垂直方向平移时间轴视图。

缩放工具
沿水平（时间）轴方向，向上拖动可放大视图，向下拖动可缩小视图。鼠标光标的水平位置决定缩放操作的中心。

注意：缩放工具有一个已知错误：当鼠标光标在时间轴视图外面时松开鼠标按钮，仍会进行缩放。如果您遇到这种情况，点击时间轴视图即可停止缩放。

计时工具
水平拖动可以标记时间间隔。间隔的时长会显示在时间轴上。如需调整间隔，请拖动其末端。如需清除间隔，可点击时间轴视图内的任意位置。

请注意，如果您选择其他某个工具，间隔标记仍会保留。

图表

TensorBoard 为模型及其性能提供了大量可视化或图形表达。将图表与 Trace Viewer 或 Streaming Trace Viewer 结合使用可以优化模型并提升其在 Cloud TPU 上的性能。

模型图

建模框架可能会通过模型生成图表。图中的数据存储在通过 --logdir 参数指定的存储分区内的 MODEL_DIR 目录中。您无需运行 capture_tpu_profile 即可查看此图。

如需查看模型的图，请选择 TensorBoard 中的图表标签页。

图片

结构图中的单个节点表示单个操作。

TPU 兼容性图

图表标签页包含一个兼容性检查器模块，用于检查和显示可能会在模型运行时引起问题的操作。

如需查看模型的 TPU 兼容性图，请选择 TensorBoard 中的图 (Graphs) 标签页，然后选择 TPU 兼容性 (TPU Compatibility) 选项。该图以绿色表示兼容（有效）操作，以红色表示不兼容（无效）操作。

图片

一个给定的节点可以显示两种颜色，每种颜色按照节点中 Cloud TPU 兼容操作的百分比显示。如需查看示例，请参阅解读兼容性结果。

图右侧显示的兼容性摘要面板展示所选节点的所有 Cloud TPU 兼容操作的百分比、这些操作的属性以及不兼容操作列表。

点击图中的任意操作即可在摘要面板中显示其属性。

图片

请注意，兼容性检查器不会评估使用手动设备配置明确分配给非 TPU 设备的任何操作。此外，检查器实际上并不编译用于执行的模型，因此请务必将结果解读为对兼容性的估算。

解读兼容性结果

性能剖析

性能剖析 (Profile) 标签页在您捕获某些模型数据后显示。您可能需要点击 TensorBoard 页面右上角的“重新加载”按钮。数据可用后，点击 性能剖析 (Profile) 标签页会显示一系列有助于性能分析的工具：

性能剖析概览页面

概览页面 (overview_page) 位于性能剖析 (Profile) 下方，其提供模型在捕获运行期间的执行情况汇总。该页面展示所有 TPU 的聚合概览页，以及整体输入流水线分析。您可以在主机下拉列表中选择单个 TPU。

该页面在以下面板中显示数据：

图片

  • 性能摘要

    • 平均单步用时 - 所有采样步的平均单步用时
    • 主机空闲时间 - 主机空闲时间百分比
    • TPU 空闲时间 - TPU 空闲时间百分比
    • FLOPS 利用率 - TPU 矩阵单元的使用百分比
    • 内存带宽利用率 - 使用的内存带宽所占百分比

  • 单步用时图。展示所有采样步的设备单步用时图（以毫秒为单位）。蓝色区域对应的是 TPU 空闲时等待主机输入数据的单步用时部分。红色区域展示 Cloud TPU 实际运行的时间。

  • TPU 上最耗时的 10 个 TensorFlow 操作。展示占用了最多时间的 TensorFlow 操作：

    每行显示操作的自用时间（以所有操作所用时间的百分比表示）、累计时间、类别、名称以及实现的 FLOPS 率。

  • 运行环境

    • 使用的主机数
    • 使用的 TPU 类型
    • TPU 核心数量
    • 训练批次大小

  • 后续步骤建议 (Recommendation for next steps)。报告模型输入是否受限，以及是否出现 Cloud TPU 问题。推荐可用于找出性能瓶颈的工具。

Input Pipeline Analyzer

Input Pipeline Analyzer 可提供有关性能结果的数据分析。该工具展示 capture_tpu_profile 工具收集的 input_pipeline.json 文件的性能结果。

该工具能够立即告知您程序是否输入受限，并可指导您完成设备端和主机端分析，以调试并确定流水线的哪个阶段存在瓶颈。

如需深入了解如何优化流水线性能，请参阅有关输入流水线性能的指南。

输入流水线

当 TensorFlow 程序从文件中读取数据时，数据以流水线的方式从 TensorFlow 图的顶部开始读取。读取过程分为多个串联连接的数据处理阶段，其中一个阶段的输出是下一个阶段的输入。这个读取系统称为输入流水线。

从文件中读取记录的典型流水线具有以下阶段：

  1. 文件读取
  2. 文件预处理（可选）
  3. 将文件从主机传输到设备

效率低下的输入流水线可能会严重降低应用的速度。当应用在其输入流水线上花费大量时间时，我们将其称为输入受限。使用 Input Pipeline Analyzer 可以确定输入流水线的低效位置。

输入流水线信息中心

如需打开 Input Pipeline Analyzer，请选择性能剖析 (Profile)，然后从工具 (Tools) 下拉列表中选择 input_pipeline_analyzer。

信息中心包含三个版块：

图片

  1. 摘要。汇总整个输入流水线的情况，其中包含您的应用是否输入受限以及受限程度（如果受限）。
  2. 设备端分析。展示详细的设备端分析结果，包括设备单步用时以及每一步中等待所有核心的输入数据所占用的设备时间的范围。
  3. 主机端分析。展示主机端的详细分析，包括主机上输入处理时间的明细。
输入流水线摘要

第一个版块显示等待主机输入所占用的设备时间百分比，以报告您的程序是否输入受限。如果您使用的是已经过测试的标准输入流水线，该工具会报告大部分输入处理时间用于何处。例如：

图片

设备端分析

第 2 个版块显示设备端分析的详细信息，让您深入了解在设备和主机上花费的时间，以及等待主机输入数据所占用的设备时间。

图片

  1. 设备单步用时统计信息。报告设备单步用时的平均值、标准偏差和范围（最小值，最大值）。
  2. 单步用时。展示所有采样步的设备单步用时图（以毫秒为单位）。蓝色区域对应的是 Cloud TPU 空闲时等待主机输入数据的单步用时部分。红色区域展示 Cloud TPU 实际运行的时间。
  3. 等待输入数据的时间百分比。报告设备等待输入数据所用的时间比例的平均值、标准偏差和范围（最小值，最大值），基于设备单步用时总时长进行归一化。
  4. 根据步编号绘制的等待所有核心的输入数据所占用的设备时间范围。展示等待输入数据处理所花费的设备时间（表示为总设备单步用时的百分比）的折线图。由于不同核心所花费时间的比例不尽相同，因此系统会针对每一步绘制不同核心的比例范围。由于单步用时由最慢的核心决定，因此这个范围应尽可能小。
主机端分析

第 3 个版块显示主机端分析的详细信息，它将主机上的输入处理时间（Dataset API 操作所用时间）细分为几类：

  • 将要传输到设备的数据加入队列。向设备传输数据之前将数据加入馈入队列所花费的时间。
  • 数据预处理。用于预处理操作的时间，例如图片解压缩。
  • 预先从文件中读取数据。读取文件所花费的时间，包括缓存、预取和交错。
  • 按需从文件中读取数据。在没有缓存、预取和交错的情况下从文件读取数据所用的时间。
  • 其他数据读取或处理。其他不使用 tf.data 的输入相关操作所花费的时间。

图片

如需查看按执行时间细分的各个输入操作及其类别的统计信息，请展开“显示输入操作统计信息”部分。

系统将显示如下所示的源数据表：

图片

每个表条目包含以下信息：

  1. 输入操作。显示输入操作的 TensorFlow 操作名称。
  2. 计数。显示在性能剖析期间执行的操作的实例总数。
  3. 总时间（以毫秒为单位）。显示每个操作实例所用时间的累计总和。
  4. 总时间百分比。显示一个操作花费的总时间占输入处理总时间的比例。
  5. 总自用时间（以毫秒为单位）。显示其中每个实例所耗自用时间的累计总和。自用时间测量在函数体内部所用的时间，不包括它调用的函数所用的时间。例如，Iterator::PaddedBatch::Filter::ForeverRepeat::Map 由 Iterator::PaddedBatch::Filter 调用，因此前者的总自用时间不包含在后者的总自用时间内。
  6. 总自用时间百分比。显示总自用时间占输入处理总时间的比例。
  7. 类别。显示输入操作的处理类别。

Op Profile

Op Profile () 是一款 Cloud TPU 工具，它可展示在性能剖析期间执行的 XLA 操作的性能统计信息。Op Profile 会显示以下内容：

  • 您的应用使用 Cloud TPU 的程度，按类别和 TPU FLOPS 利用率显示各操作所用时间的百分比。
  • 最耗时的操作。这些操作可能有待优化。
  • 各个操作的详细信息，包括形状、填充和使用该操作的表达式。

您可以借助 Op Profile 找到可能需要进行优化的目标。例如，如果您的模型的 TPU 峰值 FLOPS 仅达到 5%，您可以使用此工具找出执行时间最长的 XLA 操作并查看它们使用了多少 TPU FLOPS。

使用 Op Profile

在性能剖析文件收集期间，capture_tpu_profile 还会创建一个包含 XLA 操作的性能统计信息的 op_profile.json 文件。

如需在 TensorBoard 中查看 op_profile 中的数据，可点击屏幕顶部的性能剖析 (Profile) 标签页，然后从工具 (Tools) 下拉列表中选择 op_profile。您将看到如下所示的界面：

图片

  1. 概览部分。显示 Cloud TPU 利用率并提供优化建议。
  2. 控制面板。通过包含的控件，您可设置表格中显示的操作数量、显示哪些操作，以及这些操作的排序方式。
  3. 操作表。列出与 XLA 操作相关联的主要 TensorFlow 操作类别的表。这些操作按 Cloud TPU 使用率百分比排序。
  4. 操作详细信息卡片。将鼠标悬停在表中的某个操作上时，会出现一张卡片，上面显示有关此操作的详细信息，其中包括 FLOPS 利用率、使用该操作的表达式以及操作布局（匹配）。
XLA 操作表

操作表按照 Cloud TPU 使用率百分比从高到低的顺序列出 XLA 操作类别。对于每个类别，该表都会默认显示花费时间的百分比、操作类别名称、关联的 TensorFlow 操作的名称及其 FLOPS 利用率百分比。如需显示（或隐藏）某个类别中 10 个最耗时的 XLA 操作，请点击表格中类别名称旁边的三角形。

图片

  1. 时间。显示该类别中所有操作所用时间的总百分比。您可以点击以展开条目，查看每个操作所用时间的明细。
  2. 10 大操作。点击类别名称旁边的切换开关可以显示或隐藏此类别中 10 个最耗时的操作。如果操作列表中显示了融合操作条目，您可以将其展开以查看其包含的非融合元素级操作。
  3. TensorFlow 操作。显示与 XLA 操作关联的 TensorFlow 操作名称。
  4. FLOPS。显示 FLOPS 利用率，即以 Cloud TPU 峰值 FLOPS 的百分比表示的 FLOPS 的测量值。FLOPS 利用率百分比越高，操作运行速度越快。表单元格采用颜色编码：绿色表示 FLOPS 利用率高（好），红色表示 FLOPS 利用率低（差）。
操作详细信息卡片

选择一个表条目后，左侧会出现一张卡片，上面显示有关 XLA 操作或该操作类别的详细信息。典型的卡片如下所示：

图片

  • 名称和类别。显示突出显示的 XLA 操作名称和类别。
  • FLOPS 利用率。显示 FLOPS 利用率占总 FLOPS 的百分比。
  • 表达式。显示包含该操作的 XLA 表达式。
  • 内存利用率。显示程序峰值内存使用率的百分比。
  • 布局（仅限卷积运算。）显示张量的形状和布局，包括张量的形状是否与矩阵单元完全匹配以及矩阵的填充方式。
解读结果

对于卷积运算，TPU FLOPS 利用率低可能是由以下原因之一或全部造成的：

  • 填充（矩阵单元仅被部分使用）
  • 卷积运算受限于内存

本部分简要解释另一个 FLOPS 较低的模型中的某些数字。在这个例子中，输出融合和卷积占用了大部分执行时间，同时还有一个 FLOPS 非常低的矢量或标量操作的长尾。

此类性能剖析文件的一种优化策略是将矢量或标量操作转换为卷积运算。

在以下示例中，%convolution.399 显示的 FLOPS 利用率和内存利用率低于上例中的 %convolution.340。

图片

仔细观察布局可发现，批次大小 16 被填充至 128，并且特征大小 3 被填充至 8，这表明只有 5% 的矩阵单元得到有效利用。（此实例利用率百分比的计算方法是 ((batch_time * num_of_features) / padding_size ) / num_of_cores）。将此示例中的 FLOPS 与前例中与矩阵完全匹配的 %convolution.340 进行比较。

Pod Viewer

Pod Viewer 工具为 Pod 中的每个核心提供性能可视化，并显示 Pod 中核心之间的通信通道的状态。Pod Viewer 可以识别并突出显示潜在的瓶颈和需要优化的区域。此工具适用于完整 Pod 和所有 v2 和 v3 Pod 切片。

如需显示 Pod Viewer 工具，请执行以下操作：

  1. 从 TensorBoard 窗口右上角的菜单按钮中选择性能剖析 (Profile)。
  2. 点击窗口左侧的工具 (Tools) 菜单，然后选择 pod_viewer。

Pod Viewer 用户界面包括以下内容：

  1. 步滑块，用于选择要检查的步。
  2. 拓扑图，以交互方式直观显示整个 TPU 系统中的 TPU 核心。
  3. 通信链接图表，用于直观显示拓扑图中的发送和接收 (recv) 通道。
  4. 发送和接收通道延迟时间的条形图。将鼠标悬停在此图表中的某个柱形上会激活通信链接图表中的通信链接。左侧栏上会显示一个通道详细信息卡片，上面显示该通道的详细信息，例如传输的数据的大小、延迟时间和带宽。
  5. 步明细图表，直观显示所有核心的步明细信息。此图表可用于跟踪系统瓶颈以及确定特定核心是否降低系统速度。

图片

步滑块

使用滑块选择一个步。该工具的其余部分显示该步的统计信息，例如步明细和通信链接。

拓扑图

拓扑图按主机、芯片和核心以分层方式进行组织。最小的矩形是 TPU 核心。合并在一起的两个核心表示一个 TPU 芯片，四个芯片组合表示一个主机。

图片

拓扑图也是热图，其颜色按特定明细（例如高 FLOPS 计算、馈入、发送等）在所选步中占用时间的百分比进行标示。拓扑图下方的柱形（如下图所示）展示了核心和芯片使用率的颜色标示方式。核心的颜色从黄色向蓝色过渡，以显示不同利用率。对于高 FLOPS 计算，数字越大（颜色越深）表示计算需要花费的时间越长。对于所有其他明细，数字越小（颜色越浅）表示等待时间越短。当某个核心比其他核心颜色更暗时，说明其为潜在问题区域或热点。

点击系统名称旁边的下拉菜单选择器（图中已圈出），选择要检查的明细的特定类型。

将鼠标悬停在任何小矩形（单个核心）上，会出现一个信息提示框，其上显示核心在系统中的位置、其全局芯片 ID 和主机名。信息提示框还显示了所选明细类别（例如高 FLOPS）的时长，及其在步中的利用率百分比。

图片

通信通道

如果您的模型使用发送和接收链接在核心之间进行通信，则此工具有助于直观呈现这些链接。当您的模型包含发送和接收操作时，您可以使用通道 ID 选择器来选择通道 ID。来自源 (src) 核心和目标 (dst) 核心的链接表示通信通道。将鼠标悬停在图表的柱形上，可以在拓扑图上呈现通信通道，以显示发送和接收通道的延迟时间。

图片

左侧栏会出现一张卡片，上面显示通信通道的相关详细信息。典型的卡片如下所示：

图片

  1. 已传输的数据 (Data Transferred)，显示发送和接收通道已传输的数据量，以兆比字节 (MiB) 为单位。
  2. 延迟时间 (Latency)，显示从发送事件开始到接收完成事件结束的时长，以微秒为单位。
  3. BW，显示在运行时长内从源核心到目标核心传输的数据量，以吉比字节 (GiB) 为单位。
  4. 发送延迟 (Send Delay)，指从接收完成开始到发送开始的时长，以微秒为单位。如果接收完成操作在发送操作开始之后开始，则延迟为零。
  5. HLO 名称 (Hlo Names)，显示与此通道关联的 XLA HLO 操作名称。这些 HLO 名称与其他 TensorBoard 工具（如 op_profile 和 memory_viewer）中显示的统计信息相关联。

步明细图表

此图表提供了每个训练或评估步的详细信息。

x 轴表示全局芯片 ID，y 轴表示时间（以微秒为单位）。在此图表中，您可以看到在特定训练步中时间用于何处、何处存在瓶颈，以及所有芯片之间是否存在负载不平衡。

图片

左侧栏会出现一张卡片，上面显示有关步明细的详细信息。典型的卡片如下所示：

图片

卡片中的字段展示以下内容：

  1. 高 FLOPS 计算 (High Flops Compute)，即卷积运算或输出融合操作 (ops) 所花费的时间。
  2. 低 FLOPS 计算 (Low flops compute)，通过从总时长中减去所有其他明细来计算。
  3. 馈入 (Infeed)，这是 TPU 等待主机的时间。
  4. 馈出 (Outfeed)，即主机等待 TPU 输出的时间。
  5. AllReduce 同步 (AllReduce sync)，这是等待与其他核心同步的 CrossReplicaSum 操作所花费的时间。CrossReplicaSum 操作计算各个副本的总和。
  6. AllReduce 计算 (AllReduce compute)，这是 CrossReplicaSum 操作实际花费的计算时间。
  7. 芯片到芯片发送操作 (Chip to chip send ops)，这是发送操作所花费的时间。
  8. 芯片到芯片接收完成操作 (Chip to chip recv-done ops)，这是接收完成操作所花费的时间。

Trace viewer

Trace viewer 是性能剖析 (Profile) 下的 Cloud TPU 性能分析工具。此工具使用 Chrome 跟踪记录事件分析查看器，因此仅适用于 Chrome 浏览器。

Trace Viewer 包含一个时间轴，其中显示的内容包括：

  • 由 TensorFlow 模型执行的操作的时长。
  • 系统的哪个部分（TPU 还是主机）执行了操作。通常，主机执行馈入操作，它会预处理训练数据并将其传输到 TPU，而 TPU 执行实际的模型训练。

您可以通过 Trace Viewer 找出模型中的性能问题，然后采取措施来加以解决。例如，从总括层面来讲，您可以确定占用大部分时间的是馈入还是模型训练。展开细目，您可以确定哪些 TensorFlow 操作执行时间最长。

请注意，Trace Viewer 针对每个 Cloud TPU 仅限访问 100 万个事件。如果您需要访问更多事件，请使用 Streaming Trace Viewer。

Trace Viewer 界面

如需打开 Trace Viewer，请转到 TensorBoard，点击屏幕顶部的性能剖析 (Profile) 标签页，然后从工具下拉列表中选择 trace_viewer。Trace Viewer 会显示最近的运行：

图片

此屏幕包含以下主要元素（标有上述数字）：

  1. 运行下拉列表。包含已捕获其跟踪记录信息的所有运行。默认视图为最近的运行，您也可以打开下拉列表选择其他运行。
  2. 工具下拉列表。选择不同的性能剖析工具。
  3. 主机下拉列表。选择包含 Cloud TPU 集的主机。
  4. 时间轴窗格。显示 Cloud TPU 和主机随时间执行操作的情况。
  5. 详细信息窗格。显示时间轴窗格中所选操作的更多信息。

以下是时间轴窗格的详细信息：

图片

时间轴窗格包含以下元素：

  1. 顶栏。包含各种辅助控件。
  2. 时间轴。显示相对于跟踪记录开始时的时间。
  3. 版块和跟踪记录标签。每个版块包含多个跟踪记录，左侧有一个三角形，点击此三角形可展开和收起该版块。系统中的每个处理元素都有一个版块。
  4. 工具选择器。包含与 Trace Viewer 交互的各种工具。
  5. 事件。这些事件显示操作的执行时间或元事件（例如训练步）的时长。
  6. 垂直标签栏。它对 Cloud TPU 没有用处。此栏是 Chrome 提供的通用跟踪记录查看器工具的一部分，用于各种性能分析任务。
版块和跟踪记录

Trace Viewer 包含以下版块：

  • 每个 TPU 节点有一个版块，上面标有 TPU 芯片的编号及芯片内 TPU 节点的编号（例如“芯片 2: TPU 核心 1”）。每个 TPU 节点版块包含以下跟踪记录：
    • 步。显示 TPU 上运行的训练步的时长。
    • TensorFlow 操作。显示在 TPU 上执行的 TensorFlow 操作。
    • XLA 操作。显示在 TPU 上运行的 XLA 操作。（每个操作都会转换为一个或多个 XLA 操作。XLA 编译器会将这些 XLA 操作转换为在 TPU 上运行的代码。）
  • 针对在主机的 CPU 上运行的线程的版块，其标签为“Host Threads”。对于每个 CPU 线程，此版块都包含一个跟踪记录。注意：您可以忽略版块标签旁边显示的信息。
时间轴工具选择器

您可以使用 TensorBoard 中的时间轴工具选择器与时间轴视图进行交互。您可以点击时间轴工具或使用以下键盘快捷键激活并突出显示该工具。如需移动时间轴工具选择器，请点击顶部的虚线区域，然后将选择器拖动到所需位置。

以下是各个时间轴工具的使用方法：

 选择工具
点击事件可将其选中，拖动可选择多个事件。有关所选事件（名称、开始时间和时长）的更多信息将显示在“详细信息”窗格中。

平移工具
拖动可在水平和垂直方向平移时间轴视图。

缩放工具
沿水平（时间）轴方向，向上拖动可放大视图，向下拖动可缩小视图。鼠标光标的水平位置决定缩放操作的中心。

注意：缩放工具有一个已知错误：当鼠标光标在时间轴视图外面时松开鼠标按钮，仍会进行缩放。如果您遇到这种情况，点击时间轴视图即可停止缩放。

计时工具
水平拖动可以标记时间间隔。间隔的时长会显示在时间轴上。如需调整间隔，请拖动其末端。如需清除间隔，可点击时间轴视图内的任意位置。

请注意，如果您选择其他某个工具，间隔标记仍会保留。
事件

时间轴内的事件以不同的颜色显示；颜色本身并无特定的含义。

时间轴顶栏

时间轴窗格的顶栏包含多个辅助控件：

图片

  1. 元数据显示区。不用于 TPU。
  2. 查看选项。不用于 TPU。
  3. 搜索框。输入文本可搜索名称中包含该文本的所有事件。点击搜索框右侧的箭头按钮可在匹配的事件中前后移动，从而依次选择每个事件。
  4. 控制台按钮。不用于 TPU。
  5. 帮助按钮。点击可显示帮助摘要。
键盘快捷键

以下是您可以在 Trace Viewer 中使用的键盘快捷键。点击顶栏的帮助 (?) 按钮可查看更多键盘快捷键。

    w Zoom in
    s Zoom out
    a Pan left
    d Pan right
    f Zoom to selected event(s)
    m Mark time interval for selected event(s)
    1 Activate selection tool
    2 Activate pan tool
    3 Activate zoom tool
    4 Activate timing tool

f 快捷键非常有用。选择一个步，然后按 f 可快速对其进行缩放。

典型事件

以下是在分析 TPU 性能时非常有用的一些事件类型。

图片

  • InfeedDequeueTuple。此 TensorFlow 操作在 TPU 上运行，并接收来自主机的输入数据。如果馈入需要很长时间，则可能意味着在主机上预处理数据的 TensorFlow 操作无法跟上 TPU 使用数据的速度。您可以在名为 InfeedEnqueueTuple 的主机跟踪记录中查看相应的事件。如需查看更详细的输入流水线分析，请使用 Input Pipeline Analyzer 工具。

  • CrossReplicaSum。此 TensorFlow 操作在 TPU 上运行，并计算各个副本的总和。因为每个副本都对应于不同的 TPU 节点，所以此操作必须等待所有 TPU 节点完成一步。如果此操作花费了大量时间，这通常并不意味着求和操作本身很慢，而是 TPU 节点正在等待其他某个数据馈入较慢的 TPU 节点。

图片

  • 数据集操作。使用 Dataset API 加载数据时，Trace Viewer 会直观显示所执行的数据集操作。示例中的 Iterator::Filter::Batch::ForeverRepeat::Memory 经过编译并对应于 dataset.map() 操作。使用 Trace Viewer 检查这些加载操作有助于调试和缓解输入流水线瓶颈。

图片

  • 预取线程。使用 dataset.prefetch() 缓冲输入数据可防止因文件访问中偶尔出现的速度缓慢问题而导致输入流水线瓶颈。
常见错误

以下是使用 Trace Viewer 时需要注意的一些潜在问题：

  • 事件显示限制。Trace Viewer 最多可显示 100 万个事件。如果您捕获了更多事件，则只显示最早的 100 万个事件；以后的事件将被丢弃。如需捕获更多 TPU 事件，您可以使用 --include_dataset_ops=False 标志明确要求 capture_tpu_profile 排除该数据集操作。
  • 非常长的事件。如果事件开始于捕获开始之前或结束于捕获完成之后，则不会显示在 Trace Viewer 中。因此，非常长的事件可能会丢失。

  • 何时开始捕获跟踪记录。请务必在确定 Cloud TPU 正在运行后，再开始捕获跟踪记录。如果过早开始，您可能只会在 Trace Viewer 中看到少量事件或根本看不到任何事件。您可以使用 --duration_ms 标志增加分析时间，也可以使用 --num_tracing_attempts 标志设置自动重试。例如：

      (vm)$ capture_tpu_profile --tpu=$TPU_NAME
    --logdir=${MODEL_DIR} --duration_ms=60000 --num_tracing_attempts=10

Memory Viewer

借助 Memory Viewer，您可以直观地了解程序的峰值内存使用量，以及程序生命周期内的内存使用趋势。

Memory Viewer 界面如下所示：

图片

  1. 主机下拉列表。用于选择要直观显示的 TPU 主机和 XLA 高级优化器 (HLO) 模块。
  2. 内存概览。显示峰值内存分配和无填充的大小。
  3. 工作空间图表。显示峰值内存使用量并绘制程序生命周期内的内存使用趋势。将鼠标悬停在缓冲区图表中的缓冲区上可以为缓冲区生命周期和缓冲区详细信息卡片添加注释。
  4. 缓冲区图表。这两个图表显示峰值内存使用点（如工作空间图中的垂直线所示）处的缓冲区分配。将鼠标悬停在其中一个缓冲区图表中的缓冲区上可以在工作空间图表中显示指示缓冲区生命周期的柱形，并在左侧显示详细信息卡片。
  5. 缓冲区分配详细信息卡片。显示有关缓冲区分配的详细信息。
内存概览面板

内存概览（顶部）面板显示模块名称，以及总缓冲区分配大小达到最大值时的峰值内存分配设置。此面板还显示未填充的峰值分配大小以供比较。

图片

工作空间图表

此图表显示峰值内存使用量并绘制程序生命周期内的内存使用趋势。图表中绘制的从顶部到底部的线表示程序的峰值内存利用率。此点决定了程序能否加载到可用的全局存储空间中。

图片

上图中的每个点代表 XLA 的 HLO 程序中的“程序点”，因为它已由编译器调度。折线显示接近和远离峰值使用量的“尖锐”程度。

与缓冲区图表元素的交互

将鼠标悬停在工作空间图表下方其中一个缓冲区图表中显示的缓冲区上时，工作空间图表将显示该缓冲区的水平生命周期线。水平线与突出显示的缓冲区颜色相同。

图片

水平线的粗细程度表示缓冲区大小与峰值内存分配的相对量级。水平线的长度对应于缓冲区的生命周期，左边界为程序中缓冲区空间的分配点，右边界为此空间的释放点。

缓冲区图表

这两个图表显示了峰值使用点（由此图表上方工作空间图中的垂直线表示）处的内存使用量明细。

图片

  • 按程序顺序显示。按程序执行期间缓冲区的活动顺序从左到右显示缓冲区。活动时间最长的缓冲区位于此图表的左侧。

  • 按大小显示 (By Size)。按大小顺序以降序显示程序执行期间的活动缓冲区。在峰值内存使用点处影响最大的缓冲区位于左侧。

缓冲区分配详细信息卡片

将鼠标悬停在其中一个缓冲区图表中显示的缓冲区上时，此工具将显示缓冲区分配详细信息卡片（除了工作图表中显示的生命周期线）。典型的详细信息卡片如下所示：

图片

  1. 名称。XLA 操作的名称。
  2. 类别。操作类别。
  3. 大小。缓冲区分配的大小（包括填充）。
  4. 未填充大小。未填充的缓冲区分配的大小。
  5. 膨胀度。填充的缓冲区大小与未填充的缓冲区大小的相对量级。
  6. 额外的内存。指示填充额外使用了多少内存量。
  7. 形状。描述 N 维数组的阶、大小和数据类型。
  8. TensorFlow 操作名称。显示与缓冲区分配关联的 TensorFlow 操作名称。
  9. 分配类型。指示缓冲区分配类别。类型包括：参数、输出、线程局部和临时（例如某个融合中的缓冲区分配）。
“内存不足”错误

如果运行模型时出现“内存不足错误”，请使用以下命令捕获内存性能剖析文件并在 Memory Viewer 中进行查看。确保设置适当的 duration_ms，以使性能剖析周期与程序编译时间重叠。此命令的输出可以帮助您了解出现错误的原因：

  (vm)$ capture_tpu_profile --tpu=$TPU_NAME --logdir=${MODEL_DIR} --duration_ms=60000

Streaming Trace Viewer

Streaming Trace Viewer (trace_viewer) 是一款 Cloud TPU 性能分析工具，可用于 TensorFlow 2.17.0 或更高版本，它能够动态呈现跟踪记录。此工具使用 Chrome 跟踪记录事件分析查看器，因此仅适用于 Chrome 浏览器。

使用 capture_tpu_profile 捕获性能剖析文件时，.tracetable 文件将保存到您的 Google Cloud Storage 存储分区中。该文件包含大量跟踪记录事件，您可以在 Trace Viewer 和 Streaming Trace Viewer 中查看这些事件。

使用 Streaming Trace Viewer

如需使用 Streaming Trace Viewer (trace_viewer)，您必须关闭现有的 TensorBoard 会话，然后使用要检查的 TPU 的 IP 地址重新启动 TensorBoard。Streaming Trace Viewer 要求 TensorBoard 向 Cloud TPU 的 IP 地址发出 Google 远程过程调用 (GRPC)。GRPC 通道未加密。

您可以在 Cloud TPU 页面上找到 Cloud TPU 主机的 IP 地址。查找您的 Cloud TPU，然后在内部 IP 列中查找 IP 地址。

在您的虚拟机中，按照以下方式运行 TensorBoard，将 tpu-ip 替换为 TPU 的 IP 地址：

  (vm)$ tensorboard --logdir=${MODEL_DIR} \
    --master_tpu_unsecure_channel=tpu-ip

TensorBoard 工具会显示在工具 (Tools) 下拉列表中。

图片

在时间轴中，您可以放大和缩小以查看动态加载到浏览器中的跟踪记录事件。

图片

监控 Cloud TPU 作业

本部分介绍如何使用 capture_tpu_profile 捕获单个性能剖析文件或在命令行界面上实时持续监控 Cloud TPU 作业。将 --monitoring_level 选项设置为 0（默认值）、1 或 2，可以分别实现单个性能剖析文件、基本监控数据或详细监控数据。

打开一个新的 Cloud Shell 并使用 ssh 连接到您的虚拟机（将命令中的 vm-name 替换为您的虚拟机名称）：

  (vm)$ gcloud compute ssh vm-name \
  --ssh-flag=-L6006:localhost:6006

在新的 Cloud Shell 中，运行 capture_tpu_profile 并将 --monitoring_level 标志设置为 1 或 2，例如：

  (vm)$ capture_tpu_profile --tpu=$TPU_NAME \
   --monitoring_level=1

设置 monitoring_level=1 会生成入如下所示的输出：

    TPU type: TPU v2
    Utilization of TPU Matrix Units is (higher is better): 10.7%

设置 monitoring_level=2 会显示更详细的信息：

    TPU type: TPU v2
    Number of TPU Cores: 8
    TPU idle time (lower is better): 0.091%
    Utilization of TPU Matrix Units is (higher is better): 10.7%
    Step time: 1.95 kms (avg), 1.90kms (minute), 2.00 kms (max)
    Infeed percentage: 87.5% (avg). 87.2% (min), 87.8 (max)

监控标志

  • --tpu（必需）指定要监控的 Cloud TPU 的名称。
  • --monitoring_level。将 capture_tpu_profile 的行为从生成单个性能剖析文件更改为生成基本或详细的持续监控数据。此标志有以下三个等级： 第 0 级（默认）：生成一个性能剖析文件，然后退出。第 1 级：显示 TPU 版本和 TPU 利用率。第 2 级：显示 TPU 利用率、TPU 空闲时间和使用的 TPU 核心数量。同时提供单步用时的最小值、平均值和最大值，以及馈入百分比贡献。
  • --duration_ms（可选；默认值为 1000ms）指定在每个周期内分析 TPU 主机的时间。通常，这个时长应足够长，以便捕获至少一个训练步的数据。在大多数模型中，1 秒可捕获一个训练步，但如果模型单步用时非常大，则可以将此值设置为 step_time 的 2 倍（以毫秒为单位）。
  • --num_queries 指定运行 capture_tpu_profile 的周期数。如需持续监控 TPU 作业，请将此值设置为较大的数字。如需快速检查模型的单步用时，请将此值设置为较小的数字。