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