Profiling in PyTorch (Part 1): torch.profiler 初学者指南深度解读

背景

在深度学习领域，有一句名言：“无法剖析（Profile）的东西，就无法优化。”

无论你是试图从大语言模型（LLM）中挤出更多的每秒令牌（tokens per second）数，还是希望将推理延迟减少几毫秒，亦或是仅仅想弄明白为什么你的训练循环运行速度远低于规格说明书所承诺的性能，最终的路径都指向同一个终点：性能剖析（Profiling）。

然而，性能剖析的学习曲线相当陡峭。生成的追踪文件（Traces）是由密集排列的彩色矩形组成的“高墙”，其中的事件名称令人望而生畏。大多数教程都默认读者已经具备阅读这些复杂图表的能力。因此，即使我们知道应该进行剖析，打开一个追踪文件往往被视为一种负担，被推迟到“以后再说”或交给别人处理。

本文是 Hugging Face 博客系列《PyTorch 中的性能剖析》的第一部分，旨在降低这一入门门槛。该系列将逐步构建阅读剖析器追踪文件（Profiler Traces）的技能，并利用这些洞察驱动优化工作。

系列规划如下：

• 第 1 部分（本文）：从最简单的操作开始——矩阵乘法后接偏置加法（bias add），学习如何阅读剖析器返回的数据。
• 第 2 部分：扩展到 nn.Linear 和小型多层感知机（MLP），利用追踪文件激发优化思路，并窥探底层内核（Kernels）。
• 第 3 部分：在 transformers 库的大语言模型应用中整合所有技巧。

本文以初学者的视角记录这一过程，除了基本的 PyTorch 知识外无需其他先决条件。文章结构以问题为导向：打开追踪文件，提出“等等，为什么会出现这种情况？”，并追踪答案直到豁然开朗。

在开始之前，需要明确两个关键定义，这将有助于理解后续内容：

1. GPU 内核（GPU Kernel）：在 GPU 的许多线程上并行运行的程序。
2. CPU 调度：CPU 负责调度和启动这些内核。

通常不需要手动编写 GPU 内核；当使用 PyTorch 操作时，它会自动被转换为一个或多个在内核上执行任务的程序。

核心内容

为了演示剖析过程，我们使用一个简单的脚本 01_matmul_add.py（建议在单独标签页中打开并逐步阅读代码）。实验环境使用 NVIDIA A100-SXM4-80GB GPU。

1. 剖析对象：矩阵乘法与加法

正如 Sara Hooker 博士所言，就像人类主要由水组成一样，深度神经网络主要由矩阵乘法组成。作为最基础的操作，它是剖析之旅的理想起点。

核心函数定义如下：

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

这个操作模拟了神经元中权重和偏置的交互方式。这种简单的组合有助于理解后续编译（Compilation）的工作原理。

2. 使用 torch.profiler 进行剖析

使用 torch.profiler 模块进行剖析主要包含以下步骤：

步骤一：准备代码并添加注解 虽然注解是可选的，但强烈建议这样做。record_function 可以将我们的函数标记为 matmul_add，这在后续的追踪文件中便于导航。

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

步骤二：包裹剖析上下文 使用 torch.profiler.profile 上下文管理器包裹代码。建议运行多次以预热 GPU。

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU, # CPU 活动
        torch.profiler.ProfilerActivity.CUDA, # GPU 活动
    ],
) as prof:
    for _ in range(5):
        step()
    prof.step()

步骤三：导出剖析结果 剖析器会生成两种不同的工件：

1. 剖析表（Profiler Table）：提供算法的统计摘要，回答“什么最耗时”。这有助于识别热点（Hotspots）。热点是指耗时最长、可能是管道瓶颈或触发频率极高的事件。
2. 剖析追踪（Profiler Trace）：提供时间执行视图，回答“操作何时发生以及为什么发生”，描绘 CPU 和 GPU 上的活动。这有助于调查启动的内核、启动延迟、CPU 与 GPU 活动之间的重叠等。

导出代码如下：

# 剖析表
prof.key_averages().table(sort_by="cuda_time_total", row_limit=15)
# 剖析追踪
prof.export_chrome_trace(trace_path)

3. 解读剖析结果：小矩阵场景（Overhead Bound）

首先运行小尺寸矩阵乘法： uv run 01_matmul_add.py --size 64

生成的 .txt 文件包含剖析表。我们需要关注以下几点：

• Self vs. Total：

  • Self CPU/CUDA：仅测量事件本身花费的时间，不包括其子事件。
  • CPU/CUDA total：包括事件本身及其所有子事件的总时间。
  • 注意：matmul_add 的 "CPU total" 包含了自身时间加上它触发的子事件时间。

• 数据洞察： 在小矩阵（64x64）场景下，最后两行数据显示：

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

  CPU 时间以毫秒（ms）为单位，而 GPU 时间以微秒（us）为单位。GPU 时间（内核 ampere_bf16_s16816gemm...）仅占 CPU 时间的不到 1%。

  结论：GPU 大部分时间处于空闲状态，这是一个明显的红色警报。原因是 GPU 计算小矩阵非常快，代码大部分时间都花在准备内核、在 GPU 上启动它们、发送数据以及收集结果上。这种算法被称为开销主导型算法（Overhead-bound algorithm）。

  解决方案：最简单的方法是增大矩阵乘法的大小。

4. 解读剖析结果：大矩阵场景（Compute Bound）

运行大尺寸矩阵乘法： uv run 01_matmul_add.py --size 4096

数据显示：

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

两者均以毫秒为单位。通过增加矩阵大小，我们显著增加了 GPU 的计算时间占比。此时，大部分 CUDA 时间由 GPU 内核（ampere_bf16_s16816gemm_..）占据，而不是由启动它的 CPU 操作（matmul_add）占据。

结论：我们成功从“开销主导”转变为计算主导型（Compute Bound）。

5. 可视化调度链（Dispatch Chain）

追踪文件（.json 格式）包含了调度链的可视化信息。你可以将其上传到 Perfetto UI 查看，或使用 uvx trace-util 生成链接。

在 64x64 的追踪图中，条形的宽度表示事件的持续时间，垂直嵌套表示调用关系。这揭示了从 Python 调用一直到底层 CUDA 内核的完整事件链。

关键要点

• 剖析是优化的前提：如果不了解性能瓶颈在哪里，就无法有效地进行优化。
• 两种主要工件：
  • 剖析表（Table）：用于统计分析和识别热点（Hotspots）。
  • 剖析追踪（Trace）：用于时间序列分析，查看 CPU/GPU 活动的时序、延迟和重叠。
• Self vs. Total 的区别：
  • Self 时间仅指该函数/操作本身执行的时间。
  • Total 时间包含该操作及其所有子操作的时间。
• 开销主导（Overhead-bound）vs. 计算主导（Compute-bound）：
  • 小矩阵：GPU 计算极快，CPU