使用 Cloud TPU 工具

本文档介绍如何通过以下方式设置和运行性能剖析工具:捕获性能剖析文件,再在 TensorFlow 的 TensorBoard 控制台中用它来识别和分析 Cloud TPU 上的程序性能。本文档还介绍了如何在命令行上持续监控 TPU 作业(请参阅监控作业)。

前提条件

在使用本指南中介绍的 Cloud TPU 性能剖析工具之前,您必须完成以下操作:

使用 TensorBoard 中的 Cloud TPU 工具

TensorBoard 是一套专用于直观地呈现 TensorFlow 数据的工具。我们提供了一组 Cloud TPU 性能剖析工具,您可以在安装 Cloud TPU 分析器插件后从 TensorBoard 访问这些工具。此插件可直观呈现各种规模的 Cloud TPU 节点的性能。

只有从正在运行的 TensorFlow 模型中收集跟踪记录信息之后,才能在 TensorBoard 菜单栏上的性能剖析 (Profile) 标签页下使用 Cloud TPU 工具选择器。

下述部分介绍了如何设置计算环境、如何运行模型和捕获 Cloud TPU 性能剖析文件,以及如何从虚拟机命令行运行 TensorBoard 来使用这些工具。如需详细了解如何通过代码调用 TensorBoard,请参阅 TensorBoard 编程指南

创建 Cloud TPU 资源

Cloud TPU 的快速入门说明介绍了如何创建 Compute Engine 虚拟机和 Cloud TPU 资源。

如果您打算在创建资源后立即执行本指南中的步骤,请不要执行这些说明中的清理部分。运行模型及使用资源结束后,请按照清理说明中的步骤操作,以避免产生不必要的费用。

运行 cloud_tpu_profiler

运行 cloud-tpu-profiler 1.15 可为 TensorBoard 提供 Cloud TPU 性能剖析插件和一个脚本 capture-tpu-profile。您可运行该脚本以捕获可在 TensorBoard 中查看的性能剖析文件,或在命令行上监控 TPU 作业(请参阅监控作业)。

在 TensorBoard 中,您需要先运行模型,然后在模型运行期间运行捕获性能剖析(跟踪记录)信息性能剖析 (Profile) 标签页才会显示。

若要查看分析器版本,请使用 pip。如果您没有最新版本,请使用第二个命令进行安装:

(vm)$ pip freeze | grep cloud-tpu-profiler
    (vm)$ pip install --upgrade "cloud-tpu-profiler>=1.15"
    

您还必须按如下方式设置 PATH 环境变量:

(vm)$ export PATH="$PATH:`python -m site --user-base`/bin"
    

关于 capture_tpu_profile

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

您可以捕获性能剖析文件或跟踪记录数据,方法是运行模型、执行 capture_tpu_profile,然后在模型停止运行之前启动 TensorBoard。例如:

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

这些参数是在运行模型时设置的,具有以下含义:

  • --tpu=$TPU_NAME - 这是您在创建 Cloud TPU 时指定的名称。如果您使用 ctpu up 启动 Compute Engine 虚拟机和 Cloud TPU,则此参数默认为您的用户名。如果您使用 gcloud 或 Cloud Console 设置虚拟机和 TPU,则此参数为您创建虚拟机时指定的名称。
  • --logdir=${MODEL_DIR} - 这是存储模型和检查点的 Cloud Storage 位置。模型目录通常按如下方式进行定义
        export STORAGE_BUCKET=gs://[YOUR-BUCKET-NAME]
        export MODEL_DIR=${STORAGE_BUCKET}/OUTPUT DIRECTORY
    例如,如果您运行的是 MNIST 模型,则输出目录可能会定义为:
        export STORAGE_BUCKET=gs://[YOUR-BUCKET-NAME]
        export MODEL_DIR=${STORAGE_BUCKET}/mnist

模型运行时,Cloud TPU 服务帐号会创建一个目录(对象)并将数据写入您的 Compute Engine 存储分区。在运行模型之前,您必须在存储分区上为该帐号设置权限。

默认情况下,capture_tpu_profile 会捕获 2 秒跟踪记录。您可以使用 duration_ms 命令行选项设置跟踪记录时长,也可当运行模型时在程序中进行设置。

捕获性能剖析文件

此部分中的步骤介绍如何运行模型、执行 capture_tpu_profile,然后在模型停止运行之前启动 TensorBoard 以捕获性能剖析文件。如果按此部分所述的步骤运行 TensorBoard,则可访问除 Streaming Trace Viewer 之外的所有性能剖析工具。

如果您希望在命令行上监控 TPU 作业,请参阅监控作业

在此示例中,MNIST 教程用作模型。

  1. 运行 ctpu up。

  2. 转到 Cloud Console 信息中心(首页),然后选择 Compute Engine > TPU。

  3. 点击您创建的 Cloud TPU。

  4. 找到 Cloud TPU 服务帐号名称并复制,例如:

    service-11111111118@cloud-tpu.iam.myserviceaccount.com

  5. 在信息中心的“资源”面板中,点击“Cloud Storage”。系统将显示存储分区列表,或者提示您创建存储分区。如果尚未创建存储分区,请点击页面顶部的“创建存储分区”,然后在 TPU 所在的区域中创建一个存储分区。

  6. 点击要使用的存储分区的复选框。确保该存储分区与您创建的 Cloud TPU 位于同一区域。

  7. 在列表中选中存储分区后,选择显示信息面板,然后选择修改存储分区权限

  8. 将 TPU 服务帐号名称粘贴到该存储分区的添加成员字段中

  9. 选择 Storage Legacy,然后选中 Storage Legacy Bucket Reader、Writer 和 Owner 权限:

    图片

  10. 在虚拟机 shell 中,使用 pip 检查您的 TensorBoard 版本。

    (vm)$ pip freeze | grep tensorboard
        
  11. 如果 TensorBoard 版本低于 2.1,则升级到最新版本。

    (vm)$ pip install --upgrade -U "tensorboard>=2.1"    
    (vm)$ pip freeze | grep tensorflow
        (vm)$ pip install --upgrade -U "tensorflow>=2.1"
        
  12. 如果您尚未这样做,请设置 PATH 环境变量:

    (vm)$ export PATH="$PATH:`python -m site --user-base`/bin" 
  13. 按照 MNIST 教程设置并执行 MNIST 训练作业,例如:

        (vm)$ python /usr/share/models/official/mnist/mnist_tpu.py \
          --tpu=$TPU_NAME \
          --data_dir=${STORAGE_BUCKET}/data \
          --model_dir=${MODEL_DIR} \
          --use_tpu=True \
          --iterations=500 \
          --train_steps=5000 
  14. 在作业运行期间,打开一个新的 Cloud Shell 并运行 ctpu up;如果您使用的是非默认的 TPU 名称,请指定该名称。登录到虚拟机后,按前面的步骤设置以下环境变量:STORAGE_BUCKET、MODEL_DIR。

  15. 在作业运行期间,打开第三个 Cloud Shell 并使用 ssh 连接到您的虚拟机(将命令中的 $vm 替换为您的虚拟机名称)。

    gcloud compute ssh $vm --ssh-flag=-L6006:localhost:6006 

    登录后,按前面的步骤设置以下环境变量:TPU_NAME、STORAGE_BUCKET、MODEL_DIR。

  16. 在第二个 Cloud Shell 中,运行 capture_tpu_profile 以捕获 .tracetable 文件

    (vm)$ capture_tpu_profile --tpu=$TPU_NAME --logdir=${MODEL_DIR} 
  17. 在第三个 Cloud Shell 中,运行 TensorBoard 并将其指向模型目录:

    (vm)$ tensorboard --logdir=${MODEL_DIR} & 
  18. 点击第二个 Cloud Shell 中的网页预览按钮,并打开端口 8080 以查看 TensorBoard 输出。

TensorBoard 为模型及其性能提供了大量可视化或图形表达。您可以结合使用图和性能剖析工具,以优化模型并提升其在 Cloud TPU 上的性能。

XLA 结构图

在模型编译期间,TensorFlow 会在模型运行之前生成将在 Cloud TPU 上运行的 (XLA) 图。图中的数据存储在 MODEL_DIR 目录中。您无需运行 capture_tpu_profile 即可查看此图。

如需查看模型的 XLA 结构图,请选择 TensorBoard 中的 (Graphs) 标签页。 颜色 (Color) 的默认选项为结构 (Structure)。

图片

结构图中的单个节点表示 XLA 指令。例如,对于名为 x/y/z 的 TensorFlow 添加操作,与其对应的 XLA 在图中显示为 x/y/z/add

XLA 图显示了有关 Cloud TPU 如何执行特定模型的信息。该图还提供不同操作的输入和输出的形状。捕获模型的性能剖析文件后,您可以使用 XLA 图以及 Trace ViewerStreaming Trace Viewer 来深入了解哪些地方耗时最多。

注意:

  • 某些节点没有 TensorFlow 命名空间,因为并非所有 XLA 指令都有相应的 TensorFlow 操作(例如由 XLA 编译器注入的指令)。
  • TensorFlow 程序结构会尽可能包含在 XLA 图中。但是,由于在 Cloud TPU 上运行的 XLA 程序经过高度优化,其图结构可能与原始 TensorFlow 程序差别很大。
  • 名为融合 (fusion) 的特殊 XLA 指令可以将来自不同 TensorFlow 操作的多个指令合并为单个计算。与融合中的根指令对应的 TensorFlow 操作用作融合操作的命名空间。

TPU 兼容性图

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

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

图片

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

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

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

图片

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

前提条件
  • 设置 tf.estimator APITPUEstimatormodel_dir 属性,配置模型以将模型图数据写入文件。
  • 对于要在 Cloud TPU 中运行的操作,请移除代码中所有手动分配给 GPU 或 CPU 的部分。当结合使用 TPU 兼容性 (TPU Compatibility) 选项时,兼容性检查器会跳过明确分配给非 TPU 的所有操作。
解读兼容性结果

下面是 Abalone 模型的兼容性摘要的示例图,其中显示了多个不兼容的操作:

图片

此模型未指定手动设备配置,因此系统检查了所有操作,甚至包括那些应始终在 CPU 上运行的操作,例如显示的“save”和“report_uninitialized_variables”操作。

root_mean_squared_error 中的两个 AssignAdd 操作可能会引起问题。

您可以从源代码中看到,root_mean_squared_error 仅用作额外的评估指标:

        # Calculate root mean squared error as additional eval metric
        eval_metric_ops = {
            "rmse": tf.metrics.root_mean_squared_error(
                tf.cast(labels, tf.float64), predictions)
        }
    

除非此操作出现在训练循环内,否则它通常在 CPU 上运行,因此可以忽略错误报告。综上所述,该模型已可在 Cloud TPU 上运行。

资料

性能剖析 (Profile) 标签页在您运行 capture_tpu_profile 时创建,它仅在您捕获某些模型数据后才在 TensorBoard 中显示。数据可用后,点击性能剖析 (Profile) 标签页会显示一系列有助于性能分析的工具:

性能剖析概览页面

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

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

图片

  • 性能摘要

    • 所有采样步的平均单步用时
    • 主机空闲时间百分比
    • TPU 空闲时间百分比
    • TPU 矩阵单元利用率百分比
  • 单步用时图。展示所有采样步的设备单步用时图(以毫秒为单位)。蓝色区域对应的是 TPU 空闲时等待主机输入数据的单步用时部分。橙色区域展示 Cloud TPU 实际运行的时间。

  • TPU 上最耗时的 10 个 TensorFlow 操作。展示占用了最多时间的 TensorFlow 操作。点击显示表格 (Show table) 按钮将显示如下表格:

    图片

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

  • 运行环境

    • 使用的主机数
    • 使用的 TPU 类型
    • TPU 核心数量
    • 训练批次大小
  • 后续步骤建议。报告模型输入是否受限,以及是否出现 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. 主机端分析展示主机端的详细分析,包括主机上输入处理时间的明细。
输入流水线摘要

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

图片

设备端分析

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

图片

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

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

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

图片

如需查看按执行时间细分的各个输入操作及其类别的统计信息,请点击“显示输入操作统计信息”(Show Input Op Statistics) 按钮。

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

图片

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

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

Op Profile

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. 控制面板。包含一个设置滑块和一个切换开关,前者用于控制在操作表中显示的操作数,后者用于设置操作表,使其列出总执行时长前 90% 以内的操作。
  3. 操作表。按总时长列出与 XLA 操作相关联的主要 TensorFlow 操作类别,其中时间以 Cloud TPU 使用率百分比的形式表示,且前提是该类别中的所有操作都要执行。
  4. 操作详细信息卡片。将鼠标悬停在表中的某个操作上时,会出现一张卡片,上面显示有关此操作的详细信息,包括 FLOPS 利用率、使用该操作的表达式以及操作布局(匹配)。
XLA 操作表

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

图片

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

将鼠标悬停在表条目上时,左侧会显示一张卡片,上面显示有关 XLA 操作或操作类别的详细信息。典型的卡片如下所示:

图片

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

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

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

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

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

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

图片

仔细观察布局可发现,批次大小 16 被填充至 128,并且特征大小 3 被填充至 8,这表明只有 5% 的矩阵单元得到有效利用。(此例中利用率百分比的计算方法是将批次大小乘以特征大小,即 16 乘 3,然后用结果除以填充值,即先除以 128,再除以 8。)将此示例中的 FLOPS 与前例中的 %convolution.340 进行比较,后者与矩阵完全匹配。

Pod Viewer

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

如需显示 Pod Viewer 工具,请执行以下操作:

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

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 显示最近的运行:

图片

此屏幕包含以下主要元素(标有上述数字):

  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 操作。(每个 TensorFlow 操作都会转换为一个或多个 XLA 操作。XLA 编译器会将这些 XLA 操作转换为在 TPU 上运行的代码。)
  • 显示在主机的 CPU 上运行的线程的版块,上面标有“主机线程”。对于每个 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 1.14 或更高版本,它能够动态呈现跟踪记录。此工具使用 Chrome 跟踪记录事件分析查看器,因此仅适用于 Chrome 浏览器。

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

使用 Streaming Trace Viewer

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

要在 Cloud Console 上查找 Cloud TPU 主机的 IP 地址,请打开 TPU 页面并查看显示的表格,了解您要查看其跟踪记录的 Cloud TPU 的名称。

图片

每个 Cloud TPU 的内部 IP (Internal IP) 列包含一个 IP 地址 [TPU_IP]

图片

在您的虚拟机中,使用以下命令运行 TensorBoard:

(vm)$ tensorboard --logdir=${MODEL_DIR} --master_tpu_unsecure_channel=[TPU_IP]
    

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

图片

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

图片

监控 Cloud TPU 作业

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

  1. 打开一个新的 Cloud Shell 并使用 ssh 连接到您的虚拟机(将命令中的 $vm 替换为您的虚拟机名称):

    gcloud compute ssh $vm --ssh-flag=-L6006:localhost:6006
        
  2. 在新的 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 作业,请将此值设置为较大的数字。 如需快速检查模型的单步用时,请将此值设置为较小的数字。