跳转到主要内容
博客

面向大众的 PyTorch 跟踪分析

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

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

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

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

  1. 按维度细分
    1. 时间：将 GPU 时间细分为计算、通信、内存事件和单节点以及所有排名上的空闲时间。
    2. 空闲时间：将 GPU 空闲时间细分为等待主机、等待另一个内核或归因于未知原因。
    3. 内核：查找每个排名上持续时间最长的内核。
    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 Collective Communication Library (NCCL) 是一种广泛使用的通信库，其所有内核都带有“nccl”前缀。示例 NCCL 内核包括 NCCL_AllGather、NCCL_ReduceScatter、NCCL_AllReduce 等。
  3. 内存 (MEM) – 内存内核管理 GPU 设备上的内存分配/解除分配以及主机内存空间和 GPU 之间的数据移动。内存内核包括 Memcpy_H2D、Memcpy_D2H、Memcpy_D2D、Memset 等。此处，H 代表主机，D 代表 GPU 设备。因此，H2D、D2H、D2D 分别代表主机到设备、设备到主机和设备到设备。

由于 NVIDIA A100 GPU 这样的现代 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）内花费的时间，并按持续时间排序。我们以饼图的形式展示每种内核类型和每个排名的此信息。参见下面的图 3。

Figure 3: Pie chart of top computation and communication kernels

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

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

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

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

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

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

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。简而言之，用户需要克隆 repo 并通过 pip 安装必要的 Python 包。

使用方法

此版本的 Holistic Trace Analysis 目前处于测试阶段，我们建议在 Jupyter notebook 中使用 HTA。为方便起见，提供了演示笔记本。要开始使用，请在 Jupyter notebook 中导入 hta 包，创建一个 TraceAnalysis 对象，然后只需两行代码即可开始使用。

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

要求

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

常见问题解答

问：我如何安装 HTA？

请参阅存储库根目录中的 README。

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

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

问：您可以实现功能 X 吗？

根据该功能的广泛需求程度和实现所需的工作量，我们会考虑开发该功能。请在 Github Issue 中提出，并将其标记为 feature-request。

问：我可以修改代码吗？

请随意修改，如果您认为对其他人有用，请发送 PR。

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

请参阅此教程。

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

是的，请参阅此处的用例研究。