专栏名称: GiantPandaCV

专注于机器学习、深度学习、计算机视觉、图像处理等多个方向技术分享。团队由一群热爱技术且热衷于分享的小伙伴组成。我们坚持原创，每天一到两篇原创技术分享。希望在传播知识、分享知识的同时能够启发你，大家一起共同进步(･ω<)☆

系统调优助手，PyTorch Profiler TensorBoard 插件教程

GiantPandaCV · 公众号 · · 2024-04-16 12:05

正文

0x1. 前言

使用PyTorch Profiler进行性能分析已经一段时间了，毕竟是PyTorch提供的原生profile工具，个人感觉做系统性能分析时感觉比Nsys更方便一些，并且画的图也比较直观。这里翻译一下PyTorch Profiler TensorBoard Plugin的教程并分享一些使用经验，我使用的时候也是按照这个教程来来的，有一点不一样的是可以在vscode里面直接安装TensorBoard插件，然后Command+Shift+P打开vscode的命令行窗口输入TensorBoard启用TensorBoard插件并把PyTorch Profiler输出的日志文件所在的文件夹路径传给它就可以直接在vscode里面查看可视化Profile结果了。

如果有时候出现下面这种TensorBoard启动超时的情况可以手动杀掉vscode里和tensorboard相关的进程，比如在linux上： pkill -f tensorboard

https://github.com/pytorch/kineto/blob/main/tb_plugin/examples/resnet50_profiler_api.py 和 https://github.com/pytorch/kineto/blob/main/tb_plugin/examples/resnet50_ddp_profiler.py 分别给出了单机和多机的ResNet50模型profile脚本，核心就是下面2个红框部分。

然后，我日常是使用 tensorboard --logdir=./samples --bind_all 这条命令来把可视化结果映射到本地浏览器里面（通过打开：http://localhost:6006/#pytorch_profiler ）。然后我们也要利用好悬浮工具栏，也就是下面红色框部分，当我们点了第三个按钮之后可以通过你可以通过向上拖动鼠标并按住鼠标左键来放大。这个按钮是最常用的。

另外，在可视化界面的右上角有一个Flow Events按钮，有ac2g和fwdbwd两个按钮，前面这个按钮可以让我们获得CPU算子和CUDAKernel的映射关系。比如下图有一个aten::gt算子：

我们可以看到它对应了一个cuda的elementwise算子。

总之很强就是了，相比Nsight System可以直接看到一些分析结果。并且Trace功能非常有用，而Nsight System得肉眼去定位。

除了做训练系统的分析之外，PyTorch Profiler 同样可以用在单个算子或者推理的模型中。

我之后打算聊一些Megatron-LM的细节，其中重要的依据就是使用PyTorch Profiler 的结果，所以这里对PyTorch Profiler TensorBoard Plugin教程做一个翻译，利好初次使用的读者。

最后唠叨一句，PyTorch Profiler在渲染很大的网络的Trace图时需要的时间可能会比较久，以LLama7b为例，为了避免这个问题你可以控制Profile的step为1或者减少Transformer Block的层数为原始的1/4，这样就能很快的得到Trace视图了。如果有需要可以继续浏览0x2节的翻译，或者作为查阅的参考资料。

0x2. 翻译

PyTorch Profiler TensorBoard Plugin 教程地址见：https://github.com/pytorch/kineto/blob/main/tb_plugin/README.md

这是一个用于可视化 PyTorch 性能分析的 TensorBoard 插件。它可以解析、处理并可视化 PyTorch Profiler 导出的性能分析结果，并提供优化建议。

快速安装指南

从 pypi 安装

pip install torch-tb-profiler

若要安装支持 S3 / AzureBlob / GCS / HDFS 扩展的版本，请使用 pip install torch-tb-profiler[s3/blob/gs/hdfs] ，例如 pip install torch-tb-profiler[s3]
或者你可以从源代码安装

克隆 git 仓库：

git clone https://github.com/pytorch/kineto.git

导航到 kineto/tb_plugin 目录。

使用以下命令安装：

pip install .
构建wheel文件

python setup.py build_fe sdist bdist_wheel
** 注意 **：build_fe步骤需要设置 yarn 和 Node.js
python setup.py sdist bdist_wheel

快速启动指南

准备性能分析数据

我们已经准备了一些样本性能分析数据在 kineto/tb_plugin/samples(./samples) 你可以直接下载。或者你可以通过运行 kineto/tb_plugin/examples/resnet50_profiler_api.py(./examples/resnet50_profiler_api.py) 来生成这些样本数据。你还可以从 PyTorch Profiler(https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html?highlight=tensorboard) 学习如何对你的模型进行性能分析并生成性能分析数据。

注意：推荐的生成性能分析数据的方法是在创建 torch.profiler.profile 时，将 torch.profiler.tensorboard_trace_handler 分配给 on_trace_ready 。
启动 TensorBoard

在 TensorBoard 中指定性能分析数据文件夹为 logdir 。如果使用上述样本数据，启动 TensorBoard 命令为：

tensorboard --logdir=./samples

如果你的网络浏览器不在启动 TensorBoard 的同一台机器上，你可以添加 --bind_all 选项，如：

tensorboard --logdir=./samples --bind_all

注意：确保默认端口 6006 对浏览器的主机开放。
在 Chrome 浏览器中打开 TensorBoard

在浏览器中打开 URL http://localhost:6006 。如果你在启动 tensorboard 命令中使用了 --bind_all ，主机名可能不是 'localhost'。你可能会在命令后打印的日志中找到它。
导航到 PYTORCH_PROFILER 标签页

如果 --logdir 下的文件太大或太多，请稍等一会儿并刷新浏览器以查看最新加载的结果。
从云端加载性能分析数据这一堆就不翻译了，基本很难用到，感兴趣的可以查看原文。

快速使用指南

我们将每次启用分析器的运行视为一个“运行”。在大多数情况下，一个运行是一个单独的进程。如果启用了分布式数据并行（DDP），那么一个运行将包括多个进程。我们将每个进程称为一个“工作节点”。

每个运行对应于由 "--logdir" 指定的文件夹下的一个子文件夹。每个子文件夹包含一个或多个 chrome 跟踪文件，每个进程一个。kineto/tb_plugin/samples 是文件组织方式的一个示例。

你可以在TensorBoard左侧控制面板上选择运行的工作节点。

Runs：选择一个运行。每个运行是一次启用性能分析的 PyTorch 应用程序执行。

Views：我们将性能分析结果组织成多个视图，从粗粒度（概览级别）到细粒度（kernel级别）。

Workers：选择一个工作节点。每个工作节点是一个进程。在使用 DDP 时可能有多个工作节点。

Spans：使用 torch.profiler.schedule(https://github.com/pytorch/pytorch/blob/master/torch/profiler/profiler.py#L24) 作为 torch.profiler 的schedule，可能会生成不同跨度的多个性能跟踪文件。你可以用这个选择框选择它们。

目前我们有以下性能诊断视图：

总览视图
操作符视图
kernel 视图
跟踪视图
内存视图
分布式视图

我们将在下面描述每个视图。

总览视图

总览视图是你性能分析运行中过程的顶级视图。它显示了包括主机和 GPU 设备在内的时间成本概览。你可以在左侧面板的“工作节点”下拉菜单中选择当前工作节点。

总览视图的一个示例：

'GPU Summary' 面板显示了此次运行的 GPU 信息和使用指标，包括名称、全局内存、计算能力等。'GPU Utilization'、'Est. SM Efficiency' 和 'Est. Achieved Occupancy' 显示了此次运行在不同级别的 GPU 使用效率。'Kernel Time using Tensor Cores' 显示了 Tensor Core kernel激活的时间百分比。以上四个指标的详细信息可以在 gpu_utilization(./docs/gpu_utilization.md) 中找到。

'Step Time Breakdown' 面板显示了性能概要。我们将每次迭代（通常是一个小批量）视为一个step。每个步骤花费的时间分为以下几类：

kernel ：GPU 设备上的kernel 执行时间；
Memcpy：涉及 GPU 的内存复制时间（D2D、D2H 或 H2D）；
Memset：涉及 GPU 的内存设置时间；
通信：仅在 DDP 情况下出现的通信时间；
运行时：主机端的 CUDA 运行时执行时间；例如 cudaLaunchKernel、cudaMemcpyAsync、cudaStreamSynchronize 等；
DataLoader：在 PyTorch DataLoader 对象中的数据加载时间；
CPU 执行：主机计算时间，包括每个 PyTorch 操作符的运行时间；
其他：未包含在上述任何类别中的时间。

注意：以上所有类别的总结是端到端的实际时间。

上述列表按优先级从高到低排列。我们按优先顺序计算时间。最高优先级类别（Kernel）的时间成本最先计算

，然后是 Memcpy，然后是 Memset，...，最后是其他。在以下示例中，“kernel ”首先被计为 7-2=5 秒；然后“Memcpy”被计为 0 秒，因为它完全被“Kernel”掩盖；然后“CPU 执行”被计为 2-1=1 秒，因为 [2,3] 时间间隔被“kernel ”掩盖，只有 [1,2] 时间间隔被计算。通过这种方式，一个step中所有 7 类别计算的时间总和将与该step的总实际时间相同。

Operator视图

此视图显示了在主机或设备上执行的每个 PyTorch 操作符的性能。

每个表格行都是一个 PyTorch 操作符，这是由 C++ 实现的计算操作符，例如 "aten::relu_"、"aten::convolution"。

调用次数：此运行中操作符被调用的次数。

设备自身持续时间：在 GPU 上累计花费的时间，不包括此操作符的子操作符。

设备总持续时间：在 GPU 上累计花费的时间，包括此操作符的子操作符。

主机自身持续时间：在主机上累计花费的时间，不包括此操作符的子操作符。

主机总持续时间：在主机上累计花费的时间，包括此操作符的子操作符。

Tensor Core适用性：此操作符是否适用于使用Tensor Core。

Tensor Core自身百分比：使用Tensor Core的自身kernel时间 / 自身kernel时间。自身kernel不包括由此操作符的子操作符启动的kernel。

Tensor Core总百分比：使用Tensor Core的kernel时间 / kernel时间。

调用栈：如果已在性能分析跟踪文件中记录，则显示此操作符的所有调用栈。要转储此调用栈信息，应在 torch.profiler API 中设置 'with_stack' 参数。如果在 VSCode 中启动 TensorBoard，点击此调用栈会转到源代码的相应行，如下图：

注意：上述每个持续时间都表示实际时间。这并不意味着在此期间 GPU 或 CPU 完全被利用。

前四个饼图是上述四列持续时间的可视化。它们使得细分在一瞥间就可见。饼图中将仅显示按持续时间排序的前 N 个操作符（在文本框中可配置）。

搜索框允许按名称搜索操作符。

“分组依据”可以选择“操作符”和“操作符 + 输入形状”。“输入形状”是此操作符输入参数列表中的张量形状。空的“[]”表示具有标量类型的参数。例如，“[[32, 256, 14, 14], [1024, 256, 1, 1], [], [], [], [], [], [], []]” 表示此操作符有 9 个输入参数，第一个是尺寸为 32*256*14*14 的张量，第二个是尺寸为 1024*256*1*1 的张量，接下来的七个是标量变量。

性能建议：利用性能分析结果自动突出可能的瓶颈，并给用户提供可行的优化建议。

kernel 视图

此视图显示在 GPU 上所有kernel 的时间花费。时间通过减去kernel 的开始时间和结束时间来计算。

注意：此视图不包括 cudaMemcpy 或 cudaMemset。因为它们不是kernel 。

Tensor Core使用情况：此kernel 是否使用Tensor Core。

总持续时间：此kernel 所有调用的累计时间。

平均持续时间：所有调用的平均时间持续。即“总持续时间”除以“调用次数”。

最大持续时间：所有调用中的最长时间持续。

最小持续时间：所有调用中的最短时间持续。

注意：这些持续时间只包括 GPU 设备上kernel 的经过时间。这并不意味着 GPU 在此时间间隔内忙于执行指令。由于诸如内存访问延迟或并行度不足等原因，一些 GPU 核心可能处于空闲状态。例如，每个 SM 可用的 warps 数量可能不足以有效地隐藏内存访问延迟，或者一些 SM 可能因为块数量不足而完全空闲。请参阅 Nvidia 的最佳实践指南（https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html）。为了研究每个kernel 的效率，我们计算并显示了最后两列的'Mean Blocks Per SM' 和 'Mean Est. Achieved Occupancy' 。

Mean Blocks Per SM：每个 SM 的Block数 = 此kernel的块数 / 此 GPU 的 SM 数量。如果这个数字小于 1，表明 GPU 多处理器未被充分利用。"平均每 SM 块数" 是使用每次运行的持续时间作为权重的所有此kernel名称运行的加权平均值。

Mean Est. Achieved Occupancy：估计实现占用率的定义可以参考 gpu_utilization（https://github.com/pytorch/kineto/blob/main/tb_plugin/docs/gpu_utilization.md），它是使用每次运行的持续时间作为权重的所有此kernel名称运行的加权平均值。

左上角的饼图是“总持续时间”列的可视化。它使得细分在一瞥间就可见。饼图中将仅显示按累计时间排序的前 N 个kernel（在文本框中可配置）。右上角的饼图是使用和未使用Tensor Core的kernel时间百分比。搜索框允许按名称搜索kernel。“分组依据”可以选择“kernel名称”和“kernel属性 + 操作符名称”。“内kernel名称”将按kernel名称分组kernel。“kernel属性 + 操作符名称”将按kernel名称、启动操作符名称、网格、块、每线程寄存器数和共享内存的组合分组kernel。