由 Anupam Bhatnagar 和 Brian Coutinho 撰写

在本文中，我们分享了我们如何在不进行任何用户端代码配置的情况下，为训练工作负载收集和分析 PyTorch Profiler 跟踪。我们利用了 Dynolog——一个开源的 CPU 和 GPU 遥测守护进程来收集 PyTorch Profiler 跟踪，并使用 Holistic Trace Analysis——一个开源库来分析收集到的跟踪。这个工具链使 Meta 的工程师能够加速他们的性能优化工作流程。我们解决方案的核心是在 PyTorch 的基类 Optimizer 中实现预和后钩子。我们在一个简短的视频中演示了使用 Dynolog 进行 PyTorch 跟踪收集。

问题

Meta 的软件开发人员每天运行大量分布式训练运行。为了确保 GPU 得到有效利用，有必要对所有工作测量和分析 GPU 性能。此外，开发者需要具备检查模型并了解 CPU 和 GPU 如何交互以调试性能问题的能力。开发者使用少量 GPU 构建初始原型，而生产版本扩展到数百或数千个 GPU，服务于众多业务用例，如生成式 AI、推荐系统、广告排名等。

在 Meta 这样的规模下，需要拥有低开销且能够无缝协同工作的性能测量和监控工具链，以保持高开发者效率。

在本文中，我们描述了如何使用 PyTorch Profiler、Dynolog（一个遥测守护进程）和 Holistic Trace Analysis（一个性能调试库）来收集跟踪信息，而无需在用户端代码中进行任何仪器设置，并分析它们以识别低 GPU 利用率的作业。

解决方案

下面的图表展示了工具链如何协同工作。

  1. 用户启动一个 PyTorch 应用程序。
  2. 训练服务或用户使用 Dynolog CLI 触发一个分析会话，该 CLI 通过网络发送请求到 Dynolog 守护进程。
  3. Dynolog 守护进程将分析配置中继到 PyTorch 应用程序，将其临时设置为分析模式。
  4. PyTorch 分析器收集跟踪并将其存储到数据库中（例如，网络文件系统或 S3 存储桶）。
  5. 然后使用全面跟踪分析（HTA）对这些收集到的跟踪进行分析。

Figure 1: Dynolog, PyTorch Profiler and HTA toolchain workflow

图 1：Dynolog、PyTorch 分析器和 HTA 工具链工作流程

让我们深入探讨每个组件。

Dynolog

Dynolog 是一个轻量级的监控守护进程，适用于异构 CPU-GPU 系统。它支持对 CPU（利用率、网络带宽、每秒指令数）和 GPU（SM 占用率、DRAM 带宽、GPU 功耗）的性能指标进行持续监控。此外，dynolog 还导出 API 以收集深度分析数据，这些数据可以通过 dyno CLI 访问。

Dynolog 提供的主要集成之一是与 PyTorch Profiler 的接口。这可以通过单个命令进行远程跟踪，以跟踪数千台服务器。这可以通过使用 dyno gputrace 命令实现。

PyTorch 分析器

GPU 内核异步执行，需要 GPU 端支持来创建跟踪。NVIDIA 通过 CUPTI 库提供这种可见性。Kineto 是 Profiler 中的子系统，它与 CUPTI 接口。PyTorch Profiler 利用 Kineto 库来收集 GPU 跟踪。为了在不进行任何用户端代码配置的情况下，实现大规模训练工作负载的自动化分析，我们对 PyTorch 进行了一些基本更改。这些更改使得跟踪收集无需用户干预。

  • **注册：** 首先，我们对 PyTorch 进行了修改，以便在启动时与 Dynolog 守护进程进行注册。通过设置环境变量 KINETO_USE_DAEMON=True 来启用此功能。当此环境变量设置为 True 时，PyTorch Profiler 会定期轮询 Dynolog 以检查按需跟踪请求。
  • **迭代钩子：** 然后，我们为基本优化器类实现了预和后钩子。这使得我们能够注释训练迭代的开始/结束。然后，分析器会知道迭代次数，并可以在跟踪中安全地捕获固定数量的迭代。

**整体跟踪分析（HTA）**

机器学习研究人员和工程师常常在计算上难以扩展他们的模型，因为他们对自己的工作负载中的性能瓶颈一无所知。大规模分布式训练作业可能会生成数千个跟踪，包含的数据量对于人类来说过于庞大，难以检查。这就是整体跟踪分析（Holistic Trace Analysis）发挥作用的地方。HTA 是一个开源的性能分析库，它接受 PyTorch Profiler 跟踪作为输入，并提升其中包含的性能信息。其目标是帮助研究人员和工程师从硬件堆栈中获得最佳性能。为了帮助性能调试，HTA 提供了以下功能（部分列表）：

  • 时间分解：按计算、通信、内存事件和单个节点以及所有 rank 的空闲时间分解 GPU 时间。
  • 空闲时间分解：将 GPU 空闲时间分解为等待主机、等待另一个内核或归因于未知原因。
  • 内核分解：找到每个 rank 上持续时间最长的内核。
  • 核心持续时间分布：不同排名下最长核心的平均耗时分布。
  • 通信与计算重叠：计算通信与计算重叠的时间百分比。

欢迎您查看这些 Jupyter 笔记本，了解 HTA 能为您做什么。如果您是首次使用，我们建议从 trace_analysis_demo 笔记本开始。

总结来说，Dynolog 允许我们以可扩展的方式实时收集 PyTorch Profiler 跟踪。此外，通过利用 HTA，我们可以自动化性能分析并识别瓶颈。在 Meta，我们使用 Dynolog、PyTorch Profiler 和 HTA 工具链来加速我们的性能优化工作流程。

演示

我们分享了一个屏幕录制，展示了如何为玩具 PyTorch 程序收集跟踪，而不需要任何用户端代码配置。该演示在 Docker 容器中运行，跟踪收集是通过 Dynolog 触发的。可以使用 HTA 来随后分析收集到的跟踪。

常见问题解答

问： dyno gputrace 还能为我做什么？

The dyno gputrace 命令支持多个自定义 PyTorch Profiler 选项：

  • 捕获 Python 栈
  • 内存分析
  • 记录输入形状

请运行 dyno gputrace --help 以查看所有选项。

问：Dynolog 是否收集硬件性能指标？

Dynolog 还可用于始终在线监控：

  • 它集成了针对 NVIDIA GPU 的即插即用 GPU 性能监控，使用 DCGM。
  • Dynolog 提供基本的 Linux 内核性能指标，包括 CPU、网络和 IO 资源使用情况。
  • Dynolog 管理 Intel 和 AMD CPU 上与 CPU 缓存、TLBs 等微架构特定事件相关的硬件性能计数器。

问：我该如何构建演示中使用的 Docker 镜像？

Dockerfile 在此处可用。使用以下命令构建 Docker 镜像。

docker build -f /path/to/dynolog_repo/dynolog_hta.dockerfile -t <image_name:tag> .

如何运行 Docker 镜像？

您可以参考这份速查表来运行 Docker 镜像。