博客

面向大众的 PyTorch 跟踪分析

我们很高兴地宣布正式发布 Holistic Trace Analysis (HTA)，这是一个面向 PyTorch 用户的开源性能分析和可视化 Python 库。HTA 以 Kineto 跟踪文件（由 PyTorch Profiler 采集，通常复杂且难以解读）为输入，并对这些跟踪文件中包含的性能信息进行深度解析。该库最初由 Meta 内部开发，旨在理解和调试 GPU 上大规模分布式训练作业的性能问题。跨学科团队对 HTA 的功能进行了多项改进，并将其扩展以支持最先进的机器学习工作负载。

机器学习研究人员和系统工程师经常难以在计算上扩展其模型，因为他们无法察觉工作负载中的性能瓶颈。由于缺乏对“幕后”情况的可见性，作业申请的资源（如 GPU、内存）往往与实际所需资源不匹配。为了从硬件堆栈中获得最佳性能，必须了解分布式训练工作负载的资源利用率和瓶颈。

HTA 的初始实现专门针对基于深度学习的推荐模型 (DLRM)。为了使 HTA 的功能通用化并适用于视觉和 NLP 模型等用例，我们决定重构 HTA 代码库并向广大社区开放。这一新代码库实现了几个重要的理念，从而带来了显著的效率和性能提升。

在本篇博客中，我们将介绍 HTA 开源版本中实现的多项功能，这些功能既可以作为 Python 脚本使用，也可以在 Jupyter Notebook 中交互式使用。HTA 提供以下功能：

  1. 维度分解
    1. 时间维度：GPU 时间的细分，包括单个节点及所有 rank 上计算、通信、内存事件所占用的时间以及空闲时间。
    2. 空闲时间：将 GPU 空闲时间细分为等待主机、等待其他内核或归因于未知原因。
    3. 内核维度：查找每个 rank 上持续时间最长的内核。
    4. 通信计算重叠：计算通信与计算发生重叠的时间百分比。
  2. 统计分析
    1. 内核持续时间分布：不同 rank 上最长内核的平均耗时分布。
    2. CUDA 内核启动：持续时间极短、极长以及启动时间过长的 GPU 内核分布。
    3. 增强计数器（内存带宽、队列长度）：增强型跟踪文件，提供有关内存复制带宽和每个 CUDA 流上待处理操作数量的洞察。
  3. 模式识别
    1. 频繁 CUDA 内核：查找由任何给定 PyTorch 算子或用户自定义算子最频繁调用的 CUDA 内核。
  4. 跟踪对比
    1. 跟踪差异 (Trace Diff)：一个跟踪对比工具，用于识别和可视化不同跟踪文件之间的差异。

用户可通过 Github 获取 HTA 源代码。除了上述功能外，用户还可以使用代码库中提供的核心库和数据结构，申请新功能或构建自己的分析逻辑。

GPU 训练性能调试基础 101

为了理解分布式训练作业中的 GPU 性能，我们需要考虑模型算子如何与 GPU 设备交互，以及这些交互如何反映在某些可衡量的指标中。

从高层面来看，我们可以将模型执行中的 GPU 操作分为三大类，以下统称为内核类型：

  1. 计算 (COMP) – 计算内核执行矩阵乘法和类似数值计算的编译程序。它们负责模型执行所需的所有数值运算。
  2. 通信 (COMM) – 通信内核负责在分布式训练作业的不同 GPU 设备之间交换和同步数据。NVIDIA 集体通信库 (NCCL) 是广泛使用的通信库，其所有内核的前缀均为“nccl”。常见的 NCCL 内核包括 NCCL_AllGather、NCCL_ReduceScatter、NCCL_AllReduce 等。
  3. 内存 (MEM) – 内存内核管理 GPU 设备上的内存分配/释放，以及主机和 GPU 之间的内存空间数据移动。内存内核包括 Memcpy_H2D、Memcpy_D2H、Memcpy_D2D、Memset 等。其中，H 代表主机 (Host)，D 代表 GPU 设备 (Device)。因此，H2D、D2H、D2D 分别代表主机到设备、设备到主机和设备到设备的数据传输。

由于 NVIDIA A100 等现代 GPU 设备是能够同时运行多个内核的大规模并行设备，因此可以通过重叠计算、通信和内存内核来减少模型执行时间。实现重叠的一种常见技术是利用多个 CUDA 流。CUDA 流是按主机代码发出顺序在 GPU 设备上执行的操作序列。不同的 CUDA 流可以交错甚至并发运行，从而实现内核重叠的效果。

为了帮助理解上述概念，图 1 展示了一个 8 GPU 分布式训练作业中一次迭代的 GPU 内核时间线。在下图中，每个 rank 代表一个 GPU，每个 GPU 上的内核运行在 6 个 CUDA 流上。在图的右侧，可以看到所使用的 GPU 内核名称。在图的中间，可以看到计算内核和通信内核之间的重叠。此图是使用 HTA 中提供的 plot_timeline 示例 Notebook 创建的。

Figure 1. An example of the execution timeline of GPU Kernels across multiple ranks

图 1. 多 GPU 下 GPU 内核执行时间线示例

多 GPU 训练作业的性能受多种因素影响。其中，模型执行如何创建和编排 GPU 内核起着至关重要的作用。HTA 提供了关于模型执行如何与 GPU 设备交互的洞察，并突显了提升性能的机会。

通过我们在 HTA 中构建的功能，我们旨在让用户了解“分布式 GPU 训练的幕后究竟发生了什么”。我们将在接下来的段落中简要描述这些功能。

Holistic Trace Analysis 的功能

对于大多数用户而言，理解 GPU 训练作业的性能并非易事。因此，我们构建了这个库来简化跟踪分析任务，并通过检查模型执行跟踪文件为用户提供有用的洞察。作为第一步，我们开发了一些重要且通用的功能，以便大多数用户都能从中受益。

时间分解 (Temporal Breakdown)：我们首先要问的是，GPU 时间是在计算、通信、内存事件上，还是处于空闲状态？为了回答这个问题，时间分解功能以这些类别呈现细分情况。为了实现高训练效率，代码应最大限度地利用计算内核，并尽可能减少空闲时间和非计算时间（通信或内存内核所用的时间）。这是通过实现计算内核与通信或内存内核的并发执行来达到的。注意：在计算内核与通信/内存内核并发执行期间，通信/内存内核所花费的时间会计入计算时间中。

Figure 2: Temporal Breakdown across 8 GPUs

图 2：8 个 GPU 上的时间分解

内核分解 (Kernel Breakdown)：用户自然会问哪些内核占用了最多的时间。接下来的功能细分了每种内核类型（COMM、COMP、MEM）内花费的时间，并按持续时间进行排序。我们以饼图形式呈现每种内核类型和每个 rank 的此信息。请参阅下图 3。

Figure 3: Pie chart of top computation and communication kernels

图 3：主要计算和通信内核的饼图

内核持续时间分布 (Kernel Duration Distribution)：随后，还可以问——对于任何给定的内核，其在各 rank 上的耗时分布如何？为了回答这个问题，HTA 为所有 rank 上给定内核的平均持续时间生成条形图。此外，条形图中的误差线显示了给定内核在特定 rank 上所用的最短和最长时间。下图 4 显示了 rank 0 上的平均持续时间与其他 rank 相比存在差异。这种 rank 0 上的异常行为可以引导用户定位潜在的 Bug。

Figure 4: Average duration of NCCL AllReduce Kernel across 8 ranks

图 4：NCCL AllReduce 内核在 8 个 rank 上的平均持续时间

通信计算重叠 (Communication Computation Overlap)：在分布式训练中，大量时间花费在多个 GPU 设备之间的通信和同步事件上。为了实现高 GPU 效率（即 TFLOPS/GPU），保持 GPU 进行实际计算工作至关重要。换句话说，GPU 不应因等待来自其他 GPU 的数据而被阻塞。衡量计算被数据依赖阻塞程度的一种方法是计算计算-通信重叠。如果通信事件与计算事件重叠，则观察到更高的 GPU 效率。缺乏通信和计算重叠会导致 GPU 空闲，从而降低效率。因此，通信计算重叠功能为每个 rank 计算作业中通信与计算重叠的时间百分比，并生成条形图表示。请参阅下图。更具体地说，我们测量以下比率：

(通信期间用于计算的时间) / (通信总时间)

Figure 5: Communication computation overlap

图 5：通信计算重叠

增强计数器 (Augmented Counters)：为了辅助调试，HTA 计算了 D2H、H2D 和 D2D 内存复制 (memcpy) 和内存设置 (memset) 事件的内存带宽统计信息。此外，HTA 还计算了每个 CUDA 流上待处理的 CUDA 操作数量。我们称之为队列长度。当流上的队列长度达到 1024 或更大时，新事件无法在该流上调度，CPU 将会阻塞，直到 GPU 事件被处理完毕。此外，HTA 生成一个新的跟踪文件，其中包含带有内存带宽和队列长度时间序列的轨道。请参阅下图 6。

Figure 6: Memory Bandwidth and Queue Length

图 6：内存带宽和队列长度

这些核心功能让我们能够窥探系统性能，并帮助回答“系统中正在发生什么？”。随着 HTA 的发展，我们希望能够解决“为什么会发生 X？”的问题，并建议可能的解决方案来克服瓶颈。

安装与使用

安装

有关安装 HTA 的说明，请参考 README。简而言之，用户需要克隆 代码库 并通过 pip 安装必要的 Python 包。

使用

当前版本的 Holistic Trace Analysis 处于测试阶段，我们建议在 Jupyter Notebook 中使用 HTA。为了方便起见，我们提供了一个 演示 Notebook。入门很简单，只需在 Jupyter Notebook 中导入 hta 包，创建一个 TraceAnalysis 对象，两行代码即可开启分析之旅。

from hta.trace_analysis import TraceAnalysis
analyzer = TraceAnalysis(trace_dir = “/trace/folder/path”)

要求

  • 训练或推理作业的所有跟踪文件必须存储在一个唯一文件夹中。
  • 跟踪文件必须为 json 或 gzipped json 格式。

常见问题解答 (FAQ)

问：我该如何安装 HTA？

请参阅代码库根目录中的 README。

问：是否有关于 HTA 功能和 API 的文档？

文档和详细的 API 说明可在 此处 找到。

问：你们能实现 X 功能吗？

我们会根据该功能的普遍需求程度和实现所需的工作量考虑开发。请在 Github Issues 中提交请求，并打上 feature-request 标签。

问：我可以修改代码吗？

请随时修改，如果您认为这对其他人有帮助，欢迎 提交 PR。

问：我该如何在 PyTorch 中采集跟踪数据？

请参考此处的教程。

问：HTA 可以用于生产规模吗？

是的，请查看此处的用例研究。