1. Meta 的 AI 性能分析（MAIProf）
图 1：Meta 的 AI 性能分析（MAIProf）基础设施的简化示意图。
图 1 给出了 Meta 的 AI 性能分析基础设施的简化示意图。机器学习研究和性能工程师通过用户门户向性能分析服务提交训练作业的性能分析请求，该服务随后将请求广播到所有运行训练作业的 GPU 主机。当 GPU 主机上的监控守护进程接收到性能分析请求时，它将通知对应训练作业的 PyTorch 程序中的 Kineto GPU 跟踪器（基于 NVIDIA 的 libcupti 构建）。因此，将收集 Kineto 跟踪器并将其异步上传到对象存储（更详细地说：每个 GPU 收集一个 Kineto 跟踪器，每个跟踪器都被处理并存储为一个 blob；示例将在第 2 节中给出）。同时，MAIProf 还收集各种汇总的性能指标：每个 GPU 主机上的监控守护进程持续从 NVIDIA 的 DCGM/NVML 读取性能计数器，并将它们记录到时间序列数据库中。
一旦完成跟踪和指标收集，性能分析服务将自动从对象存储下载跟踪数据以进行跟踪分析，并从时间序列数据库下载性能指标以进行指标分析。最后，将向用户交付一份包含详细和深入分析的总体分析报告。
为了满足生产需求，我们在 MAIProf 中故意做出了以下设计选择：
- 无需修改 PyTorch 模型的源代码：通过采样用户指定时间内的未修改模型的执行来触发分析。
- 提供性能的整体视图：MAIProf 执行系统级分析，涵盖 CPU 和 GPU。在底层，它调用各种 CPU 工具（例如 Python 跟踪器、Autograd 观察器）和 GPU 工具（例如 Kineto、DCGM），并关联其结果。
- 提供多种针对广泛 AI 分区器的工具：在 Meta，有来自不同背景的工程师，他们可能需要调整他们的 AI 工作负载性能。其中一些人是 AI 专家，而其他人则是通用软件工程师。因此，MAIProf 提供了各种工具，以适应不同级别的性能调试，从高级自动跟踪理解到低级跟踪分析。
- 支持分布式 GPU 分析：MAIProf 可以从多个主机收集分析数据，每个主机都有多个 GPU。然后它显示整个系统的综合视图/分析。
- 高度可扩展：MAIProf 建立在 Meta 数据中心现有基础设施之上，如可扩展存储系统 Manifold。随着工作负载的增加，其分析能力可以通过向服务池中添加更多机器轻松扩展。
2. 案例研究：优化保护 PyTorch 模型
为了具体说明，我们以一个在生产中使用的 PyTorch 保护模型为案例研究。首先，我们讨论使用 MAIProf 识别模型性能瓶颈的步骤。然后，我们描述相应的优化措施及其影响。
2.1 性能瓶颈
步骤 1：
检查 CPU 和 GPU 在同一时间线上的利用率，如图 2 所示。
图 2：随时间变化的 CPU 使用率（上方）与随时间变化的 GPU 使用率（下方）。
我们在图 2 中注意到的第一个性能异常模式是：“GPU 空闲，GPU 活跃，GPU 空闲，GPU 活跃……”贯穿整个训练过程。总的来说，GPU 空闲时间超过训练时间的一半（这对性能不利，因为 GPU 是高性能设备，所以我们希望尽可能多地利用它）。
第 2 步：
当 GPU 空闲时，使用 MAIProf 在 CPU 上收集 Python 函数调用跟踪，如图 3 所示。
图 3：Python 调用跟踪。
Python 跟踪显示，大部分 CPU 时间都花在 Python 函数 sharded_iterrows() 内部。从模型源代码中，我们了解到这个函数并行处理一个大特征表。使用的 worker 线程数量由一个可配置的参数控制（ num_worker_threads ）。此外，经过调查特征表的生成方式，我们理解了性能异常的原因：训练数据集太大，无法一次性装入 CPU 内存；需要将其分成多个子数据集，每个子数据集有足够的数据运行 10 个 epoch。因此，每 10 个 epoch 需要从磁盘读取新的子数据集到内存中，在此期间 GPU 完全空闲。
第 3 步：
收集 GPU 性能指标，如图 4 所示。
图 4：MAIProf 中的 GPU 性能指标。
从图 4 中我们得出以下观察结果：
- 流式多处理器（SM）运行模型的 CUDA 内核。其利用率[1]为 9.1%，表明 GPU 上的并行计算单元没有得到充分利用。
- Tensor Core 利用率为 0，意味着 Tensor Core（GPU 上的混合精度计算单元[2]）完全没有使用。
- 最大 GPU 内存利用率是 47.13%，表明 GPU 内存有一半未被使用。
步骤 4：
收集训练循环的 GPU 跟踪（即 Kineto 跟踪）如图 5 所示。
图 5：训练循环的 GPU 跟踪（即 Kineto 跟踪）。
由于常用的 PyTorch 函数已经进行了注释，因此它们的名称会自动显示在追踪中。有了它们，我们可以大致将追踪分为训练迭代中的四个阶段：（1）数据加载，（2）正向传播，（3）反向传播，（4）梯度优化（注意：在图 5 中，“优化器”阶段来自前一个批次，而其他三个阶段来自当前批次）。
2.2 优化
我们对上述识别出的瓶颈进行了四次简单的优化，每个优化只需要更改一个配置参数或最多几行源代码。它们列在图 6 中。
| 优化 | 变更量 | 瓶颈问题已解决 |
|---|---|---|
通过尝试每个主机 CPU 核心的几个可能值来调整 num_worker_threads 。 |
1 源代码行 | GPU 完全空闲时间 |
| 批量大小加倍 | 2 个配置参数 | GPU 内存未充分利用 |
| 在 PyTorch 中使用自动混合精度 | 13 源代码行 | 核心张量利用率为零 |
| 在 PyTorch 中使用多张量优化器 | 1 源代码行 | 优化器中许多小的 GPU 内核 |
图 6：应用了四种简单的优化
3. 结论
在生产环境中对 PyTorch 进行性能调优越来越重要。一个强大的性能调试工具是这一过程的关键。我们通过一个生产模型的案例研究证明了 MAIProf 是识别优化机会的强大基础设施。
在 Meta，MAIProf 已被数百名工程师使用，从性能新手到专家，以识别更多类型的瓶颈。这些包括缓慢的数据加载、小而慢的 GPU 内核、分布式训练问题，如负载不平衡和过度的通信。MAIProf 覆盖了包括推荐、视觉和自然语言处理在内的主要模型类别。总之，它现在已成为调整生产 PyTorch 工作负载性能的不可或缺的工具。
参考文献列表
[1] https://docs.nvidia.com/gameworks/content/developertools/desktop/analysis/report/cudaexperiments/kernellevel/achievedoccupancy.htm
