7.0
热点
AI SCORE
技术实践2026-05-29 08:00
PyTorch 性能分析入门指南
Hugging Face Blog#PyTorch#性能优化#ML
Editor brief · 编辑速览
Hugging Face 发布 torch.profiler 教程，教你如何检测和优化模型训练性能。
无论你是在尝试从大语言模型（LLM）中榨取更多每秒令牌数，还是想减少推理的毫秒延迟，亦或只是想搞明白你的训练循环为什么比规格表承诺的速度还慢，最终都避免不了要进行性能测量。
问题是性能测量有很陡的学习曲线。痕迹是密密麻麻的彩色矩形。事件带有令人生畏的名称。大多数教程都假设你已经懂得如何阅读它们。所以即使我们知道应该进行性能测量，打开一条痕迹也会让人感觉是个最好留到以后（或交给别人）处理的琐事。这篇文章，以及它开启的系列，正是我们为降低这个学习曲线而做的尝试。
这是 PyTorch 中的性能测量系列文章的开篇，在这个系列中我们会逐步建立阅读分析器痕迹的能力，并用它来驱动优化工作：
我们从初学者的角度记录这段学习之旅。除了 PyTorch 的基础知识外，没有其他先决条件。把这篇文章当作一次轻松的阅读，期间会有一些"恍然大悟"的时刻。文章结构有意采用问题驱动的方式：我们打开一条痕迹，问"等等，为什么会这样？"，然后追踪答案直到某个地方突然点亮。到文章末尾，你应该能够：
在开始之前，有两个定义会让下面的内容读起来更清楚：
GPU 内核是在 GPU 的许多线程上并行运行的程序。
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 模块。涉及的步骤是：
准备好要测量的代码（这里是 def fn，它包装了矩阵乘法和矩阵加法）
注解算法。虽然这完全是可选的，但我们建议这样做。record_function 将我们的函数注解为 matmul_add，这样在痕迹中就很容易导航（我们稍后会提到）
def step():
with torch.profiler.record_function("matmul_add"):
return fn(x, w, b)
用 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()
# the profiler table
prof.key_averages().table(sort_by="cuda_time_total", row_limit=15)
# the profiler trace
prof.export_chrome_trace(trace_path)
分析器导出两个不同的产物：
分析器表格：提供算法的统计摘要。它回答"什么消耗了最多时间"这个问题。这对于找出瓶颈非常有帮助。瓶颈是指消耗最多时间的事件，可能是管道的瓶颈，或者是被触发很多次的事件。
分析器痕迹：提供时间执行视图。回答"何时以及为什么操作发生"这个问题，描绘在 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
.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 时间以 ms 为单位，而 GPU 时间以 us 为单位。从角度来说，在 GPU 上花费的时间（ampere_bf16_s16816gemm... 内核）不到在 CPU 上花费时间（matmul_add 操作）的 1%。GPU 大部分时间都处于闲置状态，这是一个立即的危险信号。发生这种情况的原因是 GPU 可以非常快地计算小的矩阵乘法，所以我们的代码花费大部分时间在准备内核、在 GPU 上启动它们、将数据发送进行乘法运算以及收集结果。这个概念被称为开销受限的算法。
摆脱这种状态的最简单方法是使用更大的矩阵乘法。
uv run 01_matmul_add.py --size 4096
图 2 中的最后两行是：
Self CPU time total: 4.908ms
Self CUDA time total: 4.495ms
两个时间都以 ms 为单位，这意味着我们仅通过增加矩阵乘法的大小就实现了更多 GPU 时间。如果你看图 2，你还会注意到最多的 CUDA 时间现在由 GPU 内核（ampere_bf16_s16816gemm_..）消耗，而不是启动它的 CPU 操作（matmul_add）。这意味着我们确实能够从开销受限转变为计算受限。
现在我们进入可视化派遣链，它存在于 .json 产物中。你可以将它们上传到 Perfetto UI 并查看痕迹，或者可以使用 uvx trace-util -f traces -b <hf_uname>/traces 来直接生成 Perfetto 链接。
在图 3 中，我们看到矩阵乘法和加法的分析器痕迹。这里条形宽度表示事件的持续时间，垂直嵌套表示调用层次，CPU 通道表示在 CPU 上发生的事件，而 GPU 通道显示实际的内核执行。你可能还会注意到空白区域，这些是等待或闲置时间。
脚本使用默认配置运行，这些配置是：
使用 Perfetto，我们建议使用键盘来更快地访问痕迹。你可以使用"W A S D"来导航痕迹。
图 4 中有两个通道，一个用于 CPU 活动，一个用于 GPU 活动。在 CPU 通道中，你会注意到三个分析步骤（从 ProfilerStep#2 开始）。这来自于调度。
schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
wait 跳过嘈杂的初始化操作（ProfilerStep#0），warmup 在不记录的情况下运行分析器（ProfilerStep#1），active 是在 trace 中显示的内容。可以在脚本中找到使用的 schedule。
让我们戴上侦探帽，深入调查 trace 并提出一些问题。
在图 5 中，我们注意到 ProfileStep#2 相比其他步骤耗时更长，仔细观察会发现 matmul_add 注解中也有类似的模式。关键线索在注解内部，而不是注解本身：
图 6 中显示的约 228 µs 是进入 record_function("matmul_add") 和 PyTorch 实际分派 aten::matmul 之间的"死窗口"。这可能由多个原因引起，包括工作区分配、cuBLAS（NVIDIA 的专有 GPU 加速库，用于执行基础线性代数运算）启发式算法或延迟模块加载。我们可以视而不见，或在分析之前运行更多预热步骤（这是标准做法）。
就分析而言，预热是指在实际分析之前运行事件几次。GPU 所做的预工作（包括上述指针）是一次性的工作，我们不想对其进行分析。在我们的例子中，我们有两个预热阶段，一个是在进入分析器之前实际循环该函数，另一个是在分析器内部通过 warmup 参数实现的。在本部分中，我们已启用了实际迭代以及 schedule。
uv run 01_matmul_add.py --warmup
具有预热的 64x64 的 Perfetto Trace
在图 7 中，我们看到每个分析步骤花费的时间相似，但这并不意味着我们能够优化一次性开销。我们预热了运行，以便不对开销进行分析。我们认为仓促结束这一部分而不给出解决方案的提示会对读者不公平，因此这里有一个链接可以进一步了解优化启动开销。
在图 8 中，我们看到 CPU 和 GPU 通道之间有大约 2.5 ms 的偏移：这是 CPU 提交 CUDA 核后与它们实际开始执行之间的延迟。人们可能会认为预热阶段加上 schedule 的 wait 和 warmup 应该让 GPU 保持忙碌，从而减少偏移。
为了揭示真正发生的情况，让我们稍微改变一下 schedule：
- schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
+ schedule = torch.profiler.schedule(wait=0, warmup=0, active=3, repeat=1)
图 9 显示了在任何操作之前 GPU 通道中都有一个活动缓冲区请求。让我们再放大一点。
放大到 GPU trace 后，我们注意到 ProfileStep#0 的 matmul 和 add 核（其 CPU trace 在图中不可见）一个接一个地执行，而 ProfileStep#1 的核之间有一个窗口。最好的解释是缓冲区溢出，并且在核执行期间发出了另一个缓冲区请求（在 GPU VRAM 上分配内存的请求）。
排除其他可能性的最好方法是对更多迭代进行分析，看看 trace 的其他部分是否出现类似的窗口。为此，我们以 active=20 运行。
如图 11 所示，我们在 ProfileStep#1 中看到了类似的趋势。这与我们之前的发现相一致，我们可以安全地得出结论，它确实是另一个缓冲区请求。
在图 12 中，我们看到了嵌套的 CPU 调用。这是一个重要的可视化，让人们能够理解调度链真正的样子。
我们从 ProfileStep#<id> 开始，它封装了分析步骤。由于我们注解了该步骤，我们看到了 matmul_add 行。matmul_add 由两个 aten 调用组成，一个用于矩阵乘法，一个用于矩阵加法。
aten::matmul 是用户面向 PyTorch matmul 调用的 ATen 级调度。aten::mm 是 2D 矩阵-矩阵乘法后端。
值得注意的是，如果我们向矩阵添加批处理轴，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)
在图 13 中，在将批处理轴添加到输入后，aten::matmul 现在封装了一堆其他先决条件 CUDA 运行时调用以及 aten::bmm（而不是 aten::mm）。这也暗示了 cuBLAS 需要执行的启发式算法，以便为程序调度正确的（最合适的）核。
在本文的其余部分，除非另有说明，否则我们将使用简单的 2D 矩阵。
我们注意到对于 aten::mm，有两个 CUDA 运行时调用，即 cudaOccupancyMaxActiveBlocksPerMultiprocessor（在图 14 中标注）和 cudaLaunchKernel，而对于 aten::add，只有 cudaLaunchKernel。
cudaOccupancyMaxActiveBlocksPerMultiprocessor 是一个规划调用，完全在 CPU 端。它问：给定一个核函数、一个选定的块大小和一个选定的动态共享内存大小，这个核的多少个块可以同时驻留在一个 SM（流多处理器）上？
这引出了一个问题，为什么我们需要为 matmul 规划，而不需要为 add 规划？
为了理解这一点，我们需要查看核的资源占用。如果点击 GPU 核，你将能够检查相应核的资源占用。
在图 15 中，我们注意到对于矩阵乘法，每个线程的寄存器和共享内存是动态的（基于矩阵的大小）。cuBLAS 附带数百个核变体，每个变体都有一个启发式驱动的启动路径，需要关于硬件容量的运行时信息。占用率查询是该启发式的一部分。概念上，我们可以将 GPU 加速的 matmul 视为在独立的瓦片上工作：我们使用多少个瓦片以及每个瓦片需要多大取决于矩阵和硬件。现代算法比这复杂得多，但这仍然是一个很好的参考框架。
从图 16 我们看到加法的占用说 32 个寄存器和零共享内存。这很简单。没有什么可查询的，因为没有硬件资源会限制占用率。根据设计，该核是资源轻量的。
在阅读任何 trace 时，你可以使用它作为快速诊断。扫描 CPU 通道中的 cudaOccupancyMaxActiveBlocksPerMultiprocessor。每次出现都标记一个"重型、自适应启动的"核，通常是 GEMM（通用矩阵乘法）、conv 或类似的。没有先前占用率查询的核是 PyTorch 机械启动的按元素/约化的集合。
cudaDeviceSynchronize 阻塞 CPU 直到此设备上的所有 GPU 工作完成。分析器在活动窗口结束时发出此同步以刷新事件。如果没有它，核计时将丢失。
1.78 ms 的同步覆盖 26 µs 的实际 GPU 工作告诉你这次运行 98% 的时间是空闲的。这是教科书式的开销受限症状。
我们已经从分析器表分析（上面）知道，向我们的算法提供更大的矩阵会将其从开销受限区域移出到计算受限区域。
让我们运行命令并深入探索 trace。
uv run 01_matmul_add.py --size 4096 --warmup
在图 17 中，我们注意到 ProfileStep#3 的 matmul 核在 GPU 上花费的时间比其他步骤更长。值得特别注意的是，启动的其他核是完全相同的，这意味着没有涉及 cuBLAS 启发式算法。没有调度间隙，CPU 启动是正常的，这不是分析器伪影。
图 17 中的这个 trace 说明了一个在理想化示例中容易忽视的有用要点：核运行时不是常数，即使在相同的硬件环境中运行相同的代码处理相同的数据。
让我们通过稍微修改脚本来使这更具体。我们运行迭代 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):
GPU 时钟在空闲和加速状态下的差异
驱动程序侧面的内务处理
只看平均值的读者会得出结论，矩阵乘法耗时约 ~1 ms（5 个值的平均值 = 1084 µs）；而查看追踪的读者会看到矩阵乘法耗时约 ~580 µs，除非 GPU 出现异常。这是两种非常不同的心智模型，其中只有一种是正确的。
使用 torch.compile 一直令人惊叹。你编写普通的 eager PyTorch 代码，但 PyTorch 会尝试捕获张量密集区域，将其转换为图，优化它们，并运行生成的代码。默认后端通常是 TorchInductor，广泛流程如下：
TorchDynamo 将 Python 执行捕获到 FX 图中
AOTAutograd 在涉及梯度时准备前向/后向图
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
在图 19 中，我们看到名为 Torch-Compiled Region: 0/0 的新 CPU 行，它指向正在使用的编译函数。
查看图 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，与 eager 模式使用的 cuBLAS 核心相同。所以这里的"融合"是在调度器级别，而不是在核心级别。
PyTorch 提供了 torch.addmm 函数，它用一步完成我们需要两步的工作，即乘法和加法。我们鼓励读者查看这个函数的追踪并在下面的评论中分享你的观察！
虽然我们在理论上知道编译函数时会发生什么，但同样重要的是看到它的实际运作。让我们查看反映 torch.compile 运行时架构的 CPU 端层次结构。
TorchDynamo Cache Lookup 是 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++ 时很有用。
查看图 21 中的追踪，我们非常高兴地注意到每个步骤只有一个 cudaLaunchKernel。这个观察与我们在 GPU 追踪中看到的直接矛盾。每个步骤仍然有两个核心被启动，即 Memcpy DtoD（设备到设备）和 GEMM。回到 CPU 追踪，我们注意到我们完全错过了 cudaMemcpyAsync 调度。
addmm 计算 out = α·A·B + β·C，cuBLAS 的 GEMM 带偏置加法尾声将结果写入需要已经包含偏置的目标缓冲区。尾声可以被认为是 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，通过尾声设置处理偏置"是最佳路径。
这是比较 eager 和编译运行时最容易遗漏的一点：
torch.compile 在每个步骤的 CPU 成本大约是 eager 的 2 倍。这是因为每次调用都会遍历完整的 Dynamo > AOTAutograd > Inductor 堆栈，加上我们无论如何都有的相同 aten::addmm 调度。编译管道是为拥有数十个运算的 ML 模型设计的，其中每调用开销会摊销（对于单个运算，这是一个税收）。
torch.compile 有一个 mode 参数。这是留给读者的作业，让你阅读文档并想出一种可能会降低 CPU 开销的模式。🤗
我们讲解过的模式的快速参考。其思想是：如果你在追踪中看到这个，这就是 thi