此人写作功力深厚！将这种理解与性能剖析（此处无耻地推销我的系列文章）结合，你将如虎添翼！http://hf.co/blog/torch-profiler…

PyTorch 性能剖析（第一部分）：torch.profiler 新手入门指南

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

• 矩阵乘法与加法操作
• 64x64 轨迹
  • 为什么 ProfilerStep#2 耗时那么长？
  • 为什么 CPU 和 GPU 轨道之间存在约 2.5 毫秒的偏移？
  • 事件链
  • 为什么 matmul 多了一次 CUDA runtime 调用？
  • 为什么 cudaDeviceSynchronize 这么大（约 1.78 毫秒）？
• 4096x4096 轨迹
  • 为什么同一个内核比其他内核耗时更多？
• 来看看 torch.compile 的实际效果
  • 我们是否将 matmul 和 add 内核融合成了一个？
  • torch.compile 的运行时架构
  • CUDA 启动次数是否减少了一半？
  • CPU 开销不降反升
• 轨迹解读速查表
  • 性能剖析表格
  • CPU 轨道
  • GPU 轨道
  • 调度链
  • torch.compile
• 结论

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

你无法优化的东西，就无法剖析。

无论你是想从大语言模型（LLM）中压榨出更多每秒 token 数，还是想削减推理的毫秒级延迟，抑或只是想知道为什么你的训练循环跑得比规格表承诺的慢，最终都绕不开性能剖析。但问题在于，性能剖析的入门门槛极高。轨迹图是一堵密集的彩色矩形墙，事件名称听起来吓人，大多数教程又假设你已经能读懂它们。所以，即使我们知道应该做剖析，打开轨迹图也常常让人觉得还是以后再说（或者交给别人）比较好。

本文以及它所开启的系列文章，就是我们降低门槛的尝试。我们从初学者的角度记录学习过程，除了基础的 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 运行脚本。在 Hugging Face 基础设施上设置 GPU 非常容易，可以通过 Spaces 的开发者模式 (https://huggingface.co/docs/hub/spaces-dev-mode) 来试验脚本。也可以使用 Hugging Face Jobs 管道 (https://huggingface.co/docs/huggingface_hub/en/guides/jobs) 运行脚本。

矩阵乘法与加法操作

正如 Dr. Sara Hooker (https://youtu.be/7knwihgj0fU?si=uvzGH-J9bsCHP4Nn&t=2199) 精辟指出的那样：就像我们主要由水组成一样，深度神经网络主要由矩阵乘法组成。既然矩阵乘法如此基础，用其他操作开始我们的剖析之旅未免可惜。

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

在矩阵乘法中加入矩阵加法，模拟了神经元中权重和偏置的交互。这个加法（双关语：addition 亦有“补充”之意）有助于我们理解它如何为后文的编译铺平道路。

要进行性能剖析，我们将使用 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 大小矩阵的 matmul add 剖析表格 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-profiler/profile-table-64.png)
图 1：64 大小矩阵的 matmul add 剖析表格

.txt 文件保存了剖析表格。打开文件后，如图 1 所示，你会看到一个大型表格，第一列是剖析作用域内被触发的事件。其他列与事件在 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 可以非常快速地计算小规模矩阵乘法，因此我们的代码大部分时间都花在了准备内核、启动内核、发送数据相乘以及收集结果上。这个概念被称为开销受限算法。摆脱这种状态的最简单方法是使用更大的矩阵乘法。

uv run 01_matmul_add.py --size 4096

4096 大小矩阵的 matmul add 算法剖析表格 (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

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

现在我们转向可视化调度链，它存在于 .json 产物中。你可以将其上传到 Perfetto UI (https://ui.perfetto.dev/) 查看轨迹，也可以使用 uvx trace-util -f traces -b /traces 直接生成 Perfetto 链接。

64x64 轨迹

PyTorch 剖析器轨迹图：64×64 bf16 矩阵乘法后接加法，在 CUDA GPU 上运行 (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 剖析器轨迹图，CPU 轨道和 GPU 轨道在 Perfetto 中并排标注 (https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/torch-profiler/gpu-cpu-trace.png)
图 4：PyTorch 剖析器轨迹的 CPU 和 GPU 轨道

图 4 中有两个轨道，一个用于 CPU 活动，一个用于 GPU 活动。在 CPU 轨道中，你会注意到三个剖析步骤（从 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) 。

现在，让我们戴上侦探帽，调查轨迹图并提出一些问题。

为什么 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 加速线性代数库）的启发式决策，或惰性模块加载。我们既可以忽略它，也可以在剖析之前运行更多的预热步骤（这是标准做法）。

在剖析中，预热是指在实际剖析之前将事件运行几次。GPU 完成的预工作（包括上述几点）是一次性工作，我们不想剖析它们。在我们的示例中，有两个预热阶段：一个是在进入剖析器之前对函数进行循环的预热，另一个是在剖析器内部通过 warmup 参数实现的预热。

在本节中，我们