Trace Viewer 工具

您可以使用 Trace Viewer 直观地查看性能分析会话期间发生的事件的时间轴。它会显示模型在系统的不同部分（例如主机 [CPU] 和加速器 [GPU 或 TPU]）上执行的操作的持续时间。这样，您就可以了解模型如何利用硬件资源、找出性能瓶颈，以及优化模型以加快执行速度。Trace Viewer 界面基于 chrome://tracing 中使用的界面，因此需要您使用 Chrome 浏览器。

支持的平台

TPU：受支持

GPU：受支持

GPU 的 Trace Viewer 会按 GPU 芯片和其中的流对时间轴上的事件进行分组，而对于 TPU，则会按 TPU 芯片和核心对时间轴上的事件进行分组。无论是哪种情况，Trace Viewer 还会在时间轴上显示与加速器连接的主机 CPU 的每个线程的事件。

与时间轴互动

轨迹查看器提供了多种用于浏览和检查时间轴的工具和方法：

  • 导航：您可以使用以下键盘快捷键：
    • W：放大。
    • S：缩小。
    • 答：向左平移。
    • D：向右平移。
  • 工具选择器：您可以点击工具选择器中的工具，也可以使用相应的键盘快捷键：
    • 选择工具 (1 或 !)：点击事件可将其选中，并在“Details”窗格中查看其详细信息。按住 Ctrl 键并点击多个事件可查看摘要。
    • 平移工具 (2 或 @)：拖动可水平或垂直移动时间轴。
    • 缩放工具（3 或 #）：拖动可放大时间轴的特定区域。
    • 计时工具（4 或 $）：拖动以标记时间间隔。系统会显示所标记时间段的持续时间。您还可以使用“m”键标记所选内容并确定其总时长。
  • 缩放到所选事件 (f)：选择一个或多个事件，然后按“f”键可快速缩放到时间轴的该部分。这有助于专注于特定训练步骤。

在处理大量跟踪记录事件时，Trace Viewer 会以流式传输模式运行。这意味着，当您在时间轴上平移和缩放时，它会按需加载数据，这与地图应用的工作方式类似。如果您缩放的速度快于数据加载的速度，那么在数据加载完成之前，您可能会看到数据的低分辨率表示形式。

Trace Viewer 界面组件

以下是 Trace Viewer 中的主要界面组件：

  • 时间轴位于顶部，水平延伸，显示相对于跟踪记录开始时的时间。
  • 时间轴按版块和轨道进行整理，标签位于左侧垂直轴上。每个部分代表一个处理元素（例如设备节点或主机线程），并且可以展开或收起。每个部分中都包含轨道，即特定活动的时间轴。
  • 事件是时间轴轨道上的彩色矩形块，表示操作或元事件（例如训练步）的时长。事件的颜色没有特定含义。
  • “详细信息”窗格会显示有关时间轴窗格中所选事件的更多信息，例如名称、开始时间和时长。

典型的版块和轨道

Trace Viewer 提供以下部分和轨道。

  • 每个 TPU 节点有一个版块，包含以下轨道：
    • XLA 模块：正在执行的 XLA 程序。
    • XLA 操作：显示在 TPU 核心上运行的 XLA HLO 操作。每个更高级别的框架操作（例如 JAX、TensorFlow 或 PyTorch）都会转换为一个或多个 XLA 操作，然后这些操作会被编译为在 TPU 上运行。点击 XLA 操作后，您可以在图表查看器中看到指向该操作的链接，以及有关该操作的其他信息，例如开始/停止时间、时长和源代码堆栈轨迹（如果框架和编译器提供）。
    • XLA TraceMe：用户在代码中指定的注解，用于描述他们打算跟踪的逻辑工作单元。即使您未添加任何注解，也可能会在此处看到数据；这些数据通常由 XLA（例如屏障核心）或 XProf 本身添加（例如，丢弃的轨迹条目）。
    • 步骤：显示在该 TPU 核心上运行的训练步骤的时长（如果在用户程序或框架中进行了适当注解）。
    • 框架操作：显示在 TPU 核心上执行的框架操作（例如 JAX、Tensorflow 或 PyTorch），前提是用户程序或框架中已对其进行适当注解。
    • 框架名称作用域：针对每个框架操作，显示堆栈轨迹的可视化图表。为简洁起见，此轨道仅显示单个设备。
    • 源代码：要执行的源代码的路径（如果配置文件中提供了）。
    • 标量单元：对于 TPU，在标量单元上执行的事件（如果配置文件中存在，则会显示）。
    • TensorCore 同步标志：TPU 上的同步机制，如果配置文件中存在，则会显示。
    • 主机分流：在主机内存和加速器内存之间异步移动数据的操作。XLA Ops 行上通常会显示相应的开始和停止操作，表示加速器正在为数据传输做好准备（例如，在传输期间将源/目标内存区域标记为“正在使用”）。如果有多个并行执行的卸载操作，则可能存在多个主机分流行，这需要轨迹查看器同时显示多个事件。
  • 每个 Sparsecore 节点一个版块：某些 TPU 代（例如除了密集计算 MXU 单元之外，TPU v5p 和 TPU v6e 还配备了一个或多个 SparseCore 单元；与这些核心关联的模块、运算和 TraceMe 将显示在此部分。
  • 每个 GPU 节点一个版块，包含以下轨道：
    • XLA 模块、框架操作、框架名称作用域、步骤、源代码。 这些都与 TPU 部分类似。
    • 每个数据流一个轨道，数据流名称还包含有关在数据流上执行的操作类型（Memcpy、计算等）的信息。
    • GPU 不支持 XLA TraceMe。
    • XLA 操作确实会显示在 GPU 部分，但由于这些操作是从数据流数据派生的，因此目前并不总是准确的。因此，它们无法完全说明 GPU 的执行模型，其中可能存在 XLA 运算与在不同流上执行的实际内核的 N:M 映射，以及将多个流动态调度到硬件中的不同 SM。
  • 针对在主机 CPU 上运行的每个组件（例如一个线程池）的一个版块，对于线程池，每个线程有一个轨道。如果您在配置文件收集期间启用了 Python 轨迹，还可以在该处看到这些轨迹。

请注意，只有 TPU 的 XLA 操作和 GPU 的流数据直接基于收集的配置文件；所有其他行都是“派生行”，其中涉及编译器提供的可选“边带”信息、可选的用户注释和/或 XProf 应用的启发词语。因此，这些派生行可能或可能不会出现在某些配置文件中。

其他功能

  • 您可以使用“查找活动...”搜索栏搜索特定事件名称。目前，此功能仅在屏幕上可见的时间范围内进行搜索，而不是在整个轨迹中进行搜索。
  • 数据流事件：点击顶栏中的“数据流事件”按钮可启用此选项，从而添加可视化图表，将一个线程中的事件与另一个线程中的事件相关联。例如，系统可能会从主机上用于将工作加入队列或启动加速器工作的操作，绘制一条箭头到执行该工作的加速器上的操作。XProf 通过结合使用用户注释、内置启发词语以及从不同组件（例如CUPTI 驱动程序、内核启动 ID、TPU 运行时信息等）。
  • 点击 XLA 操作后，详情窗格中会显示更多信息。 例如，它会链接到图表查看器工具中的操作。它还可以提供指向源代码和/或 Python 堆栈轨迹、导致生成此 XLA 操作的框架操作等的指针（如果配置文件中存在）。它还可能会显示 FLOPS（操作执行的浮点运算次数）和操作访问的字节数；这些信息是在编译期间从 XLA 静态获取的，而不是来自配置文件的运行时信息。