由 Anupam Bhatnagar、Xizhou Feng、Brian Coutinho、Yifan Liu、Sung-Han Lin、Louis Feng 和 Yuzhen Huang 共同撰写

我们很高兴宣布 Holistic Trace Analysis（HTA）的公开发布，这是一个为 PyTorch 用户提供的开源性能分析和可视化 Python 库。HTA 以 PyTorch 分析器收集的 Kineto 跟踪数据为输入，这些数据复杂且难以解释，并提升了这些跟踪中包含的性能信息。它最初在 Meta 内部开发，用于理解和调试在 GPU 上大规模分布式训练作业的性能问题。该多学科团队对 HTA 的功能进行了多项改进，并扩展到支持最先进的机器学习工作负载。

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

初始的 HTA 实现专门针对基于深度学习的推荐模型（DLRM）。为了使 HTA 中的功能通用并适用于分析视觉和 NLP 模型等用例，我们决定重构 HTA 代码库，并将库提供给更广泛的社区。这个新的代码库实现了几个重要的想法，这些想法导致了显著的效率和性能改进。

在本文中，我们介绍了 HTA 开源版本中实现的一些功能，这些功能可以用作 Python 脚本，也可以在 Jupyter 笔记本中交互式使用。HTA 提供了以下功能：

  1. 按维度拆分
    1. 时间：在单个节点和所有 rank 上，GPU 时间在计算、通信、内存事件和空闲时间上的分解。
    2. 空闲时间：将 GPU 空闲时间分解为等待主机、等待另一个内核或归因于未知原因。
    3. 内核：查找每个 rank 上持续时间最长的内核。
    4. 通信计算重叠：计算通信与计算重叠的时间百分比。
  2. 统计分析
    1. 内核持续时间分布：不同排名下最长内核的平均耗时分布
    2. CUDA 内核启动：非常短持续时间、长持续时间以及过长的启动时间的 GPU 内核分布
    3. 增强计数器（内存带宽、队列长度）：提供内存复制带宽和每个 CUDA 流上挂起的操作数量的增强跟踪文件
  3. 模式
    1. 常见 CUDA 内核：查找由任何给定的 PyTorch 或用户定义操作符频繁调用的 CUDA 内核。
  4. 跟踪比较
    1. 跟踪差异：一个跟踪比较工具，用于识别和可视化跟踪之间的差异。

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

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 分别代表主机到设备（Host to Device）、设备到主机（Device to Host）和设备到设备（Device to Device）。

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

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

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

图 1. 多个排名间 GPU 内核执行时间线的示例

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

通过我们在 HTA 中构建的功能，我们旨在让用户深入了解“分布式 GPU 训练内部发生了什么？”我们将在接下来的几段中简要描述这些功能。

Holistic Trace Analysis 功能

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

时间分解：我们首先询问 GPU 是否在计算、通信、内存事件上花费时间，还是处于空闲状态？为了回答这个问题，时间分解功能以这些类别为依据进行分解。为了实现高训练效率，代码应最大化计算内核的使用时间，最小化空闲时间和非计算时间（由通信或内存内核使用的时间）。这通过实现计算内核与通信或内存内核的并发执行来完成。请注意，在计算内核与通信/内存内核的并发执行过程中，通信/内存内核花费的时间计入计算时间。

Figure 2: Temporal Breakdown across 8 GPUs

图 2：8 个 GPU 的时间分解

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

Figure 3: Pie chart of top computation and communication kernels

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

内核持续时间分布：随后，人们还可以询问——对于任何给定的内核，其在各个 rank 上的时间分布是怎样的？为了回答这个问题，HTA 会为所有 rank 上的给定内核的平均持续时间生成条形图。此外，条形图中的误差线显示了给定内核在给定 rank 上所花费的最短和最长时间。图 4 显示了 rank 0 的平均持续时间与其他 rank 之间的差异。rank 0 上的这种异常行为指导用户查找可能的错误。

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

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

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

（在通信时花费的计算时间）/（花费在通信上的时间）

Figure 5: Communication computation overlap

图 5：通信计算重叠

增强型计数器（队列长度、内存带宽）：为了帮助调试，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 笔记本中使用 HTA。为了方便，提供了一个演示笔记本。要开始使用，请在 Jupyter 笔记本中导入 hta 包，创建一个 TraceAnalysis 对象，然后我们就可以用两行代码开始了。

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

要求

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

常见问题解答

问：如何安装 HTA？

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

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

文档和详细的 API 在此处可用。

问：能否实现功能 X？

根据该功能的需求范围和实现所需的努力程度，我们会考虑开发该功能。请创建一个 GitHub 问题，并标记为 feature-request 标签。

问：我可以修改代码吗？

请在过程中创建并发送一个 PR，如果你认为这对其他人有帮助。

问：如何在 PyTorch 中收集跟踪信息？

请参考此教程。

问：HTA 能否在生产规模上使用？

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