PyTorch 性能分析（第1部分）：torch.profiler 入门指南

来源：https://huggingface.co/blog/torch-profiler
返回文章列表 (https://huggingface.co/blog)

• 矩阵乘法和加法操作 (https://huggingface.co/blog/torch-profiler#the-matrix-multiplication-and-addition-operation)
• 64x64 跟踪 (https://huggingface.co/blog/torch-profiler#64x64-traces)
  • 为什么 ProfilerStep#2 耗时这么长？ (https://huggingface.co/blog/torch-profiler#why-does-the-profilerstep2-take-so-long)
  • 为什么 CPU 和 GPU 通道之间有约 2.5 毫秒的偏移？ (https://huggingface.co/blog/torch-profiler#why-is-there-an-offset-of-25-ms-between-the-cpu-and-gpu-lanes)
  • 事件链 (https://huggingface.co/blog/torch-profiler#the-chain-of-events)
  • 为什么 matmul 多了一个 CUDA 运行时调用？ (https://huggingface.co/blog/torch-profiler#why-does-matmul-have-an-extra-cuda-runtime-call)
  • 为什么 cudaDeviceSynchronize 这么大（约 1.78 毫秒）？ (https://huggingface.co/blog/torch-profiler#why-is-cudadevicesynchronize-so-big-178-ms)
• 4096x4096 跟踪 (https://huggingface.co/blog/torch-profiler#4096x4096-traces)
  • 为什么同一个内核比其他内核耗时更长？ (https://huggingface.co/blog/torch-profiler#why-does-the-same-kernel-take-more-time-compared-to-others)
• 让我们看看 torch.compile 的实际效果 (https://huggingface.co/blog/torch-profiler#lets-see-some-torch-compile-at-work)
  • 我们是否将 matmul 和 add 内核融合成了一个？ (https://huggingface.co/blog/torch-profiler#did-we-fuse-the-matmul-and-add-kernels-into-one)
  • torch.compile 的运行时架构 (https://huggingface.co/blog/torch-profiler#torchcompiles-runtime-architecture)
  • CUDA 启动次数是否减少了一半？ (https://huggingface.co/blog/torch-profiler#do-the-cuda-launches-go-down-by-half)
  • CPU 开销反而增加了 (https://huggingface.co/blog/torch-profiler#cpu-overhead-went-up-not-down)
• 跟踪速查表 (https://huggingface.co/blog/torch-profiler#trace-reading-cheatsheet)
  • 分析器表格 (https://huggingface.co/blog/torch-profiler#profiler-table)
  • CPU 通道 (https://huggingface.co/blog/torch-profiler#cpu-lane)
  • GPU 通道 (https://huggingface.co/blog/torch-profiler#gpu-lane)
  • 分发链 (https://huggingface.co/blog/torch-profiler#dispatch-chain)
  • torch.compile (https://huggingface.co/blog/torch-profiler#torchcompile)
  • 结论 (https://huggingface.co/blog/torch-profiler#conclusion)

博客缩略图 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-profiler/thumbnail.png)

你无法分析的东西，你也无法优化。

无论你是想从大型语言模型（LLM）中榨取更多每秒令牌数，缩短推理的毫秒级延迟，还是仅仅想理解为什么你的训练循环跑得比规格表承诺的慢，最终都要靠性能分析来解决。问题在于，性能分析的入门门槛很高。跟踪图是密密麻麻的彩色矩形块。事件名称令人望而生畏。大多数教程都假定你已经能读懂它们。因此，即使我们知道应该进行性能分析，打开一个跟踪文件也常常让人感觉像是一件该留到以后（或者交给别人）做的苦差事。

这篇文章，以及它开启的系列，是我们降低这个门槛的尝试。这是 PyTorch 性能分析 系列的开篇，我们将逐步培养阅读性能分析跟踪的能力，并以此驱动优化。

计划如下：

1. 第1部分（本篇）： 从最简单的操作开始——矩阵乘法后跟偏置加法——并学习如何解读分析器返回的结果。
2. 第2部分： 扩展到 nn.Linear 和一个小型 MLP，利用跟踪图推动优化，并窥探底层的 kernels（内核）。
3. 第3部分： 在大型语言模型和 transformers 上综合运用。我们将从初学者的角度记录整个旅程。除了基本的 PyTorch 知识外，无需其他先决条件。请将此视为一次轻松的阅读，途中会有些“啊哈！”时刻。

文章的结构有意采用问题驱动：我们打开一个跟踪，问“等等，为什么这个会发生？”，然后追寻答案，直到有所领悟。最终，你将了解：

• 如何设置 torch.profiler 以及它实际返回什么，
• 如何阅读分析器表格和跟踪图（CPU 通道、GPU 通道以及中间可疑的空隙），
• 从 Python 调用一直到底层 CUDA 内核的事件链，
• 当你加上 torch.compile 后，什么发生了变化（更有趣的是，什么没有变化）。

在开始之前，有两个定义能帮你更好地理解下文：

1. GPU 内核 是一个在 GPU 的许多线程上并行运行的程序。
2. CPU 调度和启动 这些内核。

你通常不需要自己编写 GPU 内核；当你使用 PyTorch 操作时，它会自动翻译成一个或多个在 GPU 上完成工作的内核。

带着这两个概念，让我们开始提问吧。

以下是本文使用的完整脚本：01_matmul_add.py (https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/01_matmul_add.py)。建议在新标签页中打开该脚本，并逐步浏览代码。我们使用 NVIDIA A100-SXM4-80GB GPU 运行脚本。

矩阵乘法和加法操作

正如 Sara Hooker 博士 (https://youtu.be/7knwihgj0fU?si=uvzGH-J9bsCHP4Nn&t=2199) 恰当地指出的那样，正如人类主要由水构成，深度神经网络主要由矩阵乘法构成。作为如此基础的操作，用其他任何操作开始我们的性能分析之旅都将是一种遗憾。

def fn(x, w, b):
    return torch.add(torch.matmul(x, w), b)

矩阵加法与矩阵乘法一起模拟了神经元中权重和偏置的相互作用。这个加法（双关语）（原文为 pun intended，中文取“加法”与“增加理解”的双关）有助于我们理解它如何为本文后面 (https://huggingface.co/blog/torch-profiler#lets-see-some-torch-compile-at-work) 的编译过程铺平道路。

为了进行性能分析，我们将使用 torch.profiler 模块。步骤如下：

1. 准备好要分析的代码 (https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/01_matmul_add.py#L26-L27)（这里是 def fn，它封装了矩阵乘法和矩阵加法）。

2. 对算法进行标注 (https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/01_matmul_add.py#L32)。虽然这是完全可选的，但我们建议这样做。record_function 将我们的函数标注为 matmul_add，这在跟踪图中很容易导航（我们稍后会提到）。

   def step():
       with torch.profiler.record_function("matmul_add"):
           return fn(x, w, b)
   

3. 用 torch.profiler.profile 上下文管理器 (https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/01_matmul_add.py#L53-L62) 包裹代码。

   with torch.profiler.profile(
       activities=[
           torch.profiler.ProfilerActivity.CPU,  # CPU 活动
           torch.profiler.ProfilerActivity.CUDA, # GPU 活动
       ],
   ) as prof:
       # 建议多次运行事件以预热 GPU
       for _ in range(5):
           step()
       prof.step()
   

4. 导出分析结果 (https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/01_matmul_add.py#L70)。

   # 分析器表格
   prof.key_averages().table(sort_by="cuda_time_total", row_limit=15)
   # 分析器跟踪
   prof.export_chrome_trace(trace_path)
   

分析器导出两个不同的产物：

1. 分析器表格： 提供算法的统计摘要。它回答“什么是最耗时的”。这对于找出热点非常有帮助。热点是指耗时最多的事件，可能是管道的瓶颈，也可能是被多次触发的事件。
2. 分析器跟踪： 提供时间执行视图。回答“一个操作在何时以及为何发生”，描绘了 CPU 和 GPU 上的活动。当我们想要调查启动的内核、启动时的任何延迟、CPU 和 GPU 活动之间的任何重叠等情况时，这很有帮助。

让我们在第一次执行中看看这两个产物的实际效果。（这里是完整的 01_matmul_add.py 脚本 (https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/01_matmul_add.py)）

建议在一台带有 GPU 的机器上运行此脚本。

uv run 01_matmul_add.py --size 64

如果你运行上述脚本（在 GPU 机器上），你会在 traces/01_matmul_add 文件夹中找到两个产物：

64_bf16_cold_eager.json
64_bf16_cold_eager.txt

矩阵乘法加法的分析器表格，矩阵大小为 64 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-profiler/profile-table-64.png)
图 1：矩阵大小为 64 的 matmul add 的分析器表格

.txt 文件保存了分析器表格。打开文件后，如图 1 所示，你会看到一个巨大的表格，第一列是 profile 作用域内触发的事件。其他列与事件在 CPU、GPU 或 torch.profiler.profile 中 activities 指定的其他设备上所花费的时间有关。

观察哪些事件耗时最长，并尝试直观理解该事件是否确实应该花那么长时间。注意“调用次数”列也很重要，它显示了事件被触发的次数。

趁此机会，我们也来谈谈“自身 CPU/CUDA”与“CPU/CUDA 总计”的区别。 “自身”列仅衡量事件本身内部花费的时间，不包括其子事件。 “总计”列则包含事件及其所有子事件的总时间。因此，如果你看 matmul_add 的“CPU 总计”，它包含了自身耗时加上它触发的子事件的耗时。这是一个需要注意的重要细微差别。

如果你看表格的最后两行，你会注意到分析器告诉我们：

Self CPU time total: 2.314ms
Self CUDA time total: 23.104us

CPU 时间是毫秒级，而 GPU 时间是微秒级。换算一下，GPU 上花费的时间（内核 ampere_bf16_s16816gemm...）不到 CPU 上花费的时间（matmul_add 操作）的 1%。GPU 大部分时间处于空闲状态，这立即是一个危险信号。发生这种情况的原因是 GPU 可以非常快速地计算一个小型矩阵乘法，因此我们的代码大部分时间花在准备内核、在 GPU 上启动它们、发送要相乘的数据以及收集结果上。这个概念被称为开销受限算法。

摆脱这种状态的最简单方法是使用更大的矩阵乘法。

uv run 01_matmul_add.py --size 4096

矩阵乘法加法算法在 4096 大小矩阵上的分析器表格 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-profiler/profiler-table-4096.png)
图 2：矩阵大小为 4096 的 matmul add 的分析器表格

图 2 的最后两行是：

Self CPU time total: 4.908ms
Self CUDA time total: 4.495ms

两个时间都是毫秒级，这意味着我们仅仅通过增大矩阵乘法的尺寸就实现了更多的 GPU 时间。如果你看图 2，你还会注意到，现在最耗时的 CUDA 时间是 GPU 内核（ampere_bf16_s16816gemm_...），而不是启动它的 CPU 操作（matmul_add）。这意味着我们确实能够从开销受限转向计算受限。

现在，我们进入可视化分发链，它位于 .json 产物中。你可以将它们上传到 Perfetto UI (https://ui.perfetto.dev/) 查看跟踪，或者使用 uvx trace-util traces -b traces 直接生成 Perfetto 链接。

64x64 跟踪

PyTorch 分析器跟踪：在 CUDA GPU 上执行 64×64 bf16 矩阵乘法后接加法 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-profiler/64-matmul-add.png)
图 3：64 大小矩阵的 matmul 和 add 的分析器跟踪

在图 3 中，我们看到了矩阵乘法和加法的分析器跟踪。这里，条的宽度表示事件的持续时间，垂直嵌套是调用层次结构，CPU 通道表示 CPU 上发生的事件，而 GPU 通道显示实际的内核执行。你可能还会注意到空白区域，这是等待或空闲时间。

该脚本使用默认配置运行，即：

• size 64：输入、权重和偏置的大小为 (64, 64)
• dtype bf16：数据类型为 bfloat16
• no compile：我们未编译 torch 操作
• no warmup：我们未在分析前预热 GPU

使用 Perfetto 时，建议使用键盘快速浏览跟踪。可以使用“W A S D”键导航。

PyTorch 分析器跟踪，在 Perfetto 中并排标注了 CPU 通道和 GPU 通道 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-profiler/gpu-cpu-trace.png)
图 4：PyTorch 分析器跟踪的 CPU 和 GPU 通道

图 4 中有两个通道，一个用于 CPU 活动，一个用于 GPU 活动。在 CPU 通道中，你会注意到三个 profile 步骤（从 ProfilerStep#2 开始）。这来自 schedule。

schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)

wait 跳过嘈杂的初始化（ProfilerStep#0），warmup 在分析器中运行但不记录（ProfilerStep#1），而 active 则是跟踪中显示的部分。你可以在脚本这里 (https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/01_matmul_add.py#L58) 找到正在使用的 schedule。

让我们戴上侦探帽，调查跟踪并问一些问题。

为什么 ProfilerStep#2 耗时这么长？

PyTorch 分析器跟踪中的 ProfileStep#2 比 ProfileStep#3 和 ProfileStep#4 看起来更宽 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-profiler/why-is-step-2-big.png)
图 5：ProfileStep#2 明显比后续步骤宽

在图 5 中，我们注意到 ProfileStep#2 比其他步骤花费了更多时间，仔细看，matmul_add 标注也出现了类似模式。烟雾弹在标注内部，而非标注本身：

profile 步骤2中 record_function matmul_add 与 aten::matmul 分发之间约228微秒的间隙 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-profiler/gap-227.png)
图 6：record_function("matmul_add") 和 aten::matmul 之间约 228 μs 的死窗口

图 6 中显示的约 228 μs 是进入 record_function("matmul_add") 与 PyTorch 实际分发 aten::matmul 之间的“死窗口”。这可能由多种原因造成，包括工作空间分配、cuBLAS (https://developer.nvidia.com/cublas)（NVIDIA 专有的 GPU 加速基础线性代数运算库）启发式方法，或延迟模块加载。我们可以忽略它，或者在分析之前运行一些更多的预热步骤 (https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/01_matmul_add.py#L35-L39)（这是标准做法）。

在性能分析中，预热是指在实际分析事件之前，先运行几次事件。GPU 进行的预工作（包括上述几点）是一次性工作，我们不希望将其纳入分析。在我们的示例中，我们有两个预热阶段：一个是在进入分析器之前实际循环调用函数，另一个是