PyTorch Profiler v1.9 现已发布！此次发布（继上一个 PyTorch Profiler 版本）的目标是为您提供最新的前沿工具，以帮助您诊断和修复机器学习性能问题，无论您是在单机还是多机环境下工作。其目的是定位耗时和/或占用内存最多的执行步骤，并可视化 GPU 与 CPU 之间的工作负载分配。
以下是本次发布的五大主要功能摘要：
- 分布式训练视图 (Distributed Training View)：这有助于您了解分布式训练作业中时间和内存的消耗情况。当您将训练模型拆分到工作节点并行运行时，往往会遇到许多问题，因为这就像一个“黑盒”。该视图的整体目标是加速模型训练，并帮助您诊断和调试单个节点内的问题。
- 内存视图 (Memory View)：此视图让您能更直观地了解内存使用情况。该工具通过显示程序运行各阶段的活动内存分配情况，帮助您避免臭名昭著的“内存溢出”(Out of Memory) 错误。
- GPU 利用率可视化 (GPU Utilization Visualization)：此工具可帮助您确保 GPU 得到充分利用。
- 云存储支持 (Cloud Storage Support)：Tensorboard 插件现在可以读取来自 Azure Blob Storage、Amazon S3 和 Google Cloud Platform 的性能分析数据。
- 跳转至源代码 (Jump to Source Code)：此功能允许您可视化堆栈跟踪信息并直接跳转至源代码。这有助于您根据分析结果快速优化和迭代代码。
PyTorch 性能分析工具入门
PyTorch 包含一项名为“PyTorch Profiler”的性能分析功能。PyTorch Profiler 的教程可查看 此处。
要为 PyTorch 代码进行性能分析埋点，您必须：
$ pip install torch-tb-profiler
import torch.profiler as profiler
With profiler.profile(XXXX)
说明：:
• 有关 CUDA 和 CPU 性能分析，请参见 此处
with torch.profiler.profile(
activities=[
torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.CUDA],
• 使用 profiler.record_function(“$NAME”)：允许为一段函数代码添加装饰器（与名称关联的标签）
• profiler.profile 下的 profile_memory=True 参数允许您分析 CPU 和 GPU 的内存占用情况
使用 PyTorch Profiler 可视化 PyTorch 模型性能
分布式训练
深度学习的最新进展强调了大数据集和大模型的重要性，这要求您将模型训练扩展到更多的计算资源上。分布式数据并行 (DDP) 和 NVIDIA 集体通信库 (NCCL) 是 PyTorch 中加速深度学习训练的广泛采用的范式。
在本次发布的 PyTorch Profiler 中，已支持带有 NCCL 后端的 DDP。
计算/通信概览
在分布式训练视图下的“计算/通信概览”中，您可以观察每个工作节点及工作节点间[负载均衡](https://en.wikipedia.org/wiki/Load_balancing_(computing)节点的计算与通信比率（按粒度衡量）。
场景 1:
如果某个工作节点的计算和重叠时间远大于其他节点，这可能表明负载均衡存在问题，或者该节点是“拖后腿”节点（straggler）。计算时间是指 GPU 上的内核总时间减去重叠时间。重叠时间是指通过在计算过程中交错进行通信所节省的时间。重叠时间越多，代表计算与通信之间的并行度越好。理想情况下，计算和通信完全重叠。通信时间是指总通信时间减去重叠时间。下图展示了此场景在 Tensorboard 中的显示方式。
图：拖后腿节点示例
场景 2:
如果批次大小 (batch size) 较小（即每个节点的计算量较小）或者传输的数据量较大，计算与通信的比率可能会很小，在性能分析器中表现为 GPU 利用率低和等待时间长。此计算/通信视图将允许您诊断代码，通过采用梯度累积来减少通信，或通过增加批次大小来降低通信占比。DDP 通信时间取决于模型大小，而批次大小与模型大小无关。因此，增加批次大小可以延长计算时间，从而提高计算与通信的比率。
同步/通信概览
在“同步/通信”视图中，您可以观察通信的效率。这是通过用“步骤时间”减去“计算时间和通信时间”来计算的。同步时间是总通信时间的一部分，用于等待并与其他工作节点进行同步。同步/通信视图包括初始化、数据加载器、CPU 计算等。您可以从该视图中得出诸如“真正用于交换数据的通信比例是多少”以及“等待其他节点数据时的空闲时间是多少”等见解。
例如，如果存在负载均衡效率低下或拖后腿的问题，您将能够在“同步/通信”视图中识别出来。该视图会显示某些工作节点的等待时间比其他节点长。
上方的表格视图允许您查看每个节点中所有通信操作的详细统计信息。这使您可以了解调用了哪些操作类型、每个操作被调用了多少次、每个操作传输的数据大小等。
内存视图：
此内存视图工具可帮助您了解模型中算子（operators）的硬件资源消耗情况。了解算子级别的计算和内存消耗，使您能够解决性能瓶颈，进而让模型执行得更快。鉴于 GPU 内存有限，优化内存使用可以：
- 支持更大的模型，从而在最终任务中可能获得更好的泛化能力。
- 支持更大的批次大小。更大的批次大小可以提高训练速度。
性能分析器会记录分析间隔内的所有内存分配。选择“设备”(Device) 后，您可以查看每个算子在 GPU 端或主机端的内存使用情况。您必须启用 profile_memory=True 才能生成如下所示的内存数据，详情见此处。
With torch.profiler.profile(
Profiler_memory=True # this will take 1 – 2 minutes to complete.
)
重要定义:
• “Size Increase”（大小增量）：显示所有分配字节的总和减去所有内存释放字节。
• “Allocation Size”（分配大小）：显示所有分配字节的总和，不考虑内存释放。
• “Self”（自身）：表示分配的内存不是来自任何子算子，而是由算子本身分配的。
时间轴上的 GPU 指标：
此功能将帮助您在发现一个或多个 GPU 利用率不足时调试性能问题。理想情况下，您的程序应该具有高 GPU 利用率（目标为 100%）、最少的 CPU 到 GPU 通信且无开销。
概览 (Overview)：概览页面重点展示了三个重要 GPU 使用指标在不同层级的结果（即 GPU 利用率、预估 SM 效率和预估已实现占用率）。本质上，每个 GPU 有许多流多处理器 (SM)，每个 SM 有许多线程束 (warps)，可以并发执行大量线程。线程束的数量取决于具体 GPU 型号。总体而言，时间轴上的 GPU 指标工具允许您查看整个堆栈，这非常有帮助。
如果 GPU 利用率结果较低，这表明模型中可能存在瓶颈。常见原因包括：
• 内核中的并行性不足（即批次大小过小）
• 循环调用了微小的内核。这意味着启动开销未能被摊销
• CPU 或 I/O 瓶颈导致 GPU 没有收到足够的任务来保持繁忙状态
查看概览页面中的性能建议部分，您会发现有关如何提高 GPU 利用率的潜在建议。在此示例中，GPU 利用率很低，因此性能建议是增加批次大小。根据性能建议，将批次大小从 4 增加到 32 后，GPU 利用率提高了 60.68%。
GPU 利用率：指性能分析器中 GPU 引擎执行工作负载的步骤间隔时间。利用率百分比越高越好。仅使用 GPU 利用率来诊断性能瓶颈的缺点是它太高级且粗糙。它无法告知您使用了多少个流多处理器。请注意，虽然此指标对于检测空闲时段很有用，但高值并不表示 GPU 的使用效率高，仅仅表示它确实在运行任务。例如，一个持续运行单个线程的内核将获得 100% 的 GPU 利用率。
预估流多处理器效率 (Est. SM Efficiency) 是一个更细粒度的指标，它指示在追踪过程的任何时刻使用了百分之多少的 SM。该指标报告了至少有一个活跃线程束在 SM 上运行的时间百分比，以及那些处于停滞状态的时间（NVIDIA 文档）。预估 SM 效率也有其局限性。例如，每个块只有一个线程的内核无法完全利用每个 SM。SM 效率并没有告诉我们每个 SM 有多忙，只告诉我们它们是否在工作，这可能包括在等待内存加载结果时停滞。为了保持 SM 繁忙，必须有足够数量的就绪线程束，以便在发生停滞时能够运行。
预估已实现占用率 (Est. Achieved Occupancy) 在诊断性能问题时比预估 SM 效率和 GPU 利用率更深一层。预估已实现占用率指示每个 SM 可以同时有多少个活跃线程束。拥有足够数量的活跃线程束通常是实现高吞吐量的关键。与 GPU 利用率和 SM 效率不同，将此值设为尽可能高并不是目标。根据经验，通过将该指标提高到 15% 及以上可以获得良好的吞吐量提升。但在某个点之后，收益会递减。例如，如果该值已经达到 30%，进一步提高收益是不确定的。此指标报告了内核执行期间所有线程束调度器的平均值（NVIDIA 文档）。预估已实现占用率值越大越好。
概览详情：Resnet50_batchsize4
概览详情：Resnet50_batchsize32
内核视图 (Kernel View)：内核具有“每 SM 块数”和“预估已实现占用率”，这是比较模型运行效果的绝佳工具。
平均每 SM 块数 (Mean Blocks per SM)
每 SM 块数 = 此内核的块数 / 此 GPU 的 SM 数量。如果此数字小于 1，则表明 GPU 多处理器未被充分利用。“平均每 SM 块数”是该内核名称所有运行的加权平均值，使用每次运行的持续时间作为权重。
平均预估已实现占用率 (Mean Est. Achieved Occupancy)
预估已实现占用率定义如上所述。“平均预估已实现占用率”是该内核名称所有运行的加权平均值，使用每次运行的持续时间作为权重。
追踪视图 (Trace View)：此追踪视图显示了一个时间轴，显示了模型中算子的持续时间以及哪个系统执行了该操作。此视图可以帮助您确定高消耗和长时间执行是因为输入还是模型训练。目前，此追踪视图在时间轴上显示 GPU 利用率和预估 SM 效率。
GPU 利用率是独立计算的，并划分为多个 10 毫秒的桶。桶的 GPU 利用率值绘制在时间轴上，范围在 0 – 100% 之间。在上面的例子中，线程 28022 繁忙期间的“ProfilerStep5”GPU 利用率高于“Optimizer.step”期间的利用率。这就是您可以放大调查原因的地方。
从上面我们可以看出，前者的内核比后者的内核长。后者的内核执行时间太短，导致 GPU 利用率较低。
预估 SM 效率：每个内核都有一个计算出的 0 – 100% 之间的预估 SM 效率。例如，下面的内核只有 64 个块，而该 GPU 中的 SM 有 80 个。那么它的“预估 SM 效率”是 64/80，即 0.8。
云存储支持
在运行 pip install tensorboard 后，若要通过这些云提供商读取数据，您可以运行：
torch-tb-profiler[blob]
torch-tb-profiler[gs]
torch-tb-profiler[s3]
pip install torch-tb-profiler[blob], pip install torch-tb-profiler[gs], 或 pip install torch-tb-profiler[S3] 来通过这些云提供商读取数据。有关更多信息，请参阅此 README。
跳转至源代码：
将 TensorBoard 和 PyTorch Profiler 直接集成到 Visual Studio Code (VS Code) 中的一大好处是，能够直接从分析器堆栈跟踪跳转到源代码（文件和行）。VS Code Python 扩展现在支持 TensorBoard 集成。
跳转至源代码功能仅在 VS Code 内启动 Tensorboard 时可用。如果在性能分析时设置了 with_stack=True，则堆栈跟踪将出现在插件 UI 上。当您点击 PyTorch Profiler 中的堆栈跟踪时，VS Code 将自动并排打开相应的文件，并直接跳转到您感兴趣的代码行进行调试。这使您可以根据性能分析结果和建议，快速对代码进行可操作的优化和更改。
GIF：使用 VS Code 插件 UI 跳转至源代码
有关如何优化批次大小性能，请查看此处的分步教程。PyTorch Profiler 也已与 PyTorch Lightning 集成，您只需通过 --trainer.profiler=pytorch 标志启动 Lightning 训练作业即可生成追踪信息。
PyTorch Profiler 的下一步计划是什么？
您刚刚了解了 PyTorch Profiler 如何帮助优化模型。现在可以通过 pip install torch-tb-profiler 尝试使用该分析器来优化您的 PyTorch 模型。
敬请期待未来推出的该教程的高级版本。我们也衷心希望能持续为 PyTorch 用户提供先进的工具以提高机器学习性能。我们非常乐意听取您的建议。欢迎在此此处提交 issue。
有关 PyTorch Profiler 即将推出的新功能，请关注 Twitter 上的 @PyTorch 并查看 pytorch.org。
致谢
作者感谢以下人员对本文的贡献。Facebook 团队：Geeta Chauhan, Gisle Dankel, Woo Kim, Sam Farahzad, 和 Mark Saroufim。微软团队：AI 框架工程师 (Teng Gao, Mike Guo, 和 Yang Gu), Guoliang Hua, 和 Thuy Nguyen。