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
使用 MAIProf 收集 GPU 空闲时 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 中得到了以下观察结果：
- Streaming multiprocessor (SM) 运行模型的 CUDA kernels。其利用率 [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 中，“optimizer”阶段来自前一个批次，而其他三个阶段来自当前批次）。
2.2 优化
我们进行了四个简单的优化，针对上面识别出的瓶颈，每个优化只需要更改一个配置参数或至多几行源代码。它们列在图 6 中。
| 优化措施 | 修改量 | 解决的瓶颈 |
|---|---|---|
通过尝试每个主机上 CPU 核心数的几个可能值来调整 num_worker_threads。 |
1 行源代码 | GPU 完全空闲时间 |
| 将批次大小加倍 | 2 个配置参数 | GPU 内存利用不足 |
| 使用 PyTorch 中的自动混合精度 | 13 行源代码 | Tensor Core 利用率为零 |
| 使用 PyTorch 中的多张量优化器 | 1 行源代码 | 优化器中的许多小型 GPU kernel |
图 6：应用的四种简单优化措施。
3. 总结
在生产环境中对 PyTorch 进行性能调优变得越来越重要。一个强大的性能调试工具是此过程的关键。我们通过一个生产模型的案例研究表明，MAIProf 是一个强大的基础设施，用于识别优化机会。
在 Meta，MAIProf 已被数百名工程师使用，从性能新手到专家，用于识别更多类型的瓶颈。这些瓶颈包括慢速数据加载、小型和/或慢速 GPU kernel、分布式训练问题（如负载不平衡和过度通信）。MAIProf 涵盖了主要类别的模型，包括推荐系统、计算机视觉和自然语言处理。总之，它现在是调优生产 PyTorch 工作负载性能不可或缺的工具。
