面向大众的 PyTorch 跟踪分析

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

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

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

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

1. 按维度细分
   • 时间：按单个节点和所有等级上花费在计算、通信、内存事件和空闲时间方面的 GPU 时间细分。
   • 空闲时间：将 GPU 空闲时间细分为等待主机、等待另一个内核或归因于未知原因。
   • 内核：查找每个等级上持续时间最长的内核。
   • 通信计算重叠：计算通信重叠计算的时间百分比。
2. 统计分析
   • 内核持续时间分布：不同等级中最长内核所用时间的分布。
   • CUDA 内核启动：持续时间非常短、持续时间很长和启动时间过长的 GPU 内核的分布。
   • 增强计数器（内存带宽、队列长度）：增强跟踪文件，可深入了解内存复制带宽和每个 CUDA 流上的未完成操作数量。
3. 模式
   • 频繁 CUDA 内核：查找任何给定的 PyTorch 或用户定义的运算符最常启动的 CUDA 内核。
4. 跟踪比较
   • 跟踪差异：跟踪比较工具，用于识别和可视化跟踪之间的差异。

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 代表主机，D 代表 GPU 设备。因此，H2D、D2H、D2D 分别代表主机到设备、设备到主机和设备到设备。

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

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

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

借助我们在 HTA 中构建的功能，我们旨在为用户提供“分布式 GPU 训练背后发生了什么？”的见解。我们将在接下来的几段中简要介绍这些功能。

整体跟踪分析的特点

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

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

内核细分：很自然地，人们会问哪些内核耗费的时间最多。下一个功能细分了每种内核类型（COMM、COMP、MEM）所花费的时间，并按持续时间进行排序。我们将每种内核类型和每个等级的信息以饼图的形式呈现。

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

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

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