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:
如果一个工作节点的计算和重叠时间远大于其他节点，这可能暗示工作负载平衡存在问题，或者该节点运行缓慢。计算时间等于 GPU 上的内核总耗时减去重叠时间。重叠时间是指在计算过程中插入通信所节省的时间。重叠时间越长，代表计算和通信之间的并行性越好。理想情况下，计算和通信完全重叠。通信时间等于总通信时间减去重叠时间。下图展示了此场景在 Tensorboard 中的呈现方式。
图：运行缓慢（straggler）示例
场景 2:
如果批量大小（Batch size）较小（即每个工作节点的计算量较少）或者传输的数据量较大，计算与通信的比率也可能很小，这会在性能分析器中表现为 GPU 利用率低和等待时间长。此计算/通信视图将允许您诊断代码，通过采用梯度累积来减少通信，或者通过增加批量大小来降低通信占比。DDP 通信时间取决于模型大小，而批量大小与模型大小无关。因此，增加批量大小可以延长计算时间，从而提高计算与通信的比率。
同步/通信概览
在同步/通信视图中，您可以观察通信的效率。计算方法是用步长时间减去计算时间和通信时间。同步时间是等待并与其他工作节点同步的总通信时间的一部分。同步/通信视图包括初始化、数据加载器、CPU 计算等。通过此视图可以洞察总通信中实际用于交换数据的比率，以及等待其他节点数据时的空闲时间。
例如，如果存在工作负载不平衡或节点运行缓慢的问题，您可以在此同步/通信视图中识别出来。该视图将显示某些工作节点的等待时间比其他节点更长。
上方的表格视图允许您查看每个节点中所有通信操作的详细统计信息。这使您可以查看调用的操作类型、每个操作被调用的次数、每个操作传输的数据大小等。
内存视图：
此内存视图工具可帮助您了解模型中算子的硬件资源消耗。了解算子级别的计算时间和内存消耗，有助于解决性能瓶颈，进而使模型执行得更快。鉴于 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 都有多个 Warp，可以并发执行多个线程。Warp 执行的数量取决于 GPU 类型。但从高层来看，此“时间轴上的 GPU 指标”工具允许您查看整个堆栈，这非常有用。
如果 GPU 利用率结果较低，这表明模型中可能存在瓶颈。常见原因包括：
•内核并行度不足（即批量大小较小）
•循环中调用了小型内核。这意味着启动开销无法被分摊
•CPU 或 I/O 瓶颈导致 GPU 没有收到足够的任务来保持繁忙
查看概览页面中的性能建议部分，您会找到关于如何提高 GPU 利用率的建议。在此示例中，GPU 利用率很低，因此性能建议是增加批量大小。按照建议将批量大小从 4 增加到 32 后，GPU 利用率提高了 60.68%。
GPU 利用率：性能分析器中 GPU 引擎执行工作负载时的步长时间间隔。利用率百分比越高越好。仅使用 GPU 利用率来诊断性能瓶颈的缺点是它太高级且粗略。它无法告诉您正在使用多少个流式多处理器。请注意，虽然此指标对于检测空闲周期很有用，但高值并不表示 GPU 利用效率高，只表示它在进行某种操作。例如，一个持续运行单个线程的内核将获得 100% 的 GPU 利用率。
预估流式多处理器效率 (Est. SM Efficiency) 是一个更精细的指标，它表示在轨迹的任何时刻正在使用多少百分比的 SM。此指标报告了至少有一个活跃 Warp 在 SM 上执行且处于停顿状态的时间百分比（NVIDIA 文档）。预估 SM 效率也有局限性。例如，每个块只有一个线程的内核无法充分利用每个 SM。SM 效率并没有告诉我们每个 SM 有多忙，只表明它们在做事情，这可能包括在等待内存加载结果时停顿。为了让 SM 保持繁忙，必须有足够数量的就绪 Warp，以便在发生停顿时可以运行。
预估已实现占用率 (Est. Achieved Occupancy) 在诊断性能问题时比预估 SM 效率和 GPU 利用率更深一层。预估已实现占用率表示每个 SM 同时可以有多少个活跃 Warp。拥有足够数量的活跃 Warp 通常是实现良好吞吐量的关键。与 GPU 利用率和 SM 效率不同，让该值尽可能高并不是目标。根据经验，将此指标提高到 15% 以上通常能获得不错的吞吐量提升。但在某一点之后，收益会递减。例如，如果该值已经达到 30%，进一步提高收益就不确定了。此指标报告了内核执行期间所有 Warp 调度器的平均值（NVIDIA 文档）。预估已实现占用率值越大越好。
概览详情：Resnet50_batchsize4
概览详情：Resnet50_batchsize32
内核视图 (Kernel View) 内核具有“每个 SM 的块数 (Blocks per SM)”和“预估已实现占用率”，这是比较模型运行情况的绝佳工具。
每个 SM 的平均块数
每个 SM 的块数 = 此内核的块数 / 此 GPU 的 SM 数量。如果此数字小于 1，则表示 GPU 多处理器未得到充分利用。“每个 SM 的平均块数”是此内核名称所有运行的加权平均值，以每次运行的持续时间作为权重。
平均预估已实现占用率
预估已实现占用率的定义如概览所述。“平均预估已实现占用率”是此内核名称所有运行的加权平均值，以每次运行的持续时间作为权重。
追踪视图 (Trace View) 此追踪视图显示了一个时间轴，展示了模型中算子的持续时间以及执行该操作的系统。此视图可以帮助您识别高消耗和长执行时间是由输入还是模型训练引起的。目前，此追踪视图在时间轴上显示 GPU 利用率和预估 SM 效率。
GPU 利用率是独立计算并划分为多个 10 毫秒的桶。桶的 GPU 利用率值显示在时间轴旁边，范围为 0-100%。在上面的示例中，线程 28022 繁忙期间的“ProfilerStep5”GPU 利用率高于“Optimizer.step”期间的利用率。在这里，您可以放大以调查原因。
从上面我们可以看到，前者的内核比后者的内核长。后者的内核执行时间太短，导致 GPU 利用率较低。
预估 SM 效率：每个内核都有一个预估 SM 效率，在 0-100% 之间计算。例如，下面的内核只有 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 将自动并排打开相应文件，并直接跳转到您感兴趣的代码行进行调试。这使您可以根据性能分析结果和建议，快速对代码进行可操作的优化和更改。
动图：使用 Visual Studio Code 插件 UI 跳转至源代码
关于如何优化批量大小性能，请查看此分步教程。PyTorch Profiler 也与 PyTorch Lightning 集成，您可以简单地通过 –trainer.profiler=pytorch 标志启动 Lightning 训练作业以生成追踪数据。
PyTorch Profiler 的下一步计划是什么？
您刚刚了解了 PyTorch Profiler 如何帮助优化模型。您现在可以通过 pip install torch-tb-profiler 试用 Profiler 来优化您的 PyTorch 模型。
敬请期待未来推出的本教程的高级版本。我们也非常激动能继续为 PyTorch 用户提供前沿工具，以提高 ML 性能。我们期待收到您的反馈。欢迎随时在此开具 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。