作者：Anupam Bhatnagar, Xizhou Feng, Brian Coutinho, Yifan Liu, Sung-Han Lin, Louis Feng, and Yuzhen Huang

我们很高兴宣布 Holistic Trace Analysis (HTA) 正式公开发布，这是一个面向 PyTorch 用户的开源性能分析和可视化 Python 库。HTA 接收 Kineto traces 作为输入，这些 traces 由 PyTorch profiler 收集，它们复杂且难以解释，而 HTA 则提取并优化了这些 traces 中包含的性能信息。它最初由 Meta 内部开发，用于理解和调试 GPU 上大规模分布式训练作业的性能问题。多学科团队对 HTA 的功能进行了多项增强，并将其扩展以支持最先进的 ML 工作负载。

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

最初的 HTA 实现专门针对基于深度学习的推荐模型 (DLRM)。为了使 HTA 中的功能具有通用性并适用于分析视觉和自然语言处理模型等用例，我们决定重构 HTA 代码库，并将该库提供给更广泛的社区。这个新的代码库实现了几个重要的想法，从而显著提高了效率和性能。

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

  1. 按维度细分
    1. 时间维度：按单节点和所有 ranks 上花费在计算、通信、内存事件和空闲时间方面的 GPU 时间细分。
    2. 空闲时间：将 GPU 空闲时间细分为等待 host、等待另一个 kernel 或未知原因造成的空闲时间。
    3. Kernel：查找每个 rank 上持续时间最长的 kernel。
    4. 通信计算重叠：计算通信与计算重叠的时间百分比。
  2. 统计分析
    1. Kernel 持续时间分布：不同 ranks 上最长 kernel 平均耗时的分布。
    2. CUDA Kernel 启动：持续时间非常短、持续时间很长以及启动时间过长的 GPU kernel 的分布。
    3. 增强计数器（内存带宽、队列长度）：增强的 trace 文件，提供关于内存复制带宽和每个 CUDA stream 上未完成操作数量的洞察。
  3. 模式
    1. 频繁调用的 CUDA Kernel：查找由任何给定 PyTorch 或用户定义算子最频繁启动的 CUDA kernel。
  4. Trace 对比
    1. Trace Diff：一种 trace 对比工具，用于识别和可视化 traces 之间的差异。

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

GPU 训练性能调试 101

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

从宏观上看，我们可以将模型执行中的 GPU 操作分为三大类，下文统称为 kernel 类型：

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

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

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

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

图 1. 跨多个 ranks 的 GPU Kernel 执行时间轴示例

多 GPU 训练作业的性能受多种因素影响。在这些因素中，模型执行如何创建和编排 GPU kernel 起着关键作用。HTA 提供了关于模型执行如何与 GPU 设备交互的洞察，并突出了性能改进的机会。

通过我们在 HTA 中构建的功能，我们旨在为用户提供对“分布式 GPU 训练底层发生了什么？”的洞察。我们将在接下来的几段中简要描述这些功能。

Holistic Trace Analysis 中的功能

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

时间维度细分：我们首先要问的是，GPU 是否在计算、通信、内存事件上花费时间，或者它是空闲的？为了回答这个问题，时间维度细分功能按这些类别提供了细分报告。为了达到高训练效率，代码应最大化计算 kernel 使用的时间，并最小化空闲时间和非计算时间（通信或内存 kernel 使用的时间）。这可以通过实现计算 kernel 与通信或内存 kernel 的并发执行来实现。请注意，在计算 kernel 与通信/内存 kernel 并发执行期间，通信/内存 kernel 花费的时间被计入计算时间。

Figure 2: Temporal Breakdown across 8 GPUs

图 2：跨越 8 个 GPU 的时间维度细分

Kernel 细分：很自然地会问，哪些 kernel 花费的时间最多。下一个功能按每个 kernel 类型（COMM、COMP、MEM）细分花费的时间，并按持续时间排序。我们将此信息按每个 kernel 类型和每个 rank 以饼状图的形式呈现。参见下图图 3。

Figure 3: Pie chart of top computation and communication kernels

图 3：顶级计算和通信 kernel 的饼状图

Kernel 持续时间分布：随后，还可以问——对于任何给定的 kernel，其在不同 ranks 上的时间花费分布是怎样的？为了回答这个问题，HTA 生成条形图，显示给定 kernel 在所有 ranks 上的平均持续时间。此外，条形图中的误差条显示了给定 kernel 在给定 rank 上花费的最短和最长时间。下图图 4 显示了 rank 0 上的平均持续时间与其他 ranks 相比存在的差异。rank 0 上的这种异常行为指导用户查找可能的 bug。

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

图 4：NCCL AllReduce Kernel 在 8 个 ranks 上的平均持续时间

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

(通信期间花费在计算上的时间) / (花费在通信上的时间)

Figure 5: Communication computation overlap

图 5：通信计算重叠

增强计数器（队列长度、内存带宽）：为了帮助调试，HTA 计算 D2H、H2D 和 D2D 内存复制 (memcpy) 和内存设置 (memset) 事件的内存带宽统计数据。此外，HTA 还计算每个 CUDA stream 上未完成的 CUDA 操作数量。我们称之为队列长度。当一个 stream 上的队列长度达到或超过 1024 时，新的事件将无法在该 stream 上安排，CPU 将会停顿，直到 GPU 事件处理完毕。此外，HTA 生成一个新的 trace 文件，其中包含显示内存带宽和队列长度时间序列的轨迹。参见下图图 6。

Figure 6: Memory Bandwidth and Queue Length

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

这些主要功能让我们得以一窥系统性能，并有助于回答“系统中正在发生什么？”这个问题。随着 HTA 的发展，我们希望能够解决“为什么会发生 X？”的问题，并提供克服瓶颈的可能解决方案。

安装和使用

安装

有关 HTA 的安装，请参阅 README 文件。简而言之，用户需要克隆 仓库 并通过 pip 安装必要的 Python 包。

使用

当前版本的 Holistic Trace Analysis 处于 Beta 阶段，我们建议在 Jupyter notebook 中使用 HTA。为了您的方便，我们提供了一个 演示 notebook。要开始使用，只需在 Jupyter notebook 中导入 hta 包，创建一个 TraceAnalysis 对象即可，总共只需要两行代码。

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

要求

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

常见问题解答

问：如何安装 HTA？

请参阅仓库根目录中的 README 文件。

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

文档和详细 API 可在此处获取：此处。

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

根据该功能的广泛需求程度和实现所需的投入，我们将考虑开发该功能。请在 Github 上创建一个 Issue，并打上 feature-request 标签。

问：我可以修改代码吗？

当然可以，如果您认为您的修改对其他人有用，请随时提交 PR。

问：如何在 PyTorch 中收集 traces？

请参阅此处提供的教程。

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

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