Hugging Face:Blog(RSS)
精选
71AI 编辑部评分,满分 100

PyTorch 性能分析系列(一):torch.profiler 入门指南

2026-05-29 08:00· 66天前
跳到正文
精选理由

PyTorch profiling 的陡峭学习曲线劝退了很多人,这篇用从零开始的方式把 trace 拆解得明明白白,想做性能优化的同学该收藏。

AI 摘要

本文是 PyTorch profiling 系列的开篇,从最简单的矩阵乘法加偏置操作出发,逐步讲解如何使用 torch.profiler 进行性能分析。涵盖 profiler 设置、导出统计表格与 Chrome trace、解读 CPU 和 GPU 活动的时序关系,以及 torch.compile 对底层 CUDA kernel 调用链的影响。实验基于 NVIDIA A100-SXM4-80GB GPU 运行,面向基本掌握 PyTorch 但缺乏 profiling 经验的读者。

正文 · AI 翻译

无法剖析,就无法优化。

无论你是想从大语言模型中榨取更多每秒 token 数、将推理时间缩短几毫秒,还是只想弄明白为什么你的训练循环跑得比规格表承诺的慢,最终都绕不开性能剖析。

问题在于,性能剖析的上手门槛很高。追踪视图是一堵由彩色矩形构成的密墙。事件名称令人望而生畏。大多数教程都假设你已经能读懂它们。所以,即使我们知道应该做性能剖析,打开一个追踪文件也常常感觉像是一件最好留到以后(或者留给别人)做的苦差事。本文,以及它所开启的系列文章,就是我们降低这个门槛的尝试。

这是《PyTorch 性能剖析》系列的开篇之作,在这个系列中,我们将逐步培养阅读性能剖析器追踪记录的能力,并以此驱动优化:

  1. PyTorch 性能剖析(第一部分):torch.profiler 入门指南(当前)
  2. PyTorch 性能剖析(第二部分):从 nn.Linear 到融合 MLP
  3. PyTorch 性能剖析(第三部分):注意力机制即一切可剖析之物

我们从初学者的视角记录这段旅程。除了基础的 PyTorch 知识外,无需任何先决条件。请把它当作一次轻松的阅读,其中穿插着一些“啊哈!”时刻。本文的结构有意采用问题驱动:我们打开一个追踪文件,问“等等,为什么会这样?”,然后追寻答案,直到豁然开朗。读完本文,你应该能了解:

  • 如何设置 torch.profiler 以及它实际返回什么,
  • 如何阅读性能剖析器表格和追踪视图(CPU 通道、GPU 通道,以及两者之间可疑的间隙),
  • 从一次 Python 调用一路向下到 CUDA 内核的事件链,
  • 当你加上 torch.compile 后,什么发生了变化(以及更有趣的是,什么没有变化)。

在开始之前,有两个定义能让下面的内容读起来更顺畅:

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

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

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

以下是我们在文章中使用的完整脚本:01_matmul_add.py。建议在新标签页中打开此脚本,并逐步浏览代码。我们使用 NVIDIA A100-SXM4-80GB GPU 来运行这些脚本。在 Hugging Face 基础设施上设置 GPU 非常容易,并且可以通过 Spaces 的开发者模式来试验这些脚本。也可以使用 Hugging Face Jobs 流水线来运行这些脚本。

矩阵乘法与加法运算

正如 Sara Hooker 博士精辟指出的那样,就像我们主要由水构成一样,深度神经网络主要由矩阵乘法构成。鉴于其基础性地位,我们的性能分析之旅从其他内容开始将是一种遗憾。

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

矩阵加法与矩阵乘法共同模拟了神经元中权重和偏置的交互方式。这个加法(双关语)将帮助我们理解它如何为后文中的编译环节铺平道路。

为了进行性能分析,我们将使用 torch.profiler 模块。涉及的步骤如下:

  1. 准备好待分析的代码(此处为 `def fn`,它封装了矩阵乘法和矩阵加法)
  2. 对算法进行标注。虽然这完全可选,但我们建议这样做。`record_function` 将我们的函数标注为 `matmul_add`,这将便于在追踪结果中导航(后文会提到)。
def step():
  with torch.profiler.record_function("matmul_add"):
    return fn(x, w, b)
  1. 将代码包裹在 `torch.profiler.profile` 上下文管理器内
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,  # the cpu activities
        torch.profiler.ProfilerActivity.CUDA, # the gpu activities
    ],
  ) as prof:
    # it is recommended to run events multiple times to warm up the GPUs
    for _ in range(5):
      step()
      prof.step()
  1. 导出分析结果
# the profiler table
prof.key_averages().table(sort_by="cuda_time_total", row_limit=15)

# the profiler trace
prof.export_chrome_trace(trace_path)

性能分析器会导出两种不同的产物:

  1. 分析器表格:提供算法的统计摘要。它回答了“什么花费时间最多”这个问题。这对于找出热点非常有帮助。热点是指耗时最长的事件,可能是流水线的瓶颈,也可能是被频繁触发的事件。
  2. 分析器追踪结果:提供时间维度的执行视图。它回答了“操作在何时以及为何发生”,展示了 CPU 和 GPU 上的活动情况。当我们想要调查被启动的内核、启动过程中的任何延迟、CPU 与 GPU 活动之间的重叠等情况时,这非常有用。

让我们通过第一次执行来看看这两者的实际效果。(这里是完整的 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
Image 13: Profiler table for matmul add on 64 sized matrices
图 1:针对 64 大小矩阵的矩阵乘法加法的性能分析器表格

`.txt` 文件保存了性能分析器表格。打开该文件后,如图 1 所示,你会看到一个大型表格,第一列包含在 `profile` 作用域内触发的事件。

其他列则与事件在 CPU、GPU 或 `torch.profiler.profile` 的 `activities` 中指定的任何其他设备上所花费的时间有关。查看哪些事件花费的时间最多,并尝试凭直觉判断该事件是否确实应该花费那么多时间。同样重要的是查看“# of Calls”列,它指明了事件被触发的次数。

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

如果你查看表格的最后两行,会发现性能分析器告诉我们:

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
Image 14: Profiler table for matmul add algorithm on 4096 sized matrices
图 2:针对 4096 大小矩阵的矩阵乘法加法的性能分析器表格

图 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 查看追踪信息,也可以使用 `uvx trace-util -f traces -b <hf_uname>/traces` 直接生成 Perfetto 链接。

64x64 追踪

Image 15: PyTorch profiler trace of a 64×64 bf16 matmul followed by an add on a CUDA GPU
图 3:64 尺寸矩阵上的矩阵乘法和加法性能分析器追踪

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

该脚本使用默认配置运行,具体如下:

  • 尺寸 64:输入、权重和偏置的尺寸均为 (64, 64)
  • 数据类型 bf16:数据类型为 bfloat16
  • 未编译:我们未对 torch 操作进行编译
  • 未预热:我们在性能分析前未对 GPU 进行预热

使用 Perfetto 时,我们建议使用键盘以便更快地浏览追踪信息。可以使用“W A S D”键来导航追踪。

Image 16: PyTorch profiler trace with the CPU lane and GPU lane labelled side by side in Perfetto
图 4:PyTorch 性能分析器追踪的 CPU 和 GPU 通道

图 4 中有两个通道,一个用于 CPU 活动,一个用于 GPU 活动。在 CPU 通道中,你会注意到三个分析步骤(从 ProfilerStep#2 开始)。这源于调度设置。

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

等待步骤会跳过嘈杂的初始化过程(ProfilerStep#0),预热步骤会在不记录的情况下运行分析器(ProfilerStep#1),而活跃步骤则是在追踪中显示的内容。你可以在此处找到脚本中使用的调度设置。

让我们戴上侦探帽,仔细研究一下这个追踪,并提出一些问题。

为什么 ProfilerStep#2 花费了这么长时间?

Image 17: ProfileStep#2 in a PyTorch profiler trace appears wider than ProfileStep#3 and ProfileStep#4
图 5:ProfileStep#2 明显比其后的步骤更宽

在图 5 中,我们注意到 ProfileStep#2 相比其他步骤耗时更长,仔细观察会发现 matmul_add 标注也呈现类似模式。问题根源在于标注内部,而非标注本身:

步骤 matmul_add 开始 aten::matmul 开始 间隔
#2 138.736 366.493 227.757 微秒
#3 517.926 523.447 5.521 微秒
#4 610.039 614.527 4.488 微秒
Image 18: 228 microsecond gap between record_function matmul_add and the aten::matmul dispatch in profile step 2
图 6:record_function("matmul_add") 与 aten::matmul 之间约 228 微秒的死窗口

图 6 中显示的约 228 微秒,是进入 record_function("matmul_add") 与 PyTorch 实际调度 aten::matmul 之间的"死窗口"。造成这一现象的原因有多种,包括工作空间分配、cuBLAS(NVIDIA 专有的 GPU 加速基础线性代数运算库)启发式策略,或延迟模块加载。我们可以选择忽略它,也可以在性能分析前多运行几次预热步骤(这是标准做法)。

在性能分析中,预热是指在实际分析之前先将事件运行若干次。GPU 完成的预处理工作(包括上述几点)是一次性开销,我们并不希望将其纳入分析。在我们的示例中,有两个预热阶段:一是在进入分析器之前实际循环执行函数,二是在分析器内部通过 warmup 参数实现的预热。在本节中,我们已结合调度启用了实际迭代。

uv run 01_matmul_add.py --warmup
Image 19: PyTorch profiler trace after warmup steps where ProfileStep#2 no longer shows cold-start overhead
图 7:预热后,每个分析步骤耗时相近

在图 7 中,我们看到每个分析步骤耗时相近,但这并不意味着我们优化了单次开销。我们只是对运行进行了预热,使得这些开销未被纳入分析。我们认为,如果不给出解决此问题的提示就草草结束本节,对读者是不公平的,因此这里提供一个链接,供进一步了解如何优化启动开销。

为什么 CPU 和 GPU 通道之间存在约 2.5 毫秒的偏移?

Image 20: 2.32 millisecond offset between the CPU lane and the GPU lane in a PyTorch profiler trace
图 8:CPU 和 GPU 通道之间约 2.5 毫秒的偏移

在图 8 中,我们看到 CPU 和 GPU 通道之间存在约 2.5 毫秒的偏移:这是 CPU 提交 CUDA 内核后到它们实际开始执行之间的延迟。有人可能会认为,预热阶段加上调度器的等待和预热应该能让 GPU 保持忙碌,从而减小这个偏移。

为了揭示真正发生了什么,让我们稍微调整一下调度器:

- schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
+ schedule = torch.profiler.schedule(wait=0, warmup=0, active=3, repeat=1)
Image 21: PyTorch profiler trace with wait=0 warmup=0 showing an Activity Buffer Request between steps
图 9:当 wait=0 且 warmup=0 时,追踪记录显示了一个活动缓冲区请求

图 9 向我们展示,在任何操作之前,GPU 通道中就存在一个活动缓冲区请求。让我们再放大一些来看。

Image 22: gap between matmul and add CUDA kernels caused by profiler buffer request
图 10:在 profile 步骤 1 上,matmul 和 add 内核之间出现了一个间隙

放大 GPU 追踪记录后,我们注意到 ProfileStep#0(其 CPU 追踪记录在该图中不可见)的 matmul 和 add 内核是连续发生的,而 ProfileStep#1 的内核之间则有一个时间窗口。对此最好的解释是,发生了缓冲区溢出,在内核执行期间又发出了另一个缓冲区请求(即在 GPU 显存上分配一些内存的请求)。

排除其他可能性的最佳方法是,对更多迭代进行性能分析,看看追踪记录的其他部分是否也出现类似的时间窗口。为此,我们设置 active=20 来运行。

Image 23: PyTorch profiler trace of 20 active iterations confirming the buffer-request gap only appears once
图 11:在 20 个 active 步骤下,该间隙只出现了一次,证实了这是一个缓冲区请求

如图 11 所示,我们在 ProfileStep#1 中观察到了类似的趋势。这与我们之前的发现一致,因此可以安全地得出结论,这确实是另一个缓冲区请求。

事件链

Image 24: nested CPU dispatch chain in PyTorch profiler: ProfileStep, matmul_add, aten::matmul, aten::mm
图 12:调度链

在图 12 中,我们看到了嵌套的 CPU 调用。这是一个重要的可视化结果,通过它我们可以理解调度链的真实面貌。

我们从 ProfileStep#<id> 开始,它封装了性能分析步骤。由于我们对步骤进行了标注,因此可以看到 matmul_add 行。matmul_add 由两个 aten 调用组成,一个用于矩阵乘法,一个用于矩阵加法。

aten::matmul 是 ATen 级别的调度,用户层面的 PyTorch matmul 调用最终都会落到这里。aten::mm 是二维矩阵-矩阵乘法的后端。

非常有趣的是,如果我们为矩阵添加批次维度,PyTorch 会调用 aten::bmm(批量矩阵乘法)。让我们绕个弯,看看 aten::bmm 的实际运行情况。

- x = torch.randn(args.size, args.size, device=device, dtype=dtype)
- w = torch.randn( args.size, args.size, device=device, dtype=dtype)
- b = torch.randn(args.size, args.size, device=device, dtype=dtype)

+ # adding a batch size of 8
+ x = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
+ w = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
+ b = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
Image 25: PyTorch profiler trace showing aten::matmul dispatching aten::bmm for 3D batched tensors
图 13:批量矩阵乘法

在图 13 中,当为输入添加批次维度后,aten::matmul 现在除了调用 aten::bmm(而非 aten::mm)之外,还封装了一系列其他必要的 CUDA 运行时调用。这也暗示了 cuBLAS 为了调度出最适合该程序的内核而需要执行的启发式策略。

在本文的其余部分,除非另有说明,我们将使用简单的二维矩阵。

为什么 matmul 会多出一个额外的 CUDA 运行时调用?

Image 26: CPU lane showing cudaOccupancyMaxActiveBlocksPerMultiprocessor preceding the matmul cudaLaunchKernel
图 14:在 matmul 内核启动之前触发了一次 CUDA 占用率查询

我们注意到,对于 aten::mm 有两个 CUDA 运行时调用,即 cudaOccupancyMaxActiveBlocksPerMultiprocessor(在图 14 中用方框标出)和 cudaLaunchKernel,而对于 aten::add 则只有 cudaLaunchKernel。

cudaOccupancyMaxActiveBlocksPerMultiprocessor 是一个规划性调用,完全在 CPU 端执行。它询问的是:“给定一个内核函数、一个选定的块大小和一个选定的动态共享内存大小,这个内核的多少个块可以同时驻留在一个 SM(流式多处理器)上?”

这就引出了一个问题:为什么 matmul 需要规划,而 add 不需要?

要理解这一点,我们必须查看内核的资源占用情况。如果你点击 GPU 内核,就可以检查相应内核的资源占用情况。

Image 27: cuBLAS matmul kernel resource footprint: registers, shared memory and block size in Perfetto Image 28: elementwise add CUDA kernel resource footprint with 32 registers and zero shared memory
图 15:Matmul 资源占用 图 16:Add 资源占用

在图 15 中,我们注意到对于矩阵乘法,每个线程的寄存器和共享内存是动态的(取决于矩阵的大小)。cuBLAS 提供了数百种内核变体,每种变体都有一个由启发式驱动的启动路径,该路径需要关于硬件容量的运行时信息。占用率查询就是该启发式策略的一部分。从概念上讲,我们可以将 GPU 加速的矩阵乘法理解为在独立的图块上工作:我们使用多少个图块以及每个图块需要多大,取决于矩阵和硬件。现代算法远比这复杂,但这仍然是一个很好的参考框架。

从图 16 可以看出,加法运算的占用空间显示为 32 个寄存器和零共享内存。这完全可以轻松容纳。无需查询任何信息,因为没有硬件资源会限制占用率。该内核在设计上就是资源轻量型的。

在阅读任何追踪信息时,你都可以将此作为快速诊断手段。扫描 CPU 通道,查找 `cudaOccupancyMaxActiveBlocksPerMultiprocessor`。每次出现都标记着一个“重量级、自适应启动”的内核,通常是 GEMM(通用矩阵乘法)、卷积或类似运算。而那些前面没有占用率查询的内核,则是 PyTorch 机械式启动的逐元素/规约类内核。

为什么 `cudaDeviceSynchronize` 耗时如此之长(约 1.78 毫秒)?

`cudaDeviceSynchronize` 会阻塞 CPU,直到该设备上的所有 GPU 工作完成。性能分析器会在活动窗口结束时发出此同步调用,以刷新事件。没有它,内核计时将会缺失。

一个耗时 1.78 毫秒的同步调用,覆盖了仅 26 微秒的实际 GPU 工作,这表明此次运行有 98% 的时间处于空闲状态。这是典型的开销瓶颈症状。

4096x4096 追踪信息

从上面的性能分析器表格分析中我们已经知道,为算法提供更大的矩阵,会使其从开销受限区域转变为计算受限区域。

让我们运行命令,更深入地探究这些追踪信息。

uv run 01_matmul_add.py --size 4096 --warmup

为什么同一个内核比其他内核耗时更长?

Image 29: 4096x4096 bf16 matmul kernel timings varying across profiler steps on the same GPU
图 17:在输入完全相同的情况下,一个矩阵乘法内核的运行时间比其他内核更长

在图 17 中,我们注意到 ProfileStep#3 的矩阵乘法内核在 GPU 上的耗时比其他步骤更长。这一点特别值得注意,因为启动的其他内核完全相同,这意味着没有涉及 cuBLAS 启发式方法。没有调度间隙,CPU 启动正常,并且这不是性能分析器的伪影。

图 17 中的这个追踪信息揭示了一个在理想化示例中容易被忽略的重要观点:即使在相同的硬件环境下,对相同的数据运行相同的代码,内核运行时间也不是恒定的。

让我们稍微修改脚本,使这一点更加具体。我们将迭代运行 20 次,并捕获每一步的数据。

- schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
+ schedule = torch.profiler.schedule(wait=0, warmup=0, active=20, repeat=1)

- for _ in range(5):
+ for _ in range(20):
Image 30: PyTorch profiler trace of 20 matmul iterations showing kernel runtime variance
图 18:在 20 次迭代中,同一个矩阵乘法内核以不同的速度运行

图 18 揭示了类似的发现。虽然每个内核完全相同,但它们的执行时间却不同。计算时间的差异可归因于多种原因:

  • GPU 时钟的空闲与升频
  • GPU 发热
  • GPU 电源管理
  • 驱动端的后台维护

只看平均值的读者会得出结论:一次矩阵乘法耗时约 1 毫秒(5 次均值 = 1084 微秒);而查看追踪记录的读者则会发现,矩阵乘法耗时约 580 微秒,除非 GPU 出现异常。这两种认知模型截然不同,且只有一种是正确的。

让我们看看 torch.compile 的实际运行

使用 torch.compile 一直让我们惊叹不已。开发者编写普通的即时执行 PyTorch 代码,但 PyTorch 会尝试捕获张量密集区域,将其转化为计算图,进行优化,并运行生成的代码。默认后端通常是 TorchInductor,其大致流程如下:

  1. TorchDynamo 将 Python 执行过程捕获为 FX 计算图
  2. 当涉及梯度时,AOTAutograd 准备前向/反向计算图
  3. Inductor 将计算图降级为优化的 CPU 或 GPU 代码。

在本节中,我们将讨论编译过程并查看性能分析器追踪记录。

uv run 01_matmul_add.py --size 4096 --warmup --compile

args.compile 标志会触发以下代码:

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

fn = torch.compile(fn) if args.compile else fn
Image 31: torch.compile region highlighted in a PyTorch profiler trace, showing TorchDynamo and Inductor frames
图 19:编译后的区域在追踪记录中显示为 TorchDynamo 和 Inductor 帧

在图 19 中,我们看到新增的 CPU 行被命名为 Torch-Compiled Region: 0/0,这指向了正在使用的编译函数。

我们是否将矩阵乘法和加法内核融合成了一个?

Image 32: Compiled trace showing aten::addmm replacing the eager aten::add and aten::mm pair
图 20:编译后的运行分派了一个单一的 aten::addmm

查看图 20,我们提出一个问题:我们是否真的将乘法和加法运算融合成了一个?

这是计算图层面的算子融合。Inductor 将我们的 torch.add(torch.matmul(x, w), b) 重写为单一的 aten::addmm(b, x, w) 调用。这里需要注意的重要一点是,它并没有生成一个新的融合 CUDA 内核。实际的 GPU 运算仍然是 ampere_bf16_s16816gemm_bf16_128x256_ldg8_f2f_stages_64x3_nn,与即时执行模式使用的 cuBLAS 内核相同。因此,这里的"融合"发生在分派器层面,而非内核层面。

PyTorch 提供了 `torch.addmm` 函数,它能一步完成我们之前分两步做的事情,即乘法和加法。我们鼓励读者查看该函数的追踪记录,并在下方评论区分享你们的观察结果!

torch.compile 的运行时架构

虽然我们在理论上知道编译函数时会发生什么,但亲眼看到它的实际运行同样重要。让我们来看看反映 torch.compile 运行时架构的 CPU 端层次结构。

TorchDynamo 缓存查找是 Dynamo 检查当前调用是否仍然与已编译内容(相同的输入形状、数据类型、设备和张量元数据)匹配的地方。如果任何内容不匹配,Dynamo 将重新编译。即使在编译之后,每次调用都会产生这一开销。

Torch-Compiled Region 是"进入"编译版本的包装器。AOTDispatcher Runtime Wrapper Prologue 是 AOT Autograd 的运行时包装器。即使我们在这里不需要梯度,AOTDispatcher 始终存在于堆栈中,负责处理张量元数据、视图追踪,并且如果 `requires_grad` 为真,它还会设置反向传播。

## Call CompiledFxGraph 是实际生成的代码运行的地方。"CompiledFxGraph" 后面的字符串是 FX 图的内容哈希值。它在所有三个活跃步骤中都是相同的,这确认了缓存命中。

你可以在磁盘上的 `/tmp/torchinductor_<user>/fxgraph` 目录下找到以此哈希值为键的生成代码,当你想要阅读 Inductor 实际生成的 Triton/C++ 代码时,这非常有用。

CUDA 启动次数是否减少了一半?

Image 33: compiled matmul trace showing Memcpy DtoD and GEMM kernels launched per step
图 21:每个编译后的步骤仍然启动两个 GPU 内核,一个设备到设备的内存拷贝和 GEMM

查看图 21 中的追踪记录,我们非常高兴地注意到每个步骤只有一个 `cudaLaunchKernel`。这一观察结果直接与我们之前在 GPU 追踪中看到的情况相矛盾。每个步骤仍然有两个内核被启动,即 Memcpy DtoD(设备到设备)和 GEMM。回到 CPU 追踪记录,我们发现自己完全漏掉了 `cudaMemcpyAsync` 的调度。

addmm 计算的是 out = α·A·B + β·C,而 cuBLAS 的带偏置加法(bias-add)后处理(epilogue)会写入一个目标缓冲区,该缓冲区需要已包含偏置值。后处理可以理解为 GEMM 之后发生的所有操作。在深度学习领域,我们不断提出各种 GEMM 后处理,例如激活函数、偏置加法、归一化等等。正因如此,才有了各种 cuBLAS GEMM 带后处理的内核变体。

如果你对 torch.compile 使用不同的模式,你会注意到启动的内核变体也不同。你可以亲自尝试一下,并在下方评论中分享你的观察结果!

因此,Inductor 生成的代码执行以下操作:

  • out = copy(C) ← 这是设备到设备(DtoD)的内存拷贝(32 MB,耗时约 33 µs)
  • out = α·(A·B) + β·out ← GEMM,其中 α=β=1,将偏置加法融合到回写操作中

从数学上讲,结果仍然相同。偏置加法并非没有代价,因为我们预先支付了一次内存拷贝,外加一次开销稍高的 GEMM 后处理。

人们可能期望的融合——即 x·w + b(此处为 out = α·A·B + β·C)能坍缩成一个内核,且没有额外的内存流量——并未发生。Inductor 保留了两次内存访问操作,只是将偏置拷贝重新标记为内存拷贝,将加法重新标记为 GEMM 后处理。

真正融合的实现会跳过内存拷贝。这正是 FlashAttention 风格的手写内核所做的,也是 Inductor 通过 Triton 代码生成所能做到的,但对于一个 4096×4096 的 bf16 矩阵乘法,Inductor 显然判定“使用 cuBLAS,通过后处理设置来完成偏置加法”是最佳路径。

CPU 开销增加了,而非减少

这是比较即时执行(eager)和编译执行时最容易忽略的一点:

步骤 即时执行耗时(毫秒) 编译执行耗时(毫秒)
#2 0.1 0.2
#3 0.07 0.1
#4 0.07 0.1

torch.compile 在 CPU 上每一步的开销大约是前者的 2 倍。这是因为每次调用都会遍历完整的 Dynamo > AOTAutograd > Inductor 堆栈,除此之外还有我们无论如何都会进行的 aten::addmm 分发。编译流水线是为包含数十个操作的机器学习模型设计的,在这些模型中,每次调用的开销可以被分摊(对于单个操作来说,这就像是一种额外税负)。

torch.compile 有一个 mode 参数。这留给读者作为课后作业,请查阅文档并提出一种能够降低 CPU 开销的模式。🤗

跟踪读取速查表

我们梳理过的模式速查表。思路是:如果在追踪记录中看到这个,它通常意味着什么。

性能分析器表格

你看到的内容 它通常意味着什么
自身 CPU 总时间 ≫ 自身 CUDA 总时间(CPU 以毫秒计,GPU 以微秒计) 开销受限。CPU 花在调度上的时间比 GPU 花在计算上的时间还多。应增大工作规模(更大的矩阵、批处理操作)或融合调用。
自身 CPU 总时间 ≈ 自身 CUDA 总时间,两者均以毫秒计 计算受限。GPU 是瓶颈,这通常是你希望看到的情况。
单个事件主导了 CUDA 总时间 那就是你的热点。从那里开始优化。
单个事件的调用次数巨大 即使每次调用开销很小,也可能是潜在瓶颈。检查是否可以融合或批处理。
某行的 CPU 总时间 ≫ 自身 CPU 时间 大部分开销在子事件中。应深入分析嵌套事件,而非父事件。

CPU 时间线

你看到的内容 它通常意味着什么
第一个 ProfileStep 比其余宽得多 冷启动开销:工作空间分配、cuBLAS 启发式算法、惰性模块加载。增加预热迭代次数和/或调度器的预热参数。
record_function("...") 开始与内部第一个 aten::* 之间存在较大间隔 同样是冷启动开销,只是放大了看。注解已进入,但调度尚未发生。
cudaLaunchKernel 之前的 cudaOccupancyMaxActiveBlocksPerMultiprocessor 一个重量级、自适应启动的内核(GEMM、卷积等)。cuBLAS 正在向驱动查询一个 SM 上能容纳多少个块,以便选择内核变体。
没有前置占用查询的 cudaLaunchKernel 一个元素级或规约内核,具有固定且资源占用小的特征。无需规划。
活动窗口末尾出现较长的 cudaDeviceSynchronize 性能分析器正在刷新事件。其持续时间主要是 GPU 完成待处理工作的时间,并非真正的 CPU 开销。覆盖微小 GPU 工作的同步是典型开销受限症状。
你并未编写的 cudaMemcpyAsync 通常是隐藏的设备到设备拷贝。常见于 addmm 在 GEMM 后处理阶段之前,用偏置值初始化其目标缓冲区时。

GPU 时间线

你看到的内容 它通常意味着什么
GPU 时间线上的活动缓冲区请求 性能分析器正在分配/重新填充其自身的事件缓冲区。第一个通常解释了初始的 CPU↔GPU 时间线偏移。
单步中两个内核之间存在间隙 很可能在执行过程中又发起了一次缓冲区请求。通过运行更多迭代来确认:如果它只出现一次,那就是性能分析器的问题,而不是你的代码。
同一个内核在不同步骤中的计时结果不同 GPU 时钟、温度、电源管理、驱动维护。要读取追踪数据,而不仅仅是平均值。
一个名为 ampere_bf16_s16816gemm_... 的内核 矩阵乘法实际的 cuBLAS GPU 工作。对于相同的形状/数据类型,该内核名称在即时执行模式和编译模式下通常是相同的。
在 GEMM 操作之前立即出现 Memcpy DtoD addmm 后置操作的偏置拷贝。这里的“融合”发生在调度器层面,而非内核层面。

调度链

你所看到的 它通常的含义
ProfileStep#N → <record_function name> → aten::* → aten::mm / aten::bmm / aten::add 标准的嵌套调用层级。自身时间不包括子项;总时间则包括子项。
aten::matmul 解析为 aten::mm 二维 × 二维矩阵乘法。
aten::matmul 解析为 aten::bmm(带有额外的 CUDA 运行时调用) 对三维及以上张量执行的批量矩阵乘法。cuBLAS 会执行更多启发式工作来选择变体。
aten::addmm(b, x, w) 替代了单独的 aten::add + aten::mm 组合 调度器层面的算子融合。GPU 内核仍然是相同的 GEMM,只是偏置加法被折叠进了后置操作中。

torch.compile

你所看到的 它通常的含义
Torch-Compiled 区域:CPU 通道中的 K/M 行 你正处于一个编译后的函数内部。
每一步都进行 TorchDynamo 缓存查找 Dynamo 正在验证形状/数据类型/设备是否与缓存的编译结果匹配。每次调用都会产生此开销,即使在编译完成后也是如此。
即使没有梯度,也存在 AOTDispatcher 运行时包装器前导 AOTAutograd 的运行时包装器始终存在于调用栈中,负责处理张量元数据和视图追踪。
## 每一步都使用相同的哈希值调用 CompiledFxGraph <hash> 命中了生成代码的缓存。生成的源代码位于 /tmp/torchinductor_<user>/fxgraph/<hash> 目录下。
对于微小的算子,使用 torch.compile 时每步的 CPU 时间高于即时执行模式 这是预期的。Dynamo → AOTAutograd → Inductor 这一调用栈本身是一种开销,只有在处理大量算子时才能被摊平。

结论

我们从一个小小的矩阵乘法加加法运算开始,并以此为切入点,学习如何阅读 PyTorch 性能分析器。在此过程中,我们掌握了一些思维模型,这些模型同样适用于更大的工作负载。这是“PyTorch 性能分析”系列的第一站。在后续的文章中,我们将逐步告别这个双算子小玩具,沿着复杂度的阶梯向上攀登,审视更大的构建模块,并最终分析真实模型。

感谢 Noe Flandre、Suvaditya Mukherjee 和 Vidit Ostwal 对本文初稿的审阅!

这篇博文使用大语言模型进行了润色。这绝不意味着我们让一个智能体在后台运行并让它生成整篇博客。我们团队中有些成员并非英语母语者,我们认为大语言模型(大多基于英语语料训练)可以纠正一些简单的语法错误,或改写那些听起来不够简洁、不够清晰的句子。希望这能有助于理解“如果这是大语言模型生成的,我为什么还要读”这个想法。🤗

PyTorch 性能分析系列(一):torch.profiler 入门指南

Hugging Face:Blog(RSS)·2026-05-29 08:00·66天前
阅读原文· huggingface.co
精选理由

PyTorch profiling 的陡峭学习曲线劝退了很多人,这篇用从零开始的方式把 trace 拆解得明明白白,想做性能优化的同学该收藏。

AI 摘要

本文是 PyTorch profiling 系列的开篇,从最简单的矩阵乘法加偏置操作出发,逐步讲解如何使用 torch.profiler 进行性能分析。涵盖 profiler 设置、导出统计表格与 Chrome trace、解读 CPU 和 GPU 活动的时序关系,以及 torch.compile 对底层 CUDA kernel 调用链的影响。实验基于 NVIDIA A100-SXM4-80GB GPU 运行,面向基本掌握 PyTorch 但缺乏 profiling 经验的读者。

正文 · AI 翻译

无法剖析,就无法优化。

无论你是想从大语言模型中榨取更多每秒 token 数、将推理时间缩短几毫秒,还是只想弄明白为什么你的训练循环跑得比规格表承诺的慢,最终都绕不开性能剖析。

问题在于,性能剖析的上手门槛很高。追踪视图是一堵由彩色矩形构成的密墙。事件名称令人望而生畏。大多数教程都假设你已经能读懂它们。所以,即使我们知道应该做性能剖析,打开一个追踪文件也常常感觉像是一件最好留到以后(或者留给别人)做的苦差事。本文,以及它所开启的系列文章,就是我们降低这个门槛的尝试。

这是《PyTorch 性能剖析》系列的开篇之作,在这个系列中,我们将逐步培养阅读性能剖析器追踪记录的能力,并以此驱动优化:

  1. PyTorch 性能剖析(第一部分):torch.profiler 入门指南(当前)
  2. PyTorch 性能剖析(第二部分):从 nn.Linear 到融合 MLP
  3. PyTorch 性能剖析(第三部分):注意力机制即一切可剖析之物

我们从初学者的视角记录这段旅程。除了基础的 PyTorch 知识外,无需任何先决条件。请把它当作一次轻松的阅读,其中穿插着一些“啊哈!”时刻。本文的结构有意采用问题驱动:我们打开一个追踪文件,问“等等,为什么会这样?”,然后追寻答案,直到豁然开朗。读完本文,你应该能了解:

  • 如何设置 torch.profiler 以及它实际返回什么,
  • 如何阅读性能剖析器表格和追踪视图(CPU 通道、GPU 通道,以及两者之间可疑的间隙),
  • 从一次 Python 调用一路向下到 CUDA 内核的事件链,
  • 当你加上 torch.compile 后,什么发生了变化(以及更有趣的是,什么没有变化)。

在开始之前,有两个定义能让下面的内容读起来更顺畅:

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

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

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

以下是我们在文章中使用的完整脚本:01_matmul_add.py。建议在新标签页中打开此脚本,并逐步浏览代码。我们使用 NVIDIA A100-SXM4-80GB GPU 来运行这些脚本。在 Hugging Face 基础设施上设置 GPU 非常容易,并且可以通过 Spaces 的开发者模式来试验这些脚本。也可以使用 Hugging Face Jobs 流水线来运行这些脚本。

矩阵乘法与加法运算

正如 Sara Hooker 博士精辟指出的那样,就像我们主要由水构成一样,深度神经网络主要由矩阵乘法构成。鉴于其基础性地位,我们的性能分析之旅从其他内容开始将是一种遗憾。

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

矩阵加法与矩阵乘法共同模拟了神经元中权重和偏置的交互方式。这个加法(双关语)将帮助我们理解它如何为后文中的编译环节铺平道路。

为了进行性能分析,我们将使用 torch.profiler 模块。涉及的步骤如下:

  1. 准备好待分析的代码(此处为 `def fn`,它封装了矩阵乘法和矩阵加法)
  2. 对算法进行标注。虽然这完全可选,但我们建议这样做。`record_function` 将我们的函数标注为 `matmul_add`,这将便于在追踪结果中导航(后文会提到)。
def step():
  with torch.profiler.record_function("matmul_add"):
    return fn(x, w, b)
  1. 将代码包裹在 `torch.profiler.profile` 上下文管理器内
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,  # the cpu activities
        torch.profiler.ProfilerActivity.CUDA, # the gpu activities
    ],
  ) as prof:
    # it is recommended to run events multiple times to warm up the GPUs
    for _ in range(5):
      step()
      prof.step()
  1. 导出分析结果
# the profiler table
prof.key_averages().table(sort_by="cuda_time_total", row_limit=15)

# the profiler trace
prof.export_chrome_trace(trace_path)

性能分析器会导出两种不同的产物:

  1. 分析器表格:提供算法的统计摘要。它回答了“什么花费时间最多”这个问题。这对于找出热点非常有帮助。热点是指耗时最长的事件,可能是流水线的瓶颈,也可能是被频繁触发的事件。
  2. 分析器追踪结果:提供时间维度的执行视图。它回答了“操作在何时以及为何发生”,展示了 CPU 和 GPU 上的活动情况。当我们想要调查被启动的内核、启动过程中的任何延迟、CPU 与 GPU 活动之间的重叠等情况时,这非常有用。

让我们通过第一次执行来看看这两者的实际效果。(这里是完整的 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
Image 13: Profiler table for matmul add on 64 sized matrices
图 1:针对 64 大小矩阵的矩阵乘法加法的性能分析器表格

`.txt` 文件保存了性能分析器表格。打开该文件后,如图 1 所示,你会看到一个大型表格,第一列包含在 `profile` 作用域内触发的事件。

其他列则与事件在 CPU、GPU 或 `torch.profiler.profile` 的 `activities` 中指定的任何其他设备上所花费的时间有关。查看哪些事件花费的时间最多,并尝试凭直觉判断该事件是否确实应该花费那么多时间。同样重要的是查看“# of Calls”列,它指明了事件被触发的次数。

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

如果你查看表格的最后两行,会发现性能分析器告诉我们:

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
Image 14: Profiler table for matmul add algorithm on 4096 sized matrices
图 2:针对 4096 大小矩阵的矩阵乘法加法的性能分析器表格

图 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 查看追踪信息,也可以使用 `uvx trace-util -f traces -b <hf_uname>/traces` 直接生成 Perfetto 链接。

64x64 追踪

Image 15: PyTorch profiler trace of a 64×64 bf16 matmul followed by an add on a CUDA GPU
图 3:64 尺寸矩阵上的矩阵乘法和加法性能分析器追踪

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

该脚本使用默认配置运行,具体如下:

  • 尺寸 64:输入、权重和偏置的尺寸均为 (64, 64)
  • 数据类型 bf16:数据类型为 bfloat16
  • 未编译:我们未对 torch 操作进行编译
  • 未预热:我们在性能分析前未对 GPU 进行预热

使用 Perfetto 时,我们建议使用键盘以便更快地浏览追踪信息。可以使用“W A S D”键来导航追踪。

Image 16: PyTorch profiler trace with the CPU lane and GPU lane labelled side by side in Perfetto
图 4:PyTorch 性能分析器追踪的 CPU 和 GPU 通道

图 4 中有两个通道,一个用于 CPU 活动,一个用于 GPU 活动。在 CPU 通道中,你会注意到三个分析步骤(从 ProfilerStep#2 开始)。这源于调度设置。

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

等待步骤会跳过嘈杂的初始化过程(ProfilerStep#0),预热步骤会在不记录的情况下运行分析器(ProfilerStep#1),而活跃步骤则是在追踪中显示的内容。你可以在此处找到脚本中使用的调度设置。

让我们戴上侦探帽,仔细研究一下这个追踪,并提出一些问题。

为什么 ProfilerStep#2 花费了这么长时间?

Image 17: ProfileStep#2 in a PyTorch profiler trace appears wider than ProfileStep#3 and ProfileStep#4
图 5:ProfileStep#2 明显比其后的步骤更宽

在图 5 中,我们注意到 ProfileStep#2 相比其他步骤耗时更长,仔细观察会发现 matmul_add 标注也呈现类似模式。问题根源在于标注内部,而非标注本身:

步骤 matmul_add 开始 aten::matmul 开始 间隔
#2 138.736 366.493 227.757 微秒
#3 517.926 523.447 5.521 微秒
#4 610.039 614.527 4.488 微秒
Image 18: 228 microsecond gap between record_function matmul_add and the aten::matmul dispatch in profile step 2
图 6:record_function("matmul_add") 与 aten::matmul 之间约 228 微秒的死窗口

图 6 中显示的约 228 微秒,是进入 record_function("matmul_add") 与 PyTorch 实际调度 aten::matmul 之间的"死窗口"。造成这一现象的原因有多种,包括工作空间分配、cuBLAS(NVIDIA 专有的 GPU 加速基础线性代数运算库)启发式策略,或延迟模块加载。我们可以选择忽略它,也可以在性能分析前多运行几次预热步骤(这是标准做法)。

在性能分析中,预热是指在实际分析之前先将事件运行若干次。GPU 完成的预处理工作(包括上述几点)是一次性开销,我们并不希望将其纳入分析。在我们的示例中,有两个预热阶段:一是在进入分析器之前实际循环执行函数,二是在分析器内部通过 warmup 参数实现的预热。在本节中,我们已结合调度启用了实际迭代。

uv run 01_matmul_add.py --warmup
Image 19: PyTorch profiler trace after warmup steps where ProfileStep#2 no longer shows cold-start overhead
图 7:预热后,每个分析步骤耗时相近

在图 7 中,我们看到每个分析步骤耗时相近,但这并不意味着我们优化了单次开销。我们只是对运行进行了预热,使得这些开销未被纳入分析。我们认为,如果不给出解决此问题的提示就草草结束本节,对读者是不公平的,因此这里提供一个链接,供进一步了解如何优化启动开销。

为什么 CPU 和 GPU 通道之间存在约 2.5 毫秒的偏移?

Image 20: 2.32 millisecond offset between the CPU lane and the GPU lane in a PyTorch profiler trace
图 8:CPU 和 GPU 通道之间约 2.5 毫秒的偏移

在图 8 中,我们看到 CPU 和 GPU 通道之间存在约 2.5 毫秒的偏移:这是 CPU 提交 CUDA 内核后到它们实际开始执行之间的延迟。有人可能会认为,预热阶段加上调度器的等待和预热应该能让 GPU 保持忙碌,从而减小这个偏移。

为了揭示真正发生了什么,让我们稍微调整一下调度器:

- schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
+ schedule = torch.profiler.schedule(wait=0, warmup=0, active=3, repeat=1)
Image 21: PyTorch profiler trace with wait=0 warmup=0 showing an Activity Buffer Request between steps
图 9:当 wait=0 且 warmup=0 时,追踪记录显示了一个活动缓冲区请求

图 9 向我们展示,在任何操作之前,GPU 通道中就存在一个活动缓冲区请求。让我们再放大一些来看。

Image 22: gap between matmul and add CUDA kernels caused by profiler buffer request
图 10:在 profile 步骤 1 上,matmul 和 add 内核之间出现了一个间隙

放大 GPU 追踪记录后,我们注意到 ProfileStep#0(其 CPU 追踪记录在该图中不可见)的 matmul 和 add 内核是连续发生的,而 ProfileStep#1 的内核之间则有一个时间窗口。对此最好的解释是,发生了缓冲区溢出,在内核执行期间又发出了另一个缓冲区请求(即在 GPU 显存上分配一些内存的请求)。

排除其他可能性的最佳方法是,对更多迭代进行性能分析,看看追踪记录的其他部分是否也出现类似的时间窗口。为此,我们设置 active=20 来运行。

Image 23: PyTorch profiler trace of 20 active iterations confirming the buffer-request gap only appears once
图 11:在 20 个 active 步骤下,该间隙只出现了一次,证实了这是一个缓冲区请求

如图 11 所示,我们在 ProfileStep#1 中观察到了类似的趋势。这与我们之前的发现一致,因此可以安全地得出结论,这确实是另一个缓冲区请求。

事件链

Image 24: nested CPU dispatch chain in PyTorch profiler: ProfileStep, matmul_add, aten::matmul, aten::mm
图 12:调度链

在图 12 中,我们看到了嵌套的 CPU 调用。这是一个重要的可视化结果,通过它我们可以理解调度链的真实面貌。

我们从 ProfileStep#<id> 开始,它封装了性能分析步骤。由于我们对步骤进行了标注,因此可以看到 matmul_add 行。matmul_add 由两个 aten 调用组成,一个用于矩阵乘法,一个用于矩阵加法。

aten::matmul 是 ATen 级别的调度,用户层面的 PyTorch matmul 调用最终都会落到这里。aten::mm 是二维矩阵-矩阵乘法的后端。

非常有趣的是,如果我们为矩阵添加批次维度,PyTorch 会调用 aten::bmm(批量矩阵乘法)。让我们绕个弯,看看 aten::bmm 的实际运行情况。

- x = torch.randn(args.size, args.size, device=device, dtype=dtype)
- w = torch.randn( args.size, args.size, device=device, dtype=dtype)
- b = torch.randn(args.size, args.size, device=device, dtype=dtype)

+ # adding a batch size of 8
+ x = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
+ w = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
+ b = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
Image 25: PyTorch profiler trace showing aten::matmul dispatching aten::bmm for 3D batched tensors
图 13:批量矩阵乘法

在图 13 中,当为输入添加批次维度后,aten::matmul 现在除了调用 aten::bmm(而非 aten::mm)之外,还封装了一系列其他必要的 CUDA 运行时调用。这也暗示了 cuBLAS 为了调度出最适合该程序的内核而需要执行的启发式策略。

在本文的其余部分,除非另有说明,我们将使用简单的二维矩阵。

为什么 matmul 会多出一个额外的 CUDA 运行时调用?

Image 26: CPU lane showing cudaOccupancyMaxActiveBlocksPerMultiprocessor preceding the matmul cudaLaunchKernel
图 14:在 matmul 内核启动之前触发了一次 CUDA 占用率查询

我们注意到,对于 aten::mm 有两个 CUDA 运行时调用,即 cudaOccupancyMaxActiveBlocksPerMultiprocessor(在图 14 中用方框标出)和 cudaLaunchKernel,而对于 aten::add 则只有 cudaLaunchKernel。

cudaOccupancyMaxActiveBlocksPerMultiprocessor 是一个规划性调用,完全在 CPU 端执行。它询问的是:“给定一个内核函数、一个选定的块大小和一个选定的动态共享内存大小,这个内核的多少个块可以同时驻留在一个 SM(流式多处理器)上?”

这就引出了一个问题:为什么 matmul 需要规划,而 add 不需要?

要理解这一点,我们必须查看内核的资源占用情况。如果你点击 GPU 内核,就可以检查相应内核的资源占用情况。

Image 27: cuBLAS matmul kernel resource footprint: registers, shared memory and block size in Perfetto Image 28: elementwise add CUDA kernel resource footprint with 32 registers and zero shared memory
图 15:Matmul 资源占用 图 16:Add 资源占用

在图 15 中,我们注意到对于矩阵乘法,每个线程的寄存器和共享内存是动态的(取决于矩阵的大小)。cuBLAS 提供了数百种内核变体,每种变体都有一个由启发式驱动的启动路径,该路径需要关于硬件容量的运行时信息。占用率查询就是该启发式策略的一部分。从概念上讲,我们可以将 GPU 加速的矩阵乘法理解为在独立的图块上工作:我们使用多少个图块以及每个图块需要多大,取决于矩阵和硬件。现代算法远比这复杂,但这仍然是一个很好的参考框架。

从图 16 可以看出,加法运算的占用空间显示为 32 个寄存器和零共享内存。这完全可以轻松容纳。无需查询任何信息,因为没有硬件资源会限制占用率。该内核在设计上就是资源轻量型的。

在阅读任何追踪信息时,你都可以将此作为快速诊断手段。扫描 CPU 通道,查找 `cudaOccupancyMaxActiveBlocksPerMultiprocessor`。每次出现都标记着一个“重量级、自适应启动”的内核,通常是 GEMM(通用矩阵乘法)、卷积或类似运算。而那些前面没有占用率查询的内核,则是 PyTorch 机械式启动的逐元素/规约类内核。

为什么 `cudaDeviceSynchronize` 耗时如此之长(约 1.78 毫秒)?

`cudaDeviceSynchronize` 会阻塞 CPU,直到该设备上的所有 GPU 工作完成。性能分析器会在活动窗口结束时发出此同步调用,以刷新事件。没有它,内核计时将会缺失。

一个耗时 1.78 毫秒的同步调用,覆盖了仅 26 微秒的实际 GPU 工作,这表明此次运行有 98% 的时间处于空闲状态。这是典型的开销瓶颈症状。

4096x4096 追踪信息

从上面的性能分析器表格分析中我们已经知道,为算法提供更大的矩阵,会使其从开销受限区域转变为计算受限区域。

让我们运行命令,更深入地探究这些追踪信息。

uv run 01_matmul_add.py --size 4096 --warmup

为什么同一个内核比其他内核耗时更长?

Image 29: 4096x4096 bf16 matmul kernel timings varying across profiler steps on the same GPU
图 17:在输入完全相同的情况下,一个矩阵乘法内核的运行时间比其他内核更长

在图 17 中,我们注意到 ProfileStep#3 的矩阵乘法内核在 GPU 上的耗时比其他步骤更长。这一点特别值得注意,因为启动的其他内核完全相同,这意味着没有涉及 cuBLAS 启发式方法。没有调度间隙,CPU 启动正常,并且这不是性能分析器的伪影。

图 17 中的这个追踪信息揭示了一个在理想化示例中容易被忽略的重要观点:即使在相同的硬件环境下,对相同的数据运行相同的代码,内核运行时间也不是恒定的。

让我们稍微修改脚本,使这一点更加具体。我们将迭代运行 20 次,并捕获每一步的数据。

- schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
+ schedule = torch.profiler.schedule(wait=0, warmup=0, active=20, repeat=1)

- for _ in range(5):
+ for _ in range(20):
Image 30: PyTorch profiler trace of 20 matmul iterations showing kernel runtime variance
图 18:在 20 次迭代中,同一个矩阵乘法内核以不同的速度运行

图 18 揭示了类似的发现。虽然每个内核完全相同,但它们的执行时间却不同。计算时间的差异可归因于多种原因:

  • GPU 时钟的空闲与升频
  • GPU 发热
  • GPU 电源管理
  • 驱动端的后台维护

只看平均值的读者会得出结论:一次矩阵乘法耗时约 1 毫秒(5 次均值 = 1084 微秒);而查看追踪记录的读者则会发现,矩阵乘法耗时约 580 微秒,除非 GPU 出现异常。这两种认知模型截然不同,且只有一种是正确的。

让我们看看 torch.compile 的实际运行

使用 torch.compile 一直让我们惊叹不已。开发者编写普通的即时执行 PyTorch 代码,但 PyTorch 会尝试捕获张量密集区域,将其转化为计算图,进行优化,并运行生成的代码。默认后端通常是 TorchInductor,其大致流程如下:

  1. TorchDynamo 将 Python 执行过程捕获为 FX 计算图
  2. 当涉及梯度时,AOTAutograd 准备前向/反向计算图
  3. Inductor 将计算图降级为优化的 CPU 或 GPU 代码。

在本节中,我们将讨论编译过程并查看性能分析器追踪记录。

uv run 01_matmul_add.py --size 4096 --warmup --compile

args.compile 标志会触发以下代码:

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

fn = torch.compile(fn) if args.compile else fn
Image 31: torch.compile region highlighted in a PyTorch profiler trace, showing TorchDynamo and Inductor frames
图 19:编译后的区域在追踪记录中显示为 TorchDynamo 和 Inductor 帧

在图 19 中,我们看到新增的 CPU 行被命名为 Torch-Compiled Region: 0/0,这指向了正在使用的编译函数。

我们是否将矩阵乘法和加法内核融合成了一个?

Image 32: Compiled trace showing aten::addmm replacing the eager aten::add and aten::mm pair
图 20:编译后的运行分派了一个单一的 aten::addmm

查看图 20,我们提出一个问题:我们是否真的将乘法和加法运算融合成了一个?

这是计算图层面的算子融合。Inductor 将我们的 torch.add(torch.matmul(x, w), b) 重写为单一的 aten::addmm(b, x, w) 调用。这里需要注意的重要一点是,它并没有生成一个新的融合 CUDA 内核。实际的 GPU 运算仍然是 ampere_bf16_s16816gemm_bf16_128x256_ldg8_f2f_stages_64x3_nn,与即时执行模式使用的 cuBLAS 内核相同。因此,这里的"融合"发生在分派器层面,而非内核层面。

PyTorch 提供了 `torch.addmm` 函数,它能一步完成我们之前分两步做的事情,即乘法和加法。我们鼓励读者查看该函数的追踪记录,并在下方评论区分享你们的观察结果!

torch.compile 的运行时架构

虽然我们在理论上知道编译函数时会发生什么,但亲眼看到它的实际运行同样重要。让我们来看看反映 torch.compile 运行时架构的 CPU 端层次结构。

TorchDynamo 缓存查找是 Dynamo 检查当前调用是否仍然与已编译内容(相同的输入形状、数据类型、设备和张量元数据)匹配的地方。如果任何内容不匹配,Dynamo 将重新编译。即使在编译之后,每次调用都会产生这一开销。

Torch-Compiled Region 是"进入"编译版本的包装器。AOTDispatcher Runtime Wrapper Prologue 是 AOT Autograd 的运行时包装器。即使我们在这里不需要梯度,AOTDispatcher 始终存在于堆栈中,负责处理张量元数据、视图追踪,并且如果 `requires_grad` 为真,它还会设置反向传播。

## Call CompiledFxGraph 是实际生成的代码运行的地方。"CompiledFxGraph" 后面的字符串是 FX 图的内容哈希值。它在所有三个活跃步骤中都是相同的,这确认了缓存命中。

你可以在磁盘上的 `/tmp/torchinductor_<user>/fxgraph` 目录下找到以此哈希值为键的生成代码,当你想要阅读 Inductor 实际生成的 Triton/C++ 代码时,这非常有用。

CUDA 启动次数是否减少了一半?

Image 33: compiled matmul trace showing Memcpy DtoD and GEMM kernels launched per step
图 21:每个编译后的步骤仍然启动两个 GPU 内核,一个设备到设备的内存拷贝和 GEMM

查看图 21 中的追踪记录,我们非常高兴地注意到每个步骤只有一个 `cudaLaunchKernel`。这一观察结果直接与我们之前在 GPU 追踪中看到的情况相矛盾。每个步骤仍然有两个内核被启动,即 Memcpy DtoD(设备到设备)和 GEMM。回到 CPU 追踪记录,我们发现自己完全漏掉了 `cudaMemcpyAsync` 的调度。

addmm 计算的是 out = α·A·B + β·C,而 cuBLAS 的带偏置加法(bias-add)后处理(epilogue)会写入一个目标缓冲区,该缓冲区需要已包含偏置值。后处理可以理解为 GEMM 之后发生的所有操作。在深度学习领域,我们不断提出各种 GEMM 后处理,例如激活函数、偏置加法、归一化等等。正因如此,才有了各种 cuBLAS GEMM 带后处理的内核变体。

如果你对 torch.compile 使用不同的模式,你会注意到启动的内核变体也不同。你可以亲自尝试一下,并在下方评论中分享你的观察结果!

因此,Inductor 生成的代码执行以下操作:

  • out = copy(C) ← 这是设备到设备(DtoD)的内存拷贝(32 MB,耗时约 33 µs)
  • out = α·(A·B) + β·out ← GEMM,其中 α=β=1,将偏置加法融合到回写操作中

从数学上讲,结果仍然相同。偏置加法并非没有代价,因为我们预先支付了一次内存拷贝,外加一次开销稍高的 GEMM 后处理。

人们可能期望的融合——即 x·w + b(此处为 out = α·A·B + β·C)能坍缩成一个内核,且没有额外的内存流量——并未发生。Inductor 保留了两次内存访问操作,只是将偏置拷贝重新标记为内存拷贝,将加法重新标记为 GEMM 后处理。

真正融合的实现会跳过内存拷贝。这正是 FlashAttention 风格的手写内核所做的,也是 Inductor 通过 Triton 代码生成所能做到的,但对于一个 4096×4096 的 bf16 矩阵乘法,Inductor 显然判定“使用 cuBLAS,通过后处理设置来完成偏置加法”是最佳路径。

CPU 开销增加了,而非减少

这是比较即时执行(eager)和编译执行时最容易忽略的一点:

步骤 即时执行耗时(毫秒) 编译执行耗时(毫秒)
#2 0.1 0.2
#3 0.07 0.1
#4 0.07 0.1

torch.compile 在 CPU 上每一步的开销大约是前者的 2 倍。这是因为每次调用都会遍历完整的 Dynamo > AOTAutograd > Inductor 堆栈,除此之外还有我们无论如何都会进行的 aten::addmm 分发。编译流水线是为包含数十个操作的机器学习模型设计的,在这些模型中,每次调用的开销可以被分摊(对于单个操作来说,这就像是一种额外税负)。

torch.compile 有一个 mode 参数。这留给读者作为课后作业,请查阅文档并提出一种能够降低 CPU 开销的模式。🤗

跟踪读取速查表

我们梳理过的模式速查表。思路是:如果在追踪记录中看到这个,它通常意味着什么。

性能分析器表格

你看到的内容 它通常意味着什么
自身 CPU 总时间 ≫ 自身 CUDA 总时间(CPU 以毫秒计,GPU 以微秒计) 开销受限。CPU 花在调度上的时间比 GPU 花在计算上的时间还多。应增大工作规模(更大的矩阵、批处理操作)或融合调用。
自身 CPU 总时间 ≈ 自身 CUDA 总时间,两者均以毫秒计 计算受限。GPU 是瓶颈,这通常是你希望看到的情况。
单个事件主导了 CUDA 总时间 那就是你的热点。从那里开始优化。
单个事件的调用次数巨大 即使每次调用开销很小,也可能是潜在瓶颈。检查是否可以融合或批处理。
某行的 CPU 总时间 ≫ 自身 CPU 时间 大部分开销在子事件中。应深入分析嵌套事件,而非父事件。

CPU 时间线

你看到的内容 它通常意味着什么
第一个 ProfileStep 比其余宽得多 冷启动开销:工作空间分配、cuBLAS 启发式算法、惰性模块加载。增加预热迭代次数和/或调度器的预热参数。
record_function("...") 开始与内部第一个 aten::* 之间存在较大间隔 同样是冷启动开销,只是放大了看。注解已进入,但调度尚未发生。
cudaLaunchKernel 之前的 cudaOccupancyMaxActiveBlocksPerMultiprocessor 一个重量级、自适应启动的内核(GEMM、卷积等)。cuBLAS 正在向驱动查询一个 SM 上能容纳多少个块,以便选择内核变体。
没有前置占用查询的 cudaLaunchKernel 一个元素级或规约内核,具有固定且资源占用小的特征。无需规划。
活动窗口末尾出现较长的 cudaDeviceSynchronize 性能分析器正在刷新事件。其持续时间主要是 GPU 完成待处理工作的时间,并非真正的 CPU 开销。覆盖微小 GPU 工作的同步是典型开销受限症状。
你并未编写的 cudaMemcpyAsync 通常是隐藏的设备到设备拷贝。常见于 addmm 在 GEMM 后处理阶段之前,用偏置值初始化其目标缓冲区时。

GPU 时间线

你看到的内容 它通常意味着什么
GPU 时间线上的活动缓冲区请求 性能分析器正在分配/重新填充其自身的事件缓冲区。第一个通常解释了初始的 CPU↔GPU 时间线偏移。
单步中两个内核之间存在间隙 很可能在执行过程中又发起了一次缓冲区请求。通过运行更多迭代来确认:如果它只出现一次,那就是性能分析器的问题,而不是你的代码。
同一个内核在不同步骤中的计时结果不同 GPU 时钟、温度、电源管理、驱动维护。要读取追踪数据,而不仅仅是平均值。
一个名为 ampere_bf16_s16816gemm_... 的内核 矩阵乘法实际的 cuBLAS GPU 工作。对于相同的形状/数据类型,该内核名称在即时执行模式和编译模式下通常是相同的。
在 GEMM 操作之前立即出现 Memcpy DtoD addmm 后置操作的偏置拷贝。这里的“融合”发生在调度器层面,而非内核层面。

调度链

你所看到的 它通常的含义
ProfileStep#N → <record_function name> → aten::* → aten::mm / aten::bmm / aten::add 标准的嵌套调用层级。自身时间不包括子项;总时间则包括子项。
aten::matmul 解析为 aten::mm 二维 × 二维矩阵乘法。
aten::matmul 解析为 aten::bmm(带有额外的 CUDA 运行时调用) 对三维及以上张量执行的批量矩阵乘法。cuBLAS 会执行更多启发式工作来选择变体。
aten::addmm(b, x, w) 替代了单独的 aten::add + aten::mm 组合 调度器层面的算子融合。GPU 内核仍然是相同的 GEMM,只是偏置加法被折叠进了后置操作中。

torch.compile

你所看到的 它通常的含义
Torch-Compiled 区域:CPU 通道中的 K/M 行 你正处于一个编译后的函数内部。
每一步都进行 TorchDynamo 缓存查找 Dynamo 正在验证形状/数据类型/设备是否与缓存的编译结果匹配。每次调用都会产生此开销,即使在编译完成后也是如此。
即使没有梯度,也存在 AOTDispatcher 运行时包装器前导 AOTAutograd 的运行时包装器始终存在于调用栈中,负责处理张量元数据和视图追踪。
## 每一步都使用相同的哈希值调用 CompiledFxGraph <hash> 命中了生成代码的缓存。生成的源代码位于 /tmp/torchinductor_<user>/fxgraph/<hash> 目录下。
对于微小的算子,使用 torch.compile 时每步的 CPU 时间高于即时执行模式 这是预期的。Dynamo → AOTAutograd → Inductor 这一调用栈本身是一种开销,只有在处理大量算子时才能被摊平。

结论

我们从一个小小的矩阵乘法加加法运算开始,并以此为切入点,学习如何阅读 PyTorch 性能分析器。在此过程中,我们掌握了一些思维模型,这些模型同样适用于更大的工作负载。这是“PyTorch 性能分析”系列的第一站。在后续的文章中,我们将逐步告别这个双算子小玩具,沿着复杂度的阶梯向上攀登,审视更大的构建模块,并最终分析真实模型。

感谢 Noe Flandre、Suvaditya Mukherjee 和 Vidit Ostwal 对本文初稿的审阅!

这篇博文使用大语言模型进行了润色。这绝不意味着我们让一个智能体在后台运行并让它生成整篇博客。我们团队中有些成员并非英语母语者,我们认为大语言模型(大多基于英语语料训练)可以纠正一些简单的语法错误,或改写那些听起来不够简洁、不够清晰的句子。希望这能有助于理解“如果这是大语言模型生成的,我为什么还要读”这个想法。🤗

阅读原文huggingface.co