1. Meta 的 AI 性能分析工具 (MAIProf)
图 1：Meta AI 性能分析 (MAIProf) 基础设施的简化示意图。
图 1 展示了 Meta AI 性能分析基础设施的简化示意图。机器学习研究人员和性能工程师通过用户门户（User Portal）向分析服务（Profiling Service）提交针对训练作业的分析请求，该服务随后将请求广播到运行训练作业的所有 GPU 主机。当 GPU 主机上的监控守护进程（Monitoring Daemon）收到分析请求时，它会通知对应训练作业 PyTorch 程序内的 Kineto GPU 追踪器（基于 NVIDIA 的 libcupti 构建）。结果，Kineto 追踪数据会被异步采集并上传到对象存储（Object Store）中（详细说明：每个 GPU 都会采集一个 Kineto 追踪文件，每个文件都被视为一个 blob 进行处理和存储；第 2 节将提供示例）。与此同时，MAIProf 还会收集各种聚合性能指标：每台 GPU 主机上的监控守护进程会持续从 NVIDIA 的 DCGM/NVML 读取性能计数器，并将其记录到时间序列数据库（Time Series DB）中。
一旦追踪数据和指标采集完成，分析服务将自动从对象存储下载追踪数据进行分析，并从时间序列数据库下载性能指标进行度量分析。最后，一份包含详尽且深刻分析的整体性能分析报告将交付给用户。
为了满足生产环境的使用需求，我们为 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 性能瓶颈
第一步：
在同一时间轴上检查 CPU 和 GPU 的利用率，如图 2 所示。
图 2：CPU 使用率随时间的变化（上）与 GPU 使用率随时间的变化（下）。
我们在图 2 中注意到的第一个性能异常是这种模式：“GPU 空闲、GPU 活动、GPU 空闲、GPU 活动……”贯穿整个训练过程。总体而言，GPU 在超过一半的训练时间内处于空闲状态（这对性能不利，因为 GPU 是高性能设备，我们希望尽可能提高其利用率）。
第二步：
使用 MAIProf 在 GPU 空闲时收集 CPU 上的 Python 函数调用追踪，如图 3 所示。
图 3：Python 调用追踪。
Python 追踪显示，大部分 CPU 时间消耗在 Python 函数 sharded_iterrows() 内部。从模型的源代码中，我们了解到该函数并行处理一个大型特征表。使用的工作线程数由可配置参数（num_worker_threads）控制。此外，在调查了特征表的生成方式后，我们理解了这一性能异常：训练数据集太大，无法一次性全部放入 CPU 内存；它需要被拆分成多个子数据集，每个子数据集包含运行 10 个 epoch 所需的足够数据。因此，每 10 个 epoch 就需要从磁盘读取一个新的子数据集到内存中，在此期间 GPU 完全处于空闲状态。
第三步：
收集 GPU 性能指标，如图 4 所示。
图 4：MAIProf 中的 GPU 性能指标。
我们从图 4 中得出了以下观察结果：
- 流式多处理器（SM）负责运行模型的 CUDA 内核。其利用率 [1] 为 9.1%，表明 GPU 上的并行计算单元未得到充分利用。
- Tensor Core 利用率为 0，意味着 Tensor Core（GPU 上的混合精度计算单元）[2] 完全未被使用。
- 最大 GPU 内存利用率为 47.13%，表明有一半的 GPU 内存处于闲置状态。
第四步：
收集训练循环的 GPU 追踪（即 Kineto 追踪），如图 5 所示。
图 5：训练循环的 GPU 追踪（即 Kineto 追踪）。
由于常用的 PyTorch 函数已经过注释，它们的名称会自动显示在追踪数据上。通过它们，我们可以将追踪大致划分为训练迭代中的四个阶段：（1）数据加载，（2）前向传播，（3）反向传播，（4）梯度优化（注意：图 5 中的“优化器”阶段来自上一个 batch，而其他三个阶段来自当前 batch）。
2.2 优化
我们针对上述识别出的瓶颈执行了四项简单的优化，每一项只需更改配置参数或最多几行源代码。它们列于图 6 中。
| 优化 | 更改量 | 解决的瓶颈 |
|---|---|---|
通过尝试每台主机 CPU 核心数范围内的几个可能值来调整 num_worker_threads。 | 1 行源代码 | GPU 完全空闲时间 |
| 加倍 batch size | 2 个配置参数 | GPU 内存利用率不足 |
| 在 PyTorch 中使用 自动混合精度 (AMP) | 13 行源代码 | Tensor Core 利用率为零 |
| 在 PyTorch 中使用 多张量优化器 (Multitensor optimizer) | 1 行源代码 | 优化器中有许多小型 GPU 内核 |
图 6：已应用的四项简单优化。
3. 结语
生产环境下 PyTorch 的性能调优正变得越来越重要。一个强大的性能调试工具是此过程的关键。我们通过一个生产模型案例研究证明，MAIProf 是一个识别优化机会的强大基础设施。
在 Meta，MAIProf 已被数百名工程师使用，从性能新手到专家，他们用它识别了更多类型的瓶颈。这些瓶颈包括数据加载缓慢、GPU 内核过小或过慢，以及分布式训练中诸如负载不均衡和过度通信等问题。MAIProf 覆盖了推荐系统、视觉和自然语言处理等主流模型类别。总之，它现在已成为调整生产环境 PyTorch 工作负载性能不可或缺的工具。