作者:Anupam Bhatnagar, Brian Coutinho

在这篇博客中,我们分享了如何在无需任何用户端代码插桩的情况下,为训练工作负载启用 PyTorch Profiler 跟踪的收集和分析。我们利用 Dynolog(一个用于 CPU 和 GPU 遥测的开源守护进程)来收集 PyTorch Profiler 跟踪,并使用 Holistic Trace Analysis(一个用于分析 PyTorch Profiler 跟踪的开源库)来分析收集到的跟踪。这套工具链使得 Meta 的工程师能够加速其性能优化工作流程。我们解决方案的关键是在 PyTorch 的基础 Optimizer 类中实现了前置和后置钩子(hooks)。我们在一个短视频中演示了如何使用 Dynolog 进行 PyTorch 跟踪收集。

问题

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

鉴于 Meta 的规模,需要一套开销低且彼此无缝协作的性能测量和监控工具链,以保持较高的开发者效率。

在这篇博客中,我们描述了如何使用 PyTorch Profiler、Dynolog(一个遥测守护进程)和 Holistic Trace Analysis(一个性能调试库)来在无需任何用户端代码插桩的情况下收集跟踪,并分析它们以识别 GPU 利用率低的作业。

解决方案

下图概述了这套工具链如何协同工作。

  1. 用户启动一个 PyTorch 应用程序。
  2. 训练服务或用户使用 Dynolog CLI 触发一个性能分析会话,Dynolog CLI 通过网络向 Dynolog 守护进程发送请求。
  3. Dynolog 守护进程将性能分析配置中继给 PyTorch 应用程序,将其临时设置为性能分析模式。
  4. PyTorch Profiler 收集跟踪并将其存储到数据库(例如,网络文件系统或 S3 存储桶)。
  5. 然后使用 Holistic Trace Analysis (HTA) 分析收集到的跟踪。

Figure 1: Dynolog, PyTorch Profiler and HTA toolchain workflow

图 1:Dynolog、PyTorch Profiler 和 HTA 工具链工作流程

让我们深入了解各个组件。

Dynolog

Dynolog 是一个用于异构 CPU-GPU 系统的轻量级监控守护进程。它支持持续监控来自 CPU(利用率、网络带宽、每秒指令数)和 GPU(SM 占用率、DRAM 带宽、GPU 功耗)的性能指标。此外,dynolog 暴露了 API 以收集可通过 dyno CLI 访问的深度性能分析数据。

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

PyTorch Profiler

GPU 内核(kernels)是异步执行的,需要 GPU 侧的支持才能创建跟踪。NVIDIA 通过 CUPTI 库提供这种可见性。Kineto 是 Profiler 内部与 CUPTI 对接的子系统。PyTorch Profiler 利用 Kineto 库收集 GPU 跟踪。为了在无需任何用户端代码插桩的情况下大规模地自动进行训练工作负载的性能分析,我们对 PyTorch 进行了几项基础性修改。这些修改使得跟踪收集无需任何用户干预。

  • 注册:** **首先,我们修改了 PyTorch,使其在启动时向 Dynolog 守护进程注册。通过设置环境变量 KINETO_USE_DAEMON=True 可以开启此功能。将此环境变量设置为 True 后,PyTorch Profiler 会定期轮询 Dynolog,检查是否有按需跟踪请求。
  • 迭代钩子:然后,我们为基础 Optimizer 类实现了前置和后置钩子(hooks)。这使得我们能够标记训练迭代的开始/结束。Profiler 随后会知道迭代次数,并可以安全地在跟踪中捕获固定数量的迭代。

Holistic Trace Analysis (HTA)

机器学习研究人员和工程师常常难以在计算上扩展其模型,因为他们不清楚其工作负载中的性能瓶颈。大型分布式训练作业可能会生成数千个跟踪,数据量巨大,人工难以检查。这就是 Holistic Trace Analysis 的用武之地。HTA 是一个用于性能分析的开源库——它以 PyTorch Profiler 跟踪作为输入,并提升其中包含的性能信息。其目标是帮助研究人员和工程师从硬件堆栈中获得最佳性能。为了协助性能调试,HTA 提供了以下功能(部分列表):

我们邀请您查阅这些Jupyter notebook,了解 HTA 能为您做什么。如果您是初次用户,我们建议从 trace_analysis_demo notebook 开始。

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

演示

我们分享了一个截屏视频,展示了如何在无需任何用户端代码插桩的情况下为一个简单的 PyTorch 程序收集跟踪。该演示在 Docker 容器中运行,跟踪收集由 Dynolog 触发。随后可以使用 HTA 分析收集到的跟踪。

常见问题

问:dyno gputrace 还能做什么?

dyno gputrace 命令支持几种自定义的 PyTorch Profiler 选项:

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

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

问:Dynolog 是否收集硬件性能指标?

Dynolog 也可以用于持续监控

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

问:如何构建演示中使用的 Docker 镜像?

dockerfile 可在此获取。使用以下命令构建 Docker 镜像。

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

问:如何运行 Docker 镜像?

您可以参考这份备忘单来运行 Docker 镜像。